Neon-Desk/components/financials/statement-matrix.tsx

'use client';

import { Fragment, memo, useMemo, useRef } from 'react';
import { useVirtualizer } from '@tanstack/react-virtual';
import { ChevronDown, ChevronRight } from 'lucide-react';
import type { FinancialStatementPeriod, SurfaceFinancialRow, DetailFinancialRow } from '@/lib/types';
import { cn } from '@/lib/utils';
import type {
  StatementInspectorSelection,
  StatementTreeNode,
  StatementTreeSection
} from '@/lib/financials/statement-view-model';

type MatrixRow = SurfaceFinancialRow | DetailFinancialRow;

type StatementMatrixProps = {
  periods: FinancialStatementPeriod[];
  sections: StatementTreeSection[];
  selectedRowRef: StatementInspectorSelection | null;
  onToggleRow: (key: string) => void;
  onSelectRow: (selection: StatementInspectorSelection) => void;
  renderCellValue: (row: MatrixRow, periodId: string, previousPeriodId: string | null) => string;
  periodLabelFormatter: (value: string) => string;
  dense?: boolean;
  virtualized?: boolean;
};

function isSurfaceNode(node: StatementTreeNode): node is Extract<StatementTreeNode, { kind: 'surface' }> {
  return node.kind === 'surface';
}

function rowSelected(
  node: StatementTreeNode,
  selectedRowRef: StatementInspectorSelection | null
) {
  if (!selectedRowRef) {
    return false;
  }

  if (node.kind === 'surface') {
    return selectedRowRef.kind === 'surface' && selectedRowRef.key === node.row.key;
  }

  return selectedRowRef.kind === 'detail'
    && selectedRowRef.key === node.row.key
    && selectedRowRef.parentKey === node.parentSurfaceKey;
}

function surfaceBadges(node: Extract<StatementTreeNode, { kind: 'surface' }>) {
  const badges: Array<{ label: string; tone: 'default' | 'warning' | 'muted' }> = [];

  if (node.row.resolutionMethod === 'formula_derived') {
    badges.push({ label: 'Formula', tone: node.row.confidence === 'low' ? 'warning' : 'default' });
  }

  if (node.row.resolutionMethod === 'not_meaningful') {
    badges.push({ label: 'N/M', tone: 'muted' });
  }

  if (node.row.confidence === 'low') {
    badges.push({ label: 'Low confidence', tone: 'warning' });
  }

  const detailCount = node.row.detailCount ?? node.directDetailCount;
  if (detailCount > 0) {
    badges.push({ label: `${detailCount} details`, tone: 'default' });
  }

  return badges;
}

function badgeClass(tone: 'default' | 'warning' | 'muted', dense?: boolean) {
  const baseClasses = dense
    ? 'rounded border px-1 py-0.5 text-[9px] uppercase tracking-[0.12em]'
    : 'rounded-full border px-2 py-0.5 text-[10px] uppercase tracking-[0.14em]';

  if (tone === 'warning') {
    return cn(baseClasses, 'border-[#84614f] bg-[rgba(112,76,54,0.22)] text-[#ffd7bf]');
  }

  if (tone === 'muted') {
    return cn(baseClasses, 'border-[color:var(--line-weak)] bg-[rgba(80,85,92,0.16)] text-[color:var(--terminal-muted)]');
  }

  return cn(baseClasses, 'border-[color:var(--line-weak)] bg-[rgba(88,102,122,0.16)] text-[color:var(--terminal-bright)]');
}

// Flatten tree nodes for virtualization
type FlattenedNode = {
  node: StatementTreeNode;
  sectionKey: string;
  sectionLabel?: string;
  isSectionHeader?: boolean;
};

function flattenSections(sections: StatementTreeSection[]): FlattenedNode[] {
  const result: FlattenedNode[] = [];

  function flattenNodes(nodes: StatementTreeNode[], sectionKey: string): void {
    for (const node of nodes) {
      result.push({ node, sectionKey });

      if (node.kind === 'surface' && node.expanded && node.children.length > 0) {
        flattenNodes(node.children, sectionKey);
      }
    }
  }

  for (const section of sections) {
    if (section.label) {
      result.push({
        node: {} as StatementTreeNode,
        sectionKey: section.key,
        sectionLabel: section.label,
        isSectionHeader: true
      });
    }
    flattenNodes(section.nodes, section.key);
  }

  return result;
}

function renderNodes(props: StatementMatrixProps & { nodes: StatementTreeNode[]; dense?: boolean }) {
  const { dense = false } = props;
  const buttonSize = dense ? 'size-7' : 'size-11';
  const buttonClass = dense ? 'rounded' : 'rounded-lg';
  const labelSize = dense ? 'text-[13px]' : 'text-sm';
  const detailLabelSize = dense ? 'text-xs' : 'text-[13px]';
  const paddingY = dense ? 'py-1.5' : 'py-2';
  const gapClass = dense ? 'gap-1' : 'gap-2';

  return props.nodes.map((node) => {
    const isSelected = rowSelected(node, props.selectedRowRef);
    const labelIndent = node.kind === 'detail' ? node.level * 16 + 16 : node.level * 16;
    const canToggle = isSurfaceNode(node) && node.expandable;
    const nextSelection: StatementInspectorSelection = node.kind === 'surface'
      ? { kind: 'surface', key: node.row.key }
      : { kind: 'detail', key: node.row.key, parentKey: node.parentSurfaceKey };

    return (
      <Fragment key={node.id}>
        <tr className={cn(isSelected && 'bg-[color:rgba(70,77,87,0.48)]')}>
          <td className="sticky left-0 z-10 bg-[color:var(--panel)]">
            <div className={cn('flex min-w-[240px] items-start', gapClass)} style={{ paddingLeft: `${labelIndent}px` }}>
              {canToggle ? (
                <button
                  type="button"
                  aria-label={`${node.expanded ? 'Collapse' : 'Expand'} ${node.row.label} details`}
                  aria-expanded={node.expanded}
                  aria-controls={`statement-children-${node.id}`}
                  className={cn(
                    'mt-0.5 inline-flex shrink-0 items-center justify-center border border-[color:var(--line-weak)] bg-[color:var(--panel-soft)] text-[color:var(--terminal-bright)] transition hover:border-[color:var(--line-strong)]',
                    buttonSize,
                    buttonClass
                  )}
                  onClick={(event) => {
                    event.stopPropagation();
                    props.onToggleRow(node.row.key);
                  }}
                >
                  {node.expanded ? <ChevronDown className="size-3.5" /> : <ChevronRight className="size-3.5" />}
                </button>
              ) : (
                <span className={cn('inline-flex shrink-0 items-center justify-center text-[color:var(--terminal-muted)]', buttonSize)} aria-hidden="true">
                  {node.kind === 'detail' ? '·' : ''}
                </span>
              )}
              <button
                type="button"
                className={cn('flex min-w-0 flex-1 flex-col items-start gap-0.5 text-left', paddingY)}
                onClick={() => props.onSelectRow(nextSelection)}
              >
                <span className={cn(
                  labelSize,
                  'text-[color:var(--terminal-bright)]',
                  node.kind === 'detail' && cn(detailLabelSize, 'text-[color:var(--terminal-soft)]'),
                  node.kind === 'surface' && node.level > 0 && 'text-[color:var(--terminal-soft)]'
                )}>
                  {node.row.label}
                </span>
                {node.kind === 'detail' ? (
                  <span className="text-[10px] text-[color:var(--terminal-muted)]">
                    {node.row.localName}
                    {node.row.residualFlag ? ' · residual' : ''}
                  </span>
                ) : (
                  <div className="flex flex-wrap gap-1">
                    {surfaceBadges(node).map((badge) => (
                      <span
                        key={`${node.row.key}-${badge.label}`}
                        className={badgeClass(badge.tone, dense)}
                      >
                        {badge.label}
                      </span>
                    ))}
                  </div>
                )}
              </button>
            </div>
          </td>
          {props.periods.map((period, index) => (
            <td key={`${node.id}-${period.id}`} className="font-mono text-xs">
              {props.renderCellValue(node.row, period.id, index > 0 ? props.periods[index - 1]?.id ?? null : null)}
            </td>
          ))}
        </tr>
        {isSurfaceNode(node) && node.expanded ? (
          <>
            <tr id={`statement-children-${node.id}`} className="sr-only">
              <td colSpan={props.periods.length + 1}>Expanded children for {node.row.label}</td>
            </tr>
            {renderNodes({
              ...props,
              nodes: node.children,
              dense
            })}
          </>
        ) : null}
      </Fragment>
    );
  });
}

export function StatementMatrix(props: StatementMatrixProps) {
  const { dense = false, virtualized = false } = props;
  const tableClass = dense ? 'data-table-dense min-w-[960px]' : 'data-table min-w-[1040px]';

  // Hooks must be called unconditionally
  const parentRef = useRef<HTMLDivElement>(null);
  const flatRows = useMemo(() => flattenSections(props.sections), [props.sections]);

  const virtualizer = useVirtualizer({
    count: flatRows.length,
    getScrollElement: () => parentRef.current,
    estimateSize: (index) => {
      const item = flatRows[index];
      if (item.isSectionHeader) return dense ? 32 : 40;
      if (item.node.kind === 'surface' && surfaceBadges(item.node as any).length > 0) return dense ? 44 : 56;
      return dense ? 36 : 48;
    },
    overscan: 10,
  });

  // Non-virtualized version (original implementation with dense support)
  if (!virtualized) {
    return (
      <div className="data-table-wrap">
        <table className={tableClass}>
          <thead>
            <tr>
              <th className="sticky left-0 z-10 bg-[color:var(--panel)]">Metric</th>
              {props.periods.map((period) => (
                <th key={period.id}>
                  <div className="flex flex-col gap-1">
                    <span>{props.periodLabelFormatter(period.periodEnd ?? period.filingDate)}</span>
                    <span className="text-[11px] normal-case tracking-normal text-[color:var(--terminal-muted)]">{period.filingType} · {period.periodLabel}</span>
                  </div>
                </th>
              ))}
            </tr>
          </thead>
          <tbody>
            {props.sections.map((section) => (
              <Fragment key={section.key}>
                {section.label ? (
                  <tr className="bg-[color:var(--panel-soft)]">
                    <td colSpan={props.periods.length + 1} className="font-semibold text-[color:var(--terminal-bright)]">
                      {section.label}
                    </td>
                  </tr>
                ) : null}
                {renderNodes({
                  ...props,
                  nodes: section.nodes,
                  dense
                })}
              </Fragment>
            ))}
          </tbody>
        </table>
      </div>
    );
  }

  // Virtualized version for large datasets

  return (
    <div ref={parentRef} className="data-table-wrap h-[600px] overflow-auto">
      <table className={tableClass}>
        <thead className="sticky top-0 z-20 bg-[color:var(--panel)]">
          <tr>
            <th className="sticky left-0 z-30 bg-[color:var(--panel)]">Metric</th>
            {props.periods.map((period) => (
              <th key={period.id}>
                <div className="flex flex-col gap-1">
                  <span>{props.periodLabelFormatter(period.periodEnd ?? period.filingDate)}</span>
                  <span className="text-[11px] normal-case tracking-normal text-[color:var(--terminal-muted)]">{period.filingType} · {period.periodLabel}</span>
                </div>
              </th>
            ))}
          </tr>
        </thead>
        <tbody
          style={{ height: `${virtualizer.getTotalSize()}px` }}
          className="relative"
        >
          {virtualizer.getVirtualItems().map((virtualRow) => {
            const item = flatRows[virtualRow.index];

            if (item.isSectionHeader) {
              return (
                <tr
                  key={`section-${item.sectionKey}`}
                  className="absolute w-full bg-[color:var(--panel-soft)]"
                  style={{ transform: `translateY(${virtualRow.start}px)` }}
                >
                  <td colSpan={props.periods.length + 1} className="font-semibold text-[color:var(--terminal-bright)]">
                    {item.sectionLabel}
                  </td>
                </tr>
              );
            }

            const node = item.node;
            const isSelected = rowSelected(node, props.selectedRowRef);
            const labelIndent = node.kind === 'detail' ? node.level * 16 + 16 : node.level * 16;
            const canToggle = isSurfaceNode(node) && node.expandable;
            const nextSelection: StatementInspectorSelection = node.kind === 'surface'
              ? { kind: 'surface', key: node.row.key }
              : { kind: 'detail', key: node.row.key, parentKey: node.parentSurfaceKey };

            const buttonSize = dense ? 'size-7' : 'size-11';
            const buttonClass = dense ? 'rounded' : 'rounded-lg';
            const labelSize = dense ? 'text-[13px]' : 'text-sm';
            const detailLabelSize = dense ? 'text-xs' : 'text-[13px]';
            const paddingY = dense ? 'py-1.5' : 'py-2';
            const gapClass = dense ? 'gap-1' : 'gap-2';

            return (
              <tr
                key={node.id}
                className={cn('absolute w-full', isSelected && 'bg-[color:rgba(70,77,87,0.48)]')}
                style={{ transform: `translateY(${virtualRow.start}px)` }}
              >
                <td className="sticky left-0 z-10 bg-[color:var(--panel)]">
                  <div className={cn('flex min-w-[240px] items-start', gapClass)} style={{ paddingLeft: `${labelIndent}px` }}>
                    {canToggle ? (
                      <button
                        type="button"
                        aria-label={`${node.expanded ? 'Collapse' : 'Expand'} ${node.row.label} details`}
                        aria-expanded={node.expanded}
                        aria-controls={`statement-children-${node.id}`}
                        className={cn(
                          'mt-0.5 inline-flex shrink-0 items-center justify-center border border-[color:var(--line-weak)] bg-[color:var(--panel-soft)] text-[color:var(--terminal-bright)] transition hover:border-[color:var(--line-strong)]',
                          buttonSize,
                          buttonClass
                        )}
                        onClick={(event) => {
                          event.stopPropagation();
                          props.onToggleRow(node.row.key);
                        }}
                      >
                        {node.expanded ? <ChevronDown className="size-3.5" /> : <ChevronRight className="size-3.5" />}
                      </button>
                    ) : (
                      <span className={cn('inline-flex shrink-0 items-center justify-center text-[color:var(--terminal-muted)]', buttonSize)} aria-hidden="true">
                        {node.kind === 'detail' ? '·' : ''}
                      </span>
                    )}
                    <button
                      type="button"
                      className={cn('flex min-w-0 flex-1 flex-col items-start gap-0.5 text-left', paddingY)}
                      onClick={() => props.onSelectRow(nextSelection)}
                    >
                      <span className={cn(
                        labelSize,
                        'text-[color:var(--terminal-bright)]',
                        node.kind === 'detail' && cn(detailLabelSize, 'text-[color:var(--terminal-soft)]'),
                        node.kind === 'surface' && node.level > 0 && 'text-[color:var(--terminal-soft)]'
                      )}>
                        {node.row.label}
                      </span>
                      {node.kind === 'detail' ? (
                        <span className="text-[10px] text-[color:var(--terminal-muted)]">
                          {node.row.localName}
                          {node.row.residualFlag ? ' · residual' : ''}
                        </span>
                      ) : (
                        <div className="flex flex-wrap gap-1">
                          {surfaceBadges(node as any).map((badge) => (
                            <span
                              key={`${node.row.key}-${badge.label}`}
                              className={badgeClass(badge.tone, dense)}
                            >
                              {badge.label}
                            </span>
                          ))}
                        </div>
                      )}
                    </button>
                  </div>
                </td>
                {props.periods.map((period, index) => (
                  <td key={`${node.id}-${period.id}`} className="font-mono text-xs">
                    {props.renderCellValue(node.row, period.id, index > 0 ? props.periods[index - 1]?.id ?? null : null)}
                  </td>
                ))}
              </tr>
            );
          })}
        </tbody>
      </table>
    </div>
  );
}