Neon-Desk/lib/financials/statement-view-model.ts

import type {
  DetailFinancialRow,
  FinancialCategory,
  FinancialSurfaceKind,
  SurfaceDetailMap,
  SurfaceFinancialRow
} from '@/lib/types';

const SURFACE_CHILDREN: Partial<Record<Extract<FinancialSurfaceKind, 'income_statement' | 'balance_sheet' | 'cash_flow_statement'>, Record<string, string[]>>> = {
  income_statement: {
    operating_expenses: [
      'selling_general_and_administrative',
      'research_and_development',
      'other_operating_expense'
    ]
  }
};

export type StatementInspectorSelection = {
  kind: 'surface' | 'detail';
  key: string;
  parentKey?: string;
};

export type StatementTreeDetailNode = {
  kind: 'detail';
  id: string;
  level: number;
  row: DetailFinancialRow;
  parentSurfaceKey: string;
  matchesSearch: boolean;
};

export type StatementTreeSurfaceNode = {
  kind: 'surface';
  id: string;
  level: number;
  row: SurfaceFinancialRow;
  childSurfaceKeys: string[];
  directDetailCount: number;
  children: StatementTreeNode[];
  expandable: boolean;
  expanded: boolean;
  autoExpanded: boolean;
  matchesSearch: boolean;
};

export type StatementTreeNode = StatementTreeSurfaceNode | StatementTreeDetailNode;

export type StatementTreeSection = {
  key: string;
  label: string | null;
  nodes: StatementTreeNode[];
};

export type StatementTreeModel = {
  sections: StatementTreeSection[];
  autoExpandedKeys: Set<string>;
  visibleNodeCount: number;
  totalNodeCount: number;
};

export type ResolvedStatementSelection =
  | {
      kind: 'surface';
      row: SurfaceFinancialRow;
      childSurfaceRows: SurfaceFinancialRow[];
      detailRows: DetailFinancialRow[];
    }
  | {
      kind: 'detail';
      row: DetailFinancialRow;
      parentSurfaceRow: SurfaceFinancialRow | null;
    };

type Categories = Array<{
  key: FinancialCategory;
  label: string;
  count: number;
}>;

const UNMAPPED_DETAIL_GROUP_KEY = 'unmapped';
const UNMAPPED_SECTION_KEY = 'unmapped_residual';
const UNMAPPED_SECTION_LABEL = 'Unmapped / Residual';

function surfaceConfigForKind(surfaceKind: Extract<FinancialSurfaceKind, 'income_statement' | 'balance_sheet' | 'cash_flow_statement'>) {
  return SURFACE_CHILDREN[surfaceKind] ?? {};
}

function detailNodeId(parentKey: string, row: DetailFinancialRow) {
  return `detail:${parentKey}:${row.key}`;
}

function normalize(value: string) {
  return value.trim().toLowerCase();
}

function searchTextForSurface(row: SurfaceFinancialRow) {
  return [
    row.label,
    row.key,
    ...(row.sourceConcepts ?? []),
    ...(row.sourceRowKeys ?? []),
    ...(row.warningCodes ?? [])
  ]
    .join(' ')
    .toLowerCase();
}

function searchTextForDetail(row: DetailFinancialRow) {
  return [
    row.label,
    row.key,
    row.parentSurfaceKey,
    row.conceptKey,
    row.qname,
    row.localName,
    ...(row.dimensionsSummary ?? [])
  ]
    .join(' ')
    .toLowerCase();
}

function sortSurfaceRows(left: SurfaceFinancialRow, right: SurfaceFinancialRow) {
  if (left.order !== right.order) {
    return left.order - right.order;
  }

  return left.label.localeCompare(right.label);
}

function sortDetailRows(left: DetailFinancialRow, right: DetailFinancialRow) {
  return left.label.localeCompare(right.label);
}

function buildUnmappedDetailNodes(input: {
  statementDetails: SurfaceDetailMap | null;
  searchQuery: string;
}) {
  const normalizedSearch = normalize(input.searchQuery);

  return [...(input.statementDetails?.[UNMAPPED_DETAIL_GROUP_KEY] ?? [])]
    .sort(sortDetailRows)
    .filter((detail) => normalizedSearch.length === 0 || searchTextForDetail(detail).includes(normalizedSearch))
    .map((detail) => ({
      kind: 'detail',
      id: detailNodeId(UNMAPPED_DETAIL_GROUP_KEY, detail),
      level: 0,
      row: detail,
      parentSurfaceKey: UNMAPPED_DETAIL_GROUP_KEY,
      matchesSearch: normalizedSearch.length > 0 && searchTextForDetail(detail).includes(normalizedSearch)
    }) satisfies StatementTreeDetailNode);
}

function countNodes(nodes: StatementTreeNode[]) {
  let count = 0;

  for (const node of nodes) {
    count += 1;
    if (node.kind === 'surface') {
      count += countNodes(node.children);
    }
  }

  return count;
}

export function buildStatementTree(input: {
  surfaceKind: Extract<FinancialSurfaceKind, 'income_statement' | 'balance_sheet' | 'cash_flow_statement'>;
  rows: SurfaceFinancialRow[];
  statementDetails: SurfaceDetailMap | null;
  categories: Categories;
  searchQuery: string;
  expandedRowKeys: Set<string>;
}): StatementTreeModel {
  const config = surfaceConfigForKind(input.surfaceKind);
  const rowByKey = new Map(input.rows.map((row) => [row.key, row]));
  const childKeySet = new Set<string>();

  for (const children of Object.values(config)) {
    for (const childKey of children) {
      if (rowByKey.has(childKey)) {
        childKeySet.add(childKey);
      }
    }
  }

  const normalizedSearch = normalize(input.searchQuery);
  const autoExpandedKeys = new Set<string>();

  const buildSurfaceNode = (row: SurfaceFinancialRow, level: number): StatementTreeSurfaceNode | null => {
    const childSurfaceRows = (config[row.key] ?? [])
      .map((key) => rowByKey.get(key))
      .filter((candidate): candidate is SurfaceFinancialRow => Boolean(candidate))
      .sort(sortSurfaceRows);
    const detailRows = [...(input.statementDetails?.[row.key] ?? [])].sort(sortDetailRows);
    const childSurfaceNodes = childSurfaceRows
      .map((childRow) => buildSurfaceNode(childRow, level + 1))
      .filter((node): node is StatementTreeSurfaceNode => Boolean(node));
    const detailNodes = detailRows
      .filter((detail) => normalizedSearch.length === 0 || searchTextForDetail(detail).includes(normalizedSearch))
      .map((detail) => ({
        kind: 'detail',
        id: detailNodeId(row.key, detail),
        level: level + 1,
        row: detail,
        parentSurfaceKey: row.key,
        matchesSearch: normalizedSearch.length > 0 && searchTextForDetail(detail).includes(normalizedSearch)
      }) satisfies StatementTreeDetailNode);
    const children = [...childSurfaceNodes, ...detailNodes];
    const matchesSearch = normalizedSearch.length > 0 && searchTextForSurface(row).includes(normalizedSearch);
    const hasMatchingDescendant = normalizedSearch.length > 0 && children.length > 0;

    if (normalizedSearch.length > 0 && !matchesSearch && !hasMatchingDescendant) {
      return null;
    }

    const childSurfaceKeys = childSurfaceRows.map((candidate) => candidate.key);
    const directDetailCount = detailRows.length;
    const autoExpanded = normalizedSearch.length > 0 && !matchesSearch && children.length > 0;
    const expanded = children.length > 0 && (input.expandedRowKeys.has(row.key) || autoExpanded);

    if (autoExpanded) {
      autoExpandedKeys.add(row.key);
    }

    return {
      kind: 'surface',
      id: row.key,
      level,
      row,
      childSurfaceKeys,
      directDetailCount,
      children,
      expandable: children.length > 0,
      expanded,
      autoExpanded,
      matchesSearch
    };
  };

  const rootNodes = input.rows
    .filter((row) => !childKeySet.has(row.key))
    .sort(sortSurfaceRows)
    .map((row) => buildSurfaceNode(row, 0))
    .filter((node): node is StatementTreeSurfaceNode => Boolean(node));

  if (input.categories.length === 0) {
    const sections: StatementTreeSection[] = rootNodes.length > 0
      ? [{ key: 'ungrouped', label: null, nodes: rootNodes }]
      : [];
    const unmappedNodes = buildUnmappedDetailNodes({
      statementDetails: input.statementDetails,
      searchQuery: input.searchQuery
    });

    if (unmappedNodes.length > 0) {
      sections.push({
        key: UNMAPPED_SECTION_KEY,
        label: UNMAPPED_SECTION_LABEL,
        nodes: unmappedNodes
      });
    }

    return {
      sections,
      autoExpandedKeys,
      visibleNodeCount: sections.reduce((sum, section) => sum + countNodes(section.nodes), 0),
      totalNodeCount: input.rows.length + Object.values(input.statementDetails ?? {}).reduce((sum, rows) => sum + rows.length, 0)
    };
  }

  const sections: StatementTreeSection[] = [];
  const categoriesByKey = new Map(input.categories.map((category) => [category.key, category.label]));

  for (const category of input.categories) {
    const nodes = rootNodes.filter((node) => node.row.category === category.key);
    if (nodes.length > 0) {
      sections.push({
        key: category.key,
        label: category.label,
        nodes
      });
    }
  }

  const uncategorized = rootNodes.filter((node) => !categoriesByKey.has(node.row.category));
  if (uncategorized.length > 0) {
    sections.push({
      key: 'uncategorized',
      label: null,
      nodes: uncategorized
    });
  }

  const unmappedNodes = buildUnmappedDetailNodes({
    statementDetails: input.statementDetails,
    searchQuery: input.searchQuery
  });
  if (unmappedNodes.length > 0) {
    sections.push({
      key: UNMAPPED_SECTION_KEY,
      label: UNMAPPED_SECTION_LABEL,
      nodes: unmappedNodes
    });
  }

  return {
    sections,
    autoExpandedKeys,
    visibleNodeCount: sections.reduce((sum, section) => sum + countNodes(section.nodes), 0),
    totalNodeCount: input.rows.length + Object.values(input.statementDetails ?? {}).reduce((sum, rows) => sum + rows.length, 0)
  };
}

export function resolveStatementSelection(input: {
  surfaceKind: Extract<FinancialSurfaceKind, 'income_statement' | 'balance_sheet' | 'cash_flow_statement'>;
  rows: SurfaceFinancialRow[];
  statementDetails: SurfaceDetailMap | null;
  selection: StatementInspectorSelection | null;
}): ResolvedStatementSelection | null {
  const selection = input.selection;

  if (!selection) {
    return null;
  }

  const rowByKey = new Map(input.rows.map((row) => [row.key, row]));
  const config = surfaceConfigForKind(input.surfaceKind);

  if (selection.kind === 'surface') {
    const row = rowByKey.get(selection.key);
    if (!row) {
      return null;
    }

    const childSurfaceRows = (config[row.key] ?? [])
      .map((key) => rowByKey.get(key))
      .filter((candidate): candidate is SurfaceFinancialRow => Boolean(candidate))
      .sort(sortSurfaceRows);

    return {
      kind: 'surface',
      row,
      childSurfaceRows,
      detailRows: [...(input.statementDetails?.[row.key] ?? [])].sort(sortDetailRows)
    };
  }

  const parentSurfaceKey = selection.parentKey ?? null;
  const detailRows = parentSurfaceKey === UNMAPPED_DETAIL_GROUP_KEY
    ? input.statementDetails?.[UNMAPPED_DETAIL_GROUP_KEY] ?? []
    : parentSurfaceKey
      ? input.statementDetails?.[parentSurfaceKey] ?? []
      : Object.values(input.statementDetails ?? {}).flat();
  const row = detailRows.find((candidate) => candidate.key === selection.key) ?? null;

  if (!row) {
    return null;
  }

  return {
    kind: 'detail',
    row,
    parentSurfaceRow: rowByKey.get(row.parentSurfaceKey) ?? null
  };
}