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"
      | "equity_statement"
      | "disclosures"
    >,
    Record<string, string[]>
  >
> = {
  income_statement: {
    operating_expenses: [
      "selling_general_and_administrative",
      "research_and_development",
      "other_operating_expense",
    ],
  },
  balance_sheet: {
    total_assets: [
      "current_assets",
      "non_current_assets",
      "assets_held_for_sale",
    ],
    current_assets: [
      "cash_and_equivalents",
      "short_term_investments",
      "accounts_receivable",
      "inventory",
      "prepaid_expenses",
      "other_current_assets",
    ],
    non_current_assets: [
      "property_plant_equipment_net",
      "goodwill",
      "intangible_assets_net",
      "long_term_investments",
      "deferred_tax_assets",
      "other_non_current_assets",
    ],
    total_liabilities: [
      "current_liabilities",
      "non_current_liabilities",
      "liabilities_held_for_sale",
    ],
    current_liabilities: [
      "accounts_payable",
      "short_term_debt",
      "current_portion_of_long_term_debt",
      "accrued_liabilities",
      "deferred_revenue",
      "other_current_liabilities",
    ],
    non_current_liabilities: [
      "long_term_debt",
      "deferred_tax_liabilities",
      "deferred_revenue_non_current",
      "other_non_current_liabilities",
    ],
    total_equity: [
      "stockholders_equity",
      "minority_interest",
      "total_liabilities_and_equity",
    ],
    stockholders_equity: [
      "common_stock",
      "retained_earnings",
      "additional_paid_in_capital",
      "treasury_stock",
      "accumulated_other_comprehensive_income",
      "other_stockholders_equity",
    ],
  },
  cash_flow_statement: {
    net_cash_from_operating: [
      "operating_cash_flow_adjustments",
      "changes_in_working_capital",
    ],
    operating_cash_flow_adjustments: [
      "depreciation_and_amortization",
      "stock_based_compensation",
      "deferred_taxes",
      "other_non_cash_items",
    ],
    changes_in_working_capital: [
      "change_in_accounts_receivable",
      "change_in_inventory",
      "change_in_accounts_payable",
      "change_in_other_working_capital",
    ],
    net_cash_from_investing: [
      "capital_expenditures",
      "acquisitions",
      "purchases_of_investments",
      "sales_of_investments",
      "other_investing_activities",
    ],
    net_cash_from_financing: [
      "debt_issuance",
      "debt_repayment",
      "stock_issuance",
      "stock_repurchase",
      "dividends_paid",
      "other_financing_activities",
    ],
  },
  equity_statement: {},
  disclosures: {},
};

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[];
};

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"
    | "equity_statement"
    | "disclosures"
  >,
) {
  return SURFACE_CHILDREN[surfaceKind] ?? {};
}

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

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

function normalizeConceptIdentity(value: string | null | undefined) {
  if (!value) {
    return null;
  }

  const trimmed = value.trim().toLowerCase();
  if (trimmed.length === 0) {
    return null;
  }

  const withoutNamespace = trimmed.includes(":")
    ? (trimmed.split(":").pop() ?? trimmed)
    : trimmed;
  const normalized = withoutNamespace.replace(/[\s_-]+/g, "");

  return normalized.length > 0 ? normalized : null;
}

function surfaceConceptIdentities(row: SurfaceFinancialRow) {
  const identities = new Set<string>();

  const addIdentity = (value: string | null | undefined) => {
    const normalized = normalizeConceptIdentity(value);
    if (normalized) {
      identities.add(normalized);
    }
  };

  addIdentity(row.key);

  for (const sourceConcept of row.sourceConcepts ?? []) {
    addIdentity(sourceConcept);
  }

  for (const sourceRowKey of row.sourceRowKeys ?? []) {
    addIdentity(sourceRowKey);
  }

  for (const resolvedSourceRowKey of Object.values(
    row.resolvedSourceRowKeys ?? {},
  )) {
    addIdentity(resolvedSourceRowKey);
  }

  return identities;
}

function detailConceptIdentities(row: DetailFinancialRow) {
  const identities = new Set<string>();

  const addIdentity = (value: string | null | undefined) => {
    const normalized = normalizeConceptIdentity(value);
    if (normalized) {
      identities.add(normalized);
    }
  };

  addIdentity(row.key);
  addIdentity(row.conceptKey);
  addIdentity(row.qname);
  addIdentity(row.localName);

  return identities;
}

function dedupeDetailRowsAgainstChildSurfaces(
  detailRows: DetailFinancialRow[],
  childSurfaceRows: SurfaceFinancialRow[],
) {
  if (detailRows.length === 0 || childSurfaceRows.length === 0) {
    return detailRows;
  }

  const childConceptIdentities = new Set<string>();
  for (const childSurfaceRow of childSurfaceRows) {
    for (const identity of surfaceConceptIdentities(childSurfaceRow)) {
      childConceptIdentities.add(identity);
    }
  }

  if (childConceptIdentities.size === 0) {
    return detailRows;
  }

  return detailRows.filter((detailRow) => {
    for (const identity of detailConceptIdentities(detailRow)) {
      if (childConceptIdentities.has(identity)) {
        return false;
      }
    }

    return true;
  });
}

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"
    | "equity_statement"
    | "disclosures"
  >;
  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 = dedupeDetailRowsAgainstChildSurfaces(
      [...(input.statementDetails?.[row.key] ?? [])].sort(sortDetailRows),
      childSurfaceRows,
    );
    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));

  const totalVisibleDetailCount = input.rows.reduce((sum, row) => {
    const childSurfaceRows = (config[row.key] ?? [])
      .map((key) => rowByKey.get(key))
      .filter((candidate): candidate is SurfaceFinancialRow =>
        Boolean(candidate),
      );
    const detailRows = dedupeDetailRowsAgainstChildSurfaces(
      [...(input.statementDetails?.[row.key] ?? [])].sort(sortDetailRows),
      childSurfaceRows,
    );

    return sum + detailRows.length;
  }, 0);
  const unmappedDetailCount =
    input.statementDetails?.[UNMAPPED_DETAIL_GROUP_KEY]?.length ?? 0;

  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 + totalVisibleDetailCount + unmappedDetailCount,
    };
  }

  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 + totalVisibleDetailCount + unmappedDetailCount,
  };
}

export function resolveStatementSelection(input: {
  surfaceKind: Extract<
    FinancialSurfaceKind,
    | "income_statement"
    | "balance_sheet"
    | "cash_flow_statement"
    | "equity_statement"
    | "disclosures"
  >;
  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);
    const detailRows = dedupeDetailRowsAgainstChildSurfaces(
      [...(input.statementDetails?.[row.key] ?? [])].sort(sortDetailRows),
      childSurfaceRows,
    );

    return {
      kind: "surface",
      row,
      childSurfaceRows,
      detailRows,
    };
  }

  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,
  };
}