import { describe, it, expect } from 'vitest';
import type { AnalyzeResult, Node, Edge, StatementMeta } from '@pondpilot/flowscope-core';
import {
  buildFlowEdges,
  buildFlowNodes,
  mergeAnalyzeResult,
  computeIsCollapsed,
  buildScriptLevelGraph,
  type MergedLineage,
} from '../src/utils/graphBuilders';
import { GRAPH_CONFIG } from '../src/constants';

/**
 * Test fixture shape mirroring the legacy per-statement lineage — statement
 * metadata plus the nodes/edges that participate in that statement. `toResult`
 * flattens these fixtures into the shared `AnalyzeResult` shape consumed by
 * the refactored graph builders.
 */
type TestStatement = StatementMeta & { nodes: Node[]; edges: Edge[] };

function toResult(statements: TestStatement[]): AnalyzeResult {
  const nodes: Node[] = [];
  const edges: Edge[] = [];
  const meta: StatementMeta[] = [];

  // Preserve per-statement node/edge instances verbatim: the flat analyzer
  // output usually dedupes same-id entries across statements, but tests that
  // exercise the merge logic intentionally emit distinct Node objects per
  // statement to carry statement-specific spans. See occurrenceCycling tests.
  for (const stmt of statements) {
    const { nodes: stmtNodes, edges: stmtEdges, ...m } = stmt;
    meta.push(m);
    for (const node of stmtNodes) {
      nodes.push({ ...node, statementIds: [stmt.statementIndex] } as Node);
    }
    for (const edge of stmtEdges) {
      edges.push({ ...edge, statementIds: [stmt.statementIndex] } as Edge);
    }
  }

  return {
    statements: meta,
    nodes,
    edges,
    issues: [],
    summary: {
      statementCount: meta.length,
      tableCount: 0,
      columnCount: 0,
      joinCount: 0,
      complexityScore: 0,
      issueCount: { errors: 0, warnings: 0, infos: 0 },
      hasErrors: false,
    },
  };
}

function mergedFrom(statements: TestStatement[]): MergedLineage {
  return mergeAnalyzeResult(toResult(statements));
}

describe('computeIsCollapsed', () => {
  it('returns true when defaultCollapsed is true and node is not in overrides', () => {
    const overrides = new Set<string>();
    expect(computeIsCollapsed('node-1', true, overrides)).toBe(true);
  });

  it('returns false when defaultCollapsed is true and node is in overrides (expanded)', () => {
    const overrides = new Set(['node-1']);
    expect(computeIsCollapsed('node-1', true, overrides)).toBe(false);
  });

  it('returns false when defaultCollapsed is false and node is not in overrides', () => {
    const overrides = new Set<string>();
    expect(computeIsCollapsed('node-1', false, overrides)).toBe(false);
  });

  it('returns true when defaultCollapsed is false and node is in overrides (collapsed)', () => {
    const overrides = new Set(['node-1']);
    expect(computeIsCollapsed('node-1', false, overrides)).toBe(true);
  });
});

const createInsertLineage = (): TestStatement => ({
  statementIndex: 0,
  statementType: 'INSERT',
  joinCount: 0,
  complexityScore: 1,
  nodes: [
    {
      id: 'table:staging.orders',
      type: 'table',
      label: 'staging.orders',
      qualifiedName: 'staging.orders',
    },
    {
      id: 'table:analytics.tgt_orders',
      type: 'table',
      label: 'analytics.tgt_orders',
      qualifiedName: 'analytics.tgt_orders',
    },
    {
      id: 'column:staging.orders.order_id',
      type: 'column',
      label: 'order_id',
      qualifiedName: 'staging.orders.order_id',
    },
    {
      id: 'column:analytics.tgt_orders.order_id',
      type: 'column',
      label: 'order_id',
      qualifiedName: 'analytics.tgt_orders.order_id',
    },
    {
      // Simulates SELECT projection feeding the INSERT target (no qualified name)
      id: 'column:projection.order_id',
      type: 'column',
      label: 'order_id',
    },
  ],
  edges: [
    {
      id: 'edge:own:src',
      from: 'table:staging.orders',
      to: 'column:staging.orders.order_id',
      type: 'ownership',
    },
    {
      id: 'edge:own:tgt',
      from: 'table:analytics.tgt_orders',
      to: 'column:analytics.tgt_orders.order_id',
      type: 'ownership',
    },
    {
      id: 'edge:data:src_to_tgt',
      from: 'column:staging.orders.order_id',
      to: 'column:analytics.tgt_orders.order_id',
      type: 'data_flow',
    },
    {
      id: 'edge:der:projection',
      from: 'column:staging.orders.order_id',
      to: 'column:projection.order_id',
      type: 'derivation',
    },
  ],
});

describe('mergeAnalyzeResult', () => {
  it('preserves edge join metadata from later statements', () => {
    const firstStmt: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      nodes: [{ id: 'table:users', type: 'table', label: 'users', qualifiedName: 'users' }],
      edges: [],
      joinCount: 0,
      complexityScore: 1,
    };

    const secondStmt: TestStatement = {
      statementIndex: 1,
      statementType: 'SELECT',
      nodes: [
        {
          id: 'table:users',
          type: 'table',
          label: 'users',
          qualifiedName: 'users',
        },
      ],
      edges: [
        {
          id: 'edge:join',
          from: 'table:users',
          to: 'output',
          type: 'data_flow',
          joinType: 'LEFT',
          joinCondition: 'u.id = o.user_id',
        },
      ],
      joinCount: 1,
      complexityScore: 1,
    };

    const merged = mergeAnalyzeResult(toResult([firstStmt, secondStmt]));
    const joinEdge = merged.edges.find((e) => e.id === 'edge:join');
    expect(joinEdge?.joinType).toBe('LEFT');
    expect(joinEdge?.joinCondition).toBe('u.id = o.user_id');
  });

  it('merges occurrence spans and source names for same-id nodes', () => {
    const firstStmt: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      sourceName: 'models/a.sql',
      nodes: [
        {
          id: 'table:users',
          type: 'table',
          label: 'users',
          qualifiedName: 'users',
          nameSpans: [{ start: 5, end: 10 }],
        },
      ],
      edges: [],
      joinCount: 0,
      complexityScore: 1,
    };

    const secondStmt: TestStatement = {
      statementIndex: 1,
      statementType: 'SELECT',
      sourceName: 'models/b.sql',
      nodes: [
        {
          id: 'table:users',
          type: 'table',
          label: 'users',
          qualifiedName: 'users',
          nameSpans: [
            { start: 20, end: 25 },
            { start: 40, end: 45 },
          ],
        },
      ],
      edges: [],
      joinCount: 0,
      complexityScore: 1,
    };

    const merged = mergeAnalyzeResult(toResult([firstStmt, secondStmt]));
    const mergedNode = merged.nodes.find((node) => node.id === 'table:users');

    expect(mergedNode?.nameSpans).toEqual([
      { start: 5, end: 10 },
      { start: 20, end: 25 },
      { start: 40, end: 45 },
    ]);
    expect(mergedNode?.metadata?.occurrenceSourceNames).toEqual([
      'models/a.sql',
      'models/b.sql',
      'models/b.sql',
    ]);
  });

  it('does not duplicate occurrences when the flat result already shares a node across statements', () => {
    const result: AnalyzeResult = {
      statements: [
        {
          statementIndex: 0,
          statementType: 'SELECT',
          sourceName: 'models/a.sql',
          joinCount: 0,
          complexityScore: 1,
        },
        {
          statementIndex: 1,
          statementType: 'SELECT',
          sourceName: 'models/b.sql',
          joinCount: 0,
          complexityScore: 1,
        },
      ],
      nodes: [
        {
          id: 'table:users',
          type: 'table',
          label: 'users',
          qualifiedName: 'users',
          statementIds: [0, 1],
          nameSpans: [
            { start: 5, end: 10 },
            { start: 20, end: 25 },
          ],
        },
      ],
      edges: [],
      issues: [],
      summary: {
        statementCount: 2,
        tableCount: 1,
        columnCount: 0,
        joinCount: 0,
        complexityScore: 1,
        issueCount: { errors: 0, warnings: 0, infos: 0 },
        hasErrors: false,
      },
    };

    const merged = mergeAnalyzeResult(result);
    const mergedNode = merged.nodes.find((node) => node.id === 'table:users');

    expect(mergedNode?.nameSpans).toEqual([
      { start: 5, end: 10 },
      { start: 20, end: 25 },
    ]);
    expect(mergedNode?.metadata?.occurrenceSourceNames).toEqual(['models/a.sql', 'models/b.sql']);
  });
});

/**
 * Creates a statement lineage resembling the customer_360 view:
 * - CTEs: user_ltv (from orders), user_engagement (from session_summary)
 * - Final SELECT joins users with both CTEs
 * - Represents: CREATE VIEW customer_360 AS WITH ... SELECT ... FROM users LEFT JOIN user_ltv LEFT JOIN user_engagement
 */
const createCustomer360Lineage = (): TestStatement => ({
  statementIndex: 0,
  statementType: 'CREATE_VIEW',
  joinCount: 2,
  complexityScore: 50,
  nodes: [
    // Tables
    { id: 'table:orders', type: 'table', label: 'orders', qualifiedName: 'orders' },
    {
      id: 'table:session_summary',
      type: 'table',
      label: 'session_summary',
      qualifiedName: 'session_summary',
    },
    { id: 'table:users', type: 'table', label: 'users', qualifiedName: 'users' },
    // CTEs
    {
      id: 'cte:user_ltv',
      type: 'cte',
      label: 'user_ltv',
      joinType: 'LEFT',
      joinCondition: 'u.user_id = ltv.user_id',
    },
    {
      id: 'cte:user_engagement',
      type: 'cte',
      label: 'user_engagement',
      joinType: 'LEFT',
      joinCondition: 'u.user_id = eng.user_id',
    },
    // View
    { id: 'view:customer_360', type: 'view', label: 'customer_360', qualifiedName: 'customer_360' },
    // Columns from orders
    {
      id: 'column:orders.user_id',
      type: 'column',
      label: 'user_id',
      qualifiedName: 'orders.user_id',
    },
    {
      id: 'column:orders.total_amount',
      type: 'column',
      label: 'total_amount',
      qualifiedName: 'orders.total_amount',
    },
    // Columns from session_summary
    {
      id: 'column:session_summary.user_id',
      type: 'column',
      label: 'user_id',
      qualifiedName: 'session_summary.user_id',
    },
    {
      id: 'column:session_summary.session_id',
      type: 'column',
      label: 'session_id',
      qualifiedName: 'session_summary.session_id',
    },
    // Columns from users
    {
      id: 'column:users.user_id',
      type: 'column',
      label: 'user_id',
      qualifiedName: 'users.user_id',
    },
    { id: 'column:users.email', type: 'column', label: 'email', qualifiedName: 'users.email' },
    // Columns from user_ltv CTE
    {
      id: 'column:user_ltv.user_id',
      type: 'column',
      label: 'user_id',
      qualifiedName: 'user_ltv.user_id',
    },
    {
      id: 'column:user_ltv.lifetime_value',
      type: 'column',
      label: 'lifetime_value',
      qualifiedName: 'user_ltv.lifetime_value',
    },
    // Columns from user_engagement CTE
    {
      id: 'column:user_engagement.user_id',
      type: 'column',
      label: 'user_id',
      qualifiedName: 'user_engagement.user_id',
    },
    {
      id: 'column:user_engagement.total_sessions',
      type: 'column',
      label: 'total_sessions',
      qualifiedName: 'user_engagement.total_sessions',
    },
    // Columns from customer_360 view (output)
    {
      id: 'column:customer_360.user_id',
      type: 'column',
      label: 'user_id',
      qualifiedName: 'customer_360.user_id',
    },
    {
      id: 'column:customer_360.email',
      type: 'column',
      label: 'email',
      qualifiedName: 'customer_360.email',
    },
    {
      id: 'column:customer_360.lifetime_value',
      type: 'column',
      label: 'lifetime_value',
      qualifiedName: 'customer_360.lifetime_value',
    },
    {
      id: 'column:customer_360.total_sessions',
      type: 'column',
      label: 'total_sessions',
      qualifiedName: 'customer_360.total_sessions',
    },
  ],
  edges: [
    // Ownership edges: table -> column
    {
      id: 'own:orders.user_id',
      from: 'table:orders',
      to: 'column:orders.user_id',
      type: 'ownership',
    },
    {
      id: 'own:orders.total_amount',
      from: 'table:orders',
      to: 'column:orders.total_amount',
      type: 'ownership',
    },
    {
      id: 'own:session_summary.user_id',
      from: 'table:session_summary',
      to: 'column:session_summary.user_id',
      type: 'ownership',
    },
    {
      id: 'own:session_summary.session_id',
      from: 'table:session_summary',
      to: 'column:session_summary.session_id',
      type: 'ownership',
    },
    { id: 'own:users.user_id', from: 'table:users', to: 'column:users.user_id', type: 'ownership' },
    { id: 'own:users.email', from: 'table:users', to: 'column:users.email', type: 'ownership' },
    {
      id: 'own:user_ltv.user_id',
      from: 'cte:user_ltv',
      to: 'column:user_ltv.user_id',
      type: 'ownership',
    },
    {
      id: 'own:user_ltv.lifetime_value',
      from: 'cte:user_ltv',
      to: 'column:user_ltv.lifetime_value',
      type: 'ownership',
    },
    {
      id: 'own:user_engagement.user_id',
      from: 'cte:user_engagement',
      to: 'column:user_engagement.user_id',
      type: 'ownership',
    },
    {
      id: 'own:user_engagement.total_sessions',
      from: 'cte:user_engagement',
      to: 'column:user_engagement.total_sessions',
      type: 'ownership',
    },
    {
      id: 'own:customer_360.user_id',
      from: 'view:customer_360',
      to: 'column:customer_360.user_id',
      type: 'ownership',
    },
    {
      id: 'own:customer_360.email',
      from: 'view:customer_360',
      to: 'column:customer_360.email',
      type: 'ownership',
    },
    {
      id: 'own:customer_360.lifetime_value',
      from: 'view:customer_360',
      to: 'column:customer_360.lifetime_value',
      type: 'ownership',
    },
    {
      id: 'own:customer_360.total_sessions',
      from: 'view:customer_360',
      to: 'column:customer_360.total_sessions',
      type: 'ownership',
    },
    // Data flow edges: orders -> user_ltv CTE
    {
      id: 'flow:orders.user_id->user_ltv.user_id',
      from: 'column:orders.user_id',
      to: 'column:user_ltv.user_id',
      type: 'derivation',
    },
    {
      id: 'flow:orders.total_amount->user_ltv.lifetime_value',
      from: 'column:orders.total_amount',
      to: 'column:user_ltv.lifetime_value',
      type: 'derivation',
    },
    // Data flow edges: session_summary -> user_engagement CTE
    {
      id: 'flow:session_summary.user_id->user_engagement.user_id',
      from: 'column:session_summary.user_id',
      to: 'column:user_engagement.user_id',
      type: 'derivation',
    },
    {
      id: 'flow:session_summary.session_id->user_engagement.total_sessions',
      from: 'column:session_summary.session_id',
      to: 'column:user_engagement.total_sessions',
      type: 'derivation',
    },
    // Data flow edges: users -> customer_360
    {
      id: 'flow:users.user_id->customer_360.user_id',
      from: 'column:users.user_id',
      to: 'column:customer_360.user_id',
      type: 'data_flow',
    },
    {
      id: 'flow:users.email->customer_360.email',
      from: 'column:users.email',
      to: 'column:customer_360.email',
      type: 'data_flow',
    },
    // Data flow edges: user_ltv -> customer_360
    {
      id: 'flow:user_ltv.lifetime_value->customer_360.lifetime_value',
      from: 'column:user_ltv.lifetime_value',
      to: 'column:customer_360.lifetime_value',
      type: 'data_flow',
    },
    // Data flow edges: user_engagement -> customer_360
    {
      id: 'flow:user_engagement.total_sessions->customer_360.total_sessions',
      from: 'column:user_engagement.total_sessions',
      to: 'column:customer_360.total_sessions',
      type: 'data_flow',
    },
  ],
});

describe('buildFlowEdges table consistency', () => {
  it('should produce same table-to-table pairs regardless of showColumnEdges', () => {
    const statement = createCustomer360Lineage();

    // Build edges in both modes
    const tableEdges = buildFlowEdges(mergedFrom([statement]), false);
    const columnEdges = buildFlowEdges(mergedFrom([statement]), true);

    // Extract unique table pairs from each (source->target)
    const tableModePairs = new Set(tableEdges.map((e) => `${e.source}->${e.target}`));
    const columnModePairs = new Set(columnEdges.map((e) => `${e.source}->${e.target}`));

    // Verify consistent table pairs
    expect(tableModePairs).toEqual(columnModePairs);

    // Also verify edge counts make sense
    // Table mode: deduplicated (1 edge per table pair)
    // Column mode: one edge per column connection
    expect(tableEdges.length).toBeLessThanOrEqual(columnEdges.length);

    // With 8 column-level data flows, column mode should have more edges
    // Table mode should have exactly 5 unique table pairs
    expect(tableEdges.length).toBe(5);
    expect(columnEdges.length).toBe(8);
  });

  it('should include expected table relationships for customer_360', () => {
    const statement = createCustomer360Lineage();
    const edges = buildFlowEdges(mergedFrom([statement]), false);

    const tablePairs = edges.map((e) => `${e.source}->${e.target}`);

    // Expected relationships based on SQL structure
    expect(tablePairs).toContain('table:orders->cte:user_ltv');
    expect(tablePairs).toContain('table:session_summary->cte:user_engagement');
    expect(tablePairs).toContain('table:users->view:customer_360');
    expect(tablePairs).toContain('cte:user_ltv->view:customer_360');
    expect(tablePairs).toContain('cte:user_engagement->view:customer_360');
  });

  it('retains join-only table edges in column view', () => {
    const statement: TestStatement = {
      statementIndex: 0,
      statementType: 'CREATE_VIEW',
      joinCount: 1,
      complexityScore: 1,
      nodes: [
        { id: 'view:report', type: 'view', label: 'report', qualifiedName: 'report' },
        { id: 'table:table1', type: 'table', label: 'table1', qualifiedName: 'table1' },
        {
          id: 'table:table2',
          type: 'table',
          label: 'table2',
          qualifiedName: 'table2',
          joinType: 'LEFT',
          joinCondition: 't1.a = t2.a',
        },
        { id: 'column:table1.a', type: 'column', label: 'a', qualifiedName: 'table1.a' },
        { id: 'column:table1.b', type: 'column', label: 'b', qualifiedName: 'table1.b' },
        { id: 'column:report.a', type: 'column', label: 'a' },
        { id: 'column:report.b', type: 'column', label: 'b' },
      ],
      edges: [
        { id: 'own:table1.a', from: 'table:table1', to: 'column:table1.a', type: 'ownership' },
        { id: 'own:table1.b', from: 'table:table1', to: 'column:table1.b', type: 'ownership' },
        { id: 'own:view.a', from: 'view:report', to: 'column:report.a', type: 'ownership' },
        { id: 'own:view.b', from: 'view:report', to: 'column:report.b', type: 'ownership' },
        {
          id: 'flow:table1->view',
          from: 'table:table1',
          to: 'view:report',
          type: 'data_flow',
        },
        {
          id: 'flow:table2->view',
          from: 'table:table2',
          to: 'view:report',
          type: 'data_flow',
          joinType: 'LEFT',
          joinCondition: 't1.a = t2.a',
        },
        {
          id: 'der:table1.a',
          from: 'column:table1.a',
          to: 'column:report.a',
          type: 'derivation',
        },
        {
          id: 'der:table1.b',
          from: 'column:table1.b',
          to: 'column:report.b',
          type: 'derivation',
        },
      ],
    };

    const columnEdges = buildFlowEdges(mergedFrom([statement]), true);
    const tablePairs = columnEdges.map((edge) => `${edge.source}->${edge.target}`);

    expect(tablePairs).toContain('table:table2->view:report');
  });
});

describe('column-lineage column filtering', () => {
  function createWideSelectLineage(): TestStatement {
    const wideColumns = Array.from({ length: 60 }, (_, index) => ({
      id: `column:wide.c${index + 1}`,
      type: 'column' as const,
      label: `c${index + 1}`,
      qualifiedName: `wide.c${index + 1}`,
    }));
    const outputColumns = ['c1', 'c2'].map((name) => ({
      id: `column:output.${name}`,
      type: 'column' as const,
      label: name,
    }));

    return {
      statementIndex: 0,
      statementType: 'SELECT',
      joinCount: 0,
      complexityScore: 1,
      nodes: [
        { id: 'table:wide', type: 'table', label: 'wide', qualifiedName: 'wide' },
        { id: 'output:1', type: 'output', label: 'Output' },
        ...wideColumns,
        ...outputColumns,
      ],
      edges: [
        ...wideColumns.map((column) => ({
          id: `own:${column.label}`,
          from: 'table:wide',
          to: column.id,
          type: 'ownership' as const,
        })),
        ...outputColumns.map((column) => ({
          id: `own:${column.label}:output`,
          from: 'output:1',
          to: column.id,
          type: 'ownership' as const,
        })),
        {
          id: 'flow:c1',
          from: 'column:wide.c1',
          to: 'column:output.c1',
          type: 'data_flow' as const,
        },
        {
          id: 'flow:c2',
          from: 'column:wide.c2',
          to: 'column:output.c2',
          type: 'data_flow' as const,
        },
      ],
    };
  }

  it('shows only connected columns for wide tables in column-lineage mode', () => {
    const statement = createWideSelectLineage();

    const merged = mergedFrom([statement]);
    const tableModeNodes = buildFlowNodes(merged, null, '', new Set<string>(), new Set<string>());
    const columnModeNodes = buildFlowNodes(
      merged,
      null,
      '',
      new Set<string>(),
      new Set<string>(),
      null,
      false,
      true
    );

    const tableModeWide = tableModeNodes.find((node) => node.id === 'table:wide');
    const columnModeWide = columnModeNodes.find((node) => node.id === 'table:wide');

    expect(tableModeWide?.data.columns).toHaveLength(60);
    expect(tableModeWide?.data.lineageHiddenColumnCount).toBeFalsy();

    expect(columnModeWide?.data.columns.map((column) => column.name)).toEqual(['c1', 'c2']);
    expect(columnModeWide?.data.lineageHiddenColumnCount).toBe(58);
  });

  it('keeps the selected hidden column visible so column selection still has an anchor', () => {
    const merged = mergedFrom([createWideSelectLineage()]);
    const nodes = buildFlowNodes(
      merged,
      'column:wide.c60',
      '',
      new Set<string>(),
      new Set<string>(),
      null,
      false,
      true
    );

    const wideNode = nodes.find((node) => node.id === 'table:wide');

    expect(wideNode?.data.columns.map((column) => column.name)).toEqual(['c1', 'c2', 'c60']);
    expect(wideNode?.data.lineageHiddenColumnCount).toBe(57);
  });

  it('keeps search-matching hidden columns visible so focus/search can still target them', () => {
    const merged = mergedFrom([createWideSelectLineage()]);
    const nodes = buildFlowNodes(
      merged,
      null,
      'c60',
      new Set<string>(),
      new Set<string>(),
      null,
      false,
      true
    );

    const wideNode = nodes.find((node) => node.id === 'table:wide');

    expect(wideNode?.data.columns.map((column) => column.name)).toEqual(['c1', 'c2', 'c60']);
    expect(wideNode?.data.lineageHiddenColumnCount).toBe(57);
  });

  it('drops wide-table columns when their only lineage connection collapses to a table edge', () => {
    const sourceColumns = Array.from({ length: 60 }, (_, index) => ({
      id: `column:source.c${index + 1}`,
      type: 'column' as const,
      label: `c${index + 1}`,
      qualifiedName: `source.c${index + 1}`,
    }));
    const targetColumns = Array.from({ length: 60 }, (_, index) => ({
      id: `column:target.c${index + 1}`,
      type: 'column' as const,
      label: `c${index + 1}`,
      qualifiedName: `target.c${index + 1}`,
    }));

    const statement: TestStatement = {
      statementIndex: 0,
      statementType: 'CREATE_VIEW',
      joinCount: 0,
      complexityScore: 1,
      nodes: [
        { id: 'table:source', type: 'table', label: 'source', qualifiedName: 'source' },
        { id: 'table:target', type: 'table', label: 'target', qualifiedName: 'target' },
        ...sourceColumns,
        ...targetColumns,
      ],
      edges: [
        ...sourceColumns.map((column) => ({
          id: `own:source:${column.label}`,
          from: 'table:source',
          to: column.id,
          type: 'ownership' as const,
        })),
        ...targetColumns.map((column) => ({
          id: `own:target:${column.label}`,
          from: 'table:target',
          to: column.id,
          type: 'ownership' as const,
        })),
        {
          id: 'flow:c1',
          from: 'column:source.c1',
          to: 'column:target.c1',
          type: 'data_flow' as const,
        },
      ],
    };

    const nodes = buildFlowNodes(
      mergedFrom([statement]),
      null,
      '',
      new Set<string>(['table:target']),
      new Set<string>(),
      null,
      false,
      true
    );

    const sourceNode = nodes.find((node) => node.id === 'table:source');

    expect(sourceNode?.data.columns).toEqual([]);
    expect(sourceNode?.data.lineageHiddenColumnCount).toBe(60);
  });
});

describe('graphBuilders DML handling', () => {
  it('renders INSERT lineage into the real target even with unqualified columns present', () => {
    const statement = createInsertLineage();

    const flowEdges = buildFlowEdges(mergedFrom([statement]));
    expect(flowEdges).toHaveLength(1);
    expect(flowEdges[0]).toMatchObject({
      source: 'table:staging.orders',
      target: 'table:analytics.tgt_orders',
    });

    const flowNodes = buildFlowNodes(
      mergedFrom([statement]),
      null,
      '',
      new Set<string>(),
      new Set<string>()
    );
    const outputNode = flowNodes.find((node) => node.id === GRAPH_CONFIG.VIRTUAL_OUTPUT_NODE_ID);
    expect(outputNode).toBeUndefined();
  });

  it('keeps SELECT-style output edges even when other statements introduce table-level edges', () => {
    const statement: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      joinCount: 0,
      complexityScore: 1,
      nodes: [
        { id: 'table:source', type: 'table', label: 'source', qualifiedName: 'source' },
        { id: 'table:target', type: 'table', label: 'target', qualifiedName: 'target' },
        { id: 'column:source.id', type: 'column', label: 'id', qualifiedName: 'source.id' },
        { id: 'column:target.id', type: 'column', label: 'id', qualifiedName: 'target.id' },
        { id: 'column:output.total', type: 'column', label: 'total' },
      ],
      edges: [
        {
          id: 'own:source',
          from: 'table:source',
          to: 'column:source.id',
          type: 'ownership',
        },
        {
          id: 'own:target',
          from: 'table:target',
          to: 'column:target.id',
          type: 'ownership',
        },
        {
          id: 'flow:source_to_target',
          from: 'column:source.id',
          to: 'column:target.id',
          type: 'data_flow',
        },
        {
          id: 'flow:source_to_output',
          from: 'column:source.id',
          to: 'column:output.total',
          type: 'derivation',
        },
      ],
    };

    const edges = buildFlowEdges(mergedFrom([statement]));
    const dmlEdge = edges.find(
      (edge) => edge.source === 'table:source' && edge.target === 'table:target'
    );
    expect(dmlEdge, 'should keep DML-style edge').toBeDefined();

    const selectEdge = edges.find(
      (edge) =>
        edge.target === GRAPH_CONFIG.VIRTUAL_OUTPUT_NODE_ID && edge.source === 'table:source'
    );
    expect(selectEdge, 'should add SELECT output edge').toBeDefined();

    const nodes = buildFlowNodes(
      mergedFrom([statement]),
      null,
      '',
      new Set<string>(),
      new Set<string>()
    );
    const outputNode = nodes.find((node) => node.id === GRAPH_CONFIG.VIRTUAL_OUTPUT_NODE_ID);
    expect(outputNode, 'virtual Output node should exist for SELECT projections').toBeDefined();
  });

  it('renders relation-to-output-column lineage as a table-to-output edge', () => {
    const statement: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      joinCount: 1,
      complexityScore: 20,
      nodes: [
        { id: 'output:1', type: 'output', label: 'Output' },
        { id: 'table:users', type: 'table', label: 'users', qualifiedName: 'users' },
        {
          id: 'table:orders',
          type: 'table',
          label: 'orders',
          qualifiedName: 'orders',
        },
        { id: 'column:count', type: 'column', label: 'count', expression: 'COUNT(*)' },
      ],
      edges: [
        { id: 'own:output:count', from: 'output:1', to: 'column:count', type: 'ownership' },
        {
          id: 'der:users:count',
          from: 'table:users',
          to: 'column:count',
          type: 'derivation',
          expression: 'COUNT(*)',
        },
        {
          id: 'join:orders:output',
          from: 'table:orders',
          to: 'output:1',
          type: 'join_dependency',
          joinType: 'LEFT',
          joinCondition: 'u.id = o.user_id',
        },
      ],
    };

    const tableEdges = buildFlowEdges(mergedFrom([statement]), false);
    const columnEdges = buildFlowEdges(mergedFrom([statement]), true);

    expect(
      tableEdges.find((edge) => edge.source === 'table:users' && edge.target === 'output:1'),
      'table mode should surface the base table edge'
    ).toBeDefined();
    expect(
      columnEdges.find((edge) => edge.source === 'table:users' && edge.target === 'output:1'),
      'column mode should surface the base table edge'
    ).toBeDefined();
    expect(
      tableEdges.find((edge) => edge.source === 'table:orders' && edge.target === 'output:1'),
      'join-only dependency should still be preserved'
    ).toBeDefined();
  });

  it('keeps separate explicit output nodes when merged statements include multiple models', () => {
    const customersStmt: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      joinCount: 0,
      complexityScore: 1,
      sourceName: 'models/customers.sql',
      nodes: [
        { id: 'output:customers', type: 'output', label: 'customers' },
        {
          id: 'cte:stmt0:stg_customers',
          type: 'cte',
          label: 'stg_customers',
          qualifiedName: 'stg_customers',
        },
        {
          id: 'column:customers.id',
          type: 'column',
          label: 'customer_id',
        },
      ],
      edges: [
        {
          id: 'own:customers',
          from: 'output:customers',
          to: 'column:customers.id',
          type: 'ownership',
        },
      ],
    };

    const ordersStmt: TestStatement = {
      statementIndex: 1,
      statementType: 'SELECT',
      joinCount: 0,
      complexityScore: 1,
      sourceName: 'models/orders.sql',
      nodes: [
        { id: 'output:orders', type: 'output', label: 'orders' },
        {
          id: 'cte:stmt1:stg_customers',
          type: 'cte',
          label: 'stg_customers',
          qualifiedName: 'stg_customers',
        },
        {
          id: 'column:orders.id',
          type: 'column',
          label: 'order_id',
        },
      ],
      edges: [
        {
          id: 'own:orders',
          from: 'output:orders',
          to: 'column:orders.id',
          type: 'ownership',
        },
      ],
    };

    const merged = mergeAnalyzeResult(toResult([customersStmt, ordersStmt]));
    const nodes = buildFlowNodes(merged, null, '', new Set<string>(), new Set<string>());

    const outputNodes = nodes.filter((node) => node.data.nodeType === 'virtualOutput');
    const cteNodes = nodes.filter(
      (node) => node.id === 'cte:stmt0:stg_customers' || node.id === 'cte:stmt1:stg_customers'
    );

    expect(outputNodes.map((node) => node.id).sort()).toEqual([
      'output:customers',
      'output:orders',
    ]);
    expect(outputNodes.map((node) => node.data.label).sort()).toEqual(['customers', 'orders']);
    expect(cteNodes).toHaveLength(2);
  });

  it('preserves virtual output fallback when merged with an explicit-output statement', () => {
    const explicitStmt: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      joinCount: 0,
      complexityScore: 1,
      sourceName: 'models/explicit.sql',
      nodes: [
        { id: 'output:explicit', type: 'output', label: 'explicit_output' },
        {
          id: 'table:explicit_source',
          type: 'table',
          label: 'explicit_source',
          qualifiedName: 'explicit_source',
        },
        {
          id: 'column:explicit_source.id',
          type: 'column',
          label: 'id',
          qualifiedName: 'explicit_source.id',
        },
        {
          id: 'column:explicit.output_id',
          type: 'column',
          label: 'id',
        },
      ],
      edges: [
        {
          id: 'own:explicit:source',
          from: 'table:explicit_source',
          to: 'column:explicit_source.id',
          type: 'ownership',
        },
        {
          id: 'own:explicit:output',
          from: 'output:explicit',
          to: 'column:explicit.output_id',
          type: 'ownership',
        },
        {
          id: 'flow:explicit',
          from: 'column:explicit_source.id',
          to: 'column:explicit.output_id',
          type: 'data_flow',
        },
      ],
    };

    const fallbackStmt: TestStatement = {
      statementIndex: 1,
      statementType: 'SELECT',
      joinCount: 0,
      complexityScore: 1,
      sourceName: 'models/fallback.sql',
      nodes: [
        {
          id: 'table:fallback_source',
          type: 'table',
          label: 'fallback_source',
          qualifiedName: 'fallback_source',
        },
        {
          id: 'column:fallback_source.id',
          type: 'column',
          label: 'id',
          qualifiedName: 'fallback_source.id',
        },
        {
          id: 'column:fallback.output_id',
          type: 'column',
          label: 'id',
        },
      ],
      edges: [
        {
          id: 'own:fallback:source',
          from: 'table:fallback_source',
          to: 'column:fallback_source.id',
          type: 'ownership',
        },
        {
          id: 'flow:fallback',
          from: 'column:fallback_source.id',
          to: 'column:fallback.output_id',
          type: 'data_flow',
        },
      ],
    };

    const merged = mergeAnalyzeResult(toResult([explicitStmt, fallbackStmt]));
    const nodes = buildFlowNodes(merged, null, '', new Set<string>(), new Set<string>());
    const edges = buildFlowEdges(merged);

    expect(nodes.find((node) => node.id === 'output:explicit')).toBeDefined();
    expect(nodes.find((node) => node.id === GRAPH_CONFIG.VIRTUAL_OUTPUT_NODE_ID)).toBeDefined();
    expect(
      edges.find(
        (edge) =>
          edge.source === 'table:fallback_source' &&
          edge.target === GRAPH_CONFIG.VIRTUAL_OUTPUT_NODE_ID
      )
    ).toBeDefined();
  });

  it('does not create a virtual output node when an explicit output node lacks ownership edges', () => {
    const statement: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      joinCount: 0,
      complexityScore: 1,
      nodes: [
        { id: 'output:legacy', type: 'output', label: 'legacy_output' },
        { id: 'table:source', type: 'table', label: 'source', qualifiedName: 'source' },
        { id: 'column:source.id', type: 'column', label: 'id', qualifiedName: 'source.id' },
        { id: 'column:output.id', type: 'column', label: 'id' },
      ],
      edges: [
        {
          id: 'own:source:id',
          from: 'table:source',
          to: 'column:source.id',
          type: 'ownership',
        },
        {
          id: 'flow:source:output',
          from: 'column:source.id',
          to: 'column:output.id',
          type: 'data_flow',
        },
      ],
    };

    const nodes = buildFlowNodes(
      mergedFrom([statement]),
      null,
      '',
      new Set<string>(),
      new Set<string>()
    );
    const edges = buildFlowEdges(mergedFrom([statement]));

    expect(nodes.find((node) => node.id === 'output:legacy')).toBeDefined();
    expect(nodes.find((node) => node.id === GRAPH_CONFIG.VIRTUAL_OUTPUT_NODE_ID)).toBeUndefined();
    expect(
      edges.find((edge) => edge.target === GRAPH_CONFIG.VIRTUAL_OUTPUT_NODE_ID)
    ).toBeUndefined();
  });

  it('only marks physical tables as base tables when joins exist', () => {
    const statement: TestStatement = {
      statementIndex: 0,
      statementType: 'SELECT',
      nodes: [
        { id: 'table:users', type: 'table', label: 'users', qualifiedName: 'users' },
        { id: 'cte:recent_orders', type: 'cte', label: 'recent_orders' },
        { id: 'view:active_users', type: 'view', label: 'active_users' },
        {
          id: 'table:orders',
          type: 'table',
          label: 'orders',
          qualifiedName: 'orders',
        },
        {
          id: 'column:orders.total',
          type: 'column',
          label: 'total',
          qualifiedName: 'orders.total',
        },
        {
          id: 'output:1',
          type: 'output',
          label: 'Output',
        },
        {
          id: 'column:output.total',
          type: 'column',
          label: 'total',
        },
      ],
      edges: [
        {
          id: 'edge:orders:owns:total',
          from: 'table:orders',
          to: 'column:orders.total',
          type: 'ownership',
        },
        {
          id: 'edge:output:owns:total',
          from: 'output:1',
          to: 'column:output.total',
          type: 'ownership',
        },
        {
          id: 'edge:orders:total:output',
          from: 'column:orders.total',
          to: 'column:output.total',
          type: 'data_flow',
          joinType: 'INNER',
        },
      ],
      joinCount: 1,
      complexityScore: 1,
    };

    const nodes = buildFlowNodes(
      mergedFrom([statement]),
      null,
      '',
      new Set<string>(),
      new Set<string>()
    );
    const usersNode = nodes.find((node) => node.id === 'table:users');
    const ordersNode = nodes.find((node) => node.id === 'table:orders');
    const recentOrdersNode = nodes.find((node) => node.id === 'cte:recent_orders');
    const viewNode = nodes.find((node) => node.id === 'view:active_users');

    expect(usersNode?.data.isBaseTable).toBe(true);
    expect(ordersNode?.data.isBaseTable).toBe(false);
    expect(recentOrdersNode?.data.isBaseTable).toBeFalsy();
    expect(viewNode?.data.isBaseTable).toBeFalsy();
  });

  it('emits distinct synthetic edge ids for distinct relation pairs containing _to_', () => {
    const statement: TestStatement = {
      statementIndex: 0,
      statementType: 'INSERT',
      joinCount: 0,
      complexityScore: 1,
      nodes: [
        { id: 'table:a', type: 'table', label: 'a', qualifiedName: 'a' },
        { id: 'table:b_to_c', type: 'table', label: 'b_to_c', qualifiedName: 'b_to_c' },
        { id: 'table:a_to_b', type: 'table', label: 'a_to_b', qualifiedName: 'a_to_b' },
        { id: 'table:c', type: 'table', label: 'c', qualifiedName: 'c' },
        { id: 'column:a.id', type: 'column', label: 'id', qualifiedName: 'a.id' },
        {
          id: 'column:b_to_c.id',
          type: 'column',
          label: 'id',
          qualifiedName: 'b_to_c.id',
        },
        {
          id: 'column:a_to_b.id',
          type: 'column',
          label: 'id',
          qualifiedName: 'a_to_b.id',
        },
        { id: 'column:c.id', type: 'column', label: 'id', qualifiedName: 'c.id' },
      ],
      edges: [
        { id: 'own:a', from: 'table:a', to: 'column:a.id', type: 'ownership' },
        {
          id: 'own:b_to_c',
          from: 'table:b_to_c',
          to: 'column:b_to_c.id',
          type: 'ownership',
        },
        {
          id: 'own:a_to_b',
          from: 'table:a_to_b',
          to: 'column:a_to_b.id',
          type: 'ownership',
        },
        { id: 'own:c', from: 'table:c', to: 'column:c.id', type: 'ownership' },
        {
          id: 'flow:a:b_to_c',
          from: 'column:a.id',
          to: 'column:b_to_c.id',
          type: 'data_flow',
        },
        {
          id: 'flow:a_to_b:c',
          from: 'column:a_to_b.id',
          to: 'column:c.id',
          type: 'data_flow',
        },
      ],
    };

    const edges = buildFlowEdges(mergedFrom([statement]));
    const edgeIds = edges.map((edge) => edge.id);

    expect(edges).toHaveLength(2);
    expect(new Set(edgeIds).size).toBe(2);
    expect(edges.map((edge) => `${edge.source}->${edge.target}`).sort()).toEqual([
      'table:a->table:b_to_c',
      'table:a_to_b->table:c',
    ]);
  });

  it('includes explicit output nodes as written relations in script graph mode', () => {
    const statements: TestStatement[] = [
      {
        statementIndex: 0,
        statementType: 'WITH',
        sourceName: 'scratchpad.sql',
        joinCount: 0,
        complexityScore: 1,
        nodes: [
          {
            id: 'output:scratchpad',
            type: 'output',
            label: 'scratchpad',
            qualifiedName: 'scratchpad',
          },
          {
            id: 'table:raw_orders',
            type: 'table',
            label: 'raw_orders',
            qualifiedName: 'jaffle_shop.raw_orders',
          },
        ],
        edges: [],
      },
    ];

    const { nodes, edges } = buildScriptLevelGraph(toResult(statements), null, '', true);

    expect(nodes.find((node) => node.id === 'table:scratchpad')).toBeDefined();
    expect(
      edges.find(
        (edge) => edge.source === 'script:scratchpad.sql' && edge.target === 'table:scratchpad'
      )
    ).toBeDefined();
    expect(
      edges.find(
        (edge) =>
          edge.source === 'table:jaffle_shop.raw_orders' && edge.target === 'script:scratchpad.sql'
      )
    ).toBeDefined();
  });

  it('includes dbt relation sinks as written relations in script graph mode', () => {
    const statements: TestStatement[] = [
      {
        statementIndex: 0,
        statementType: 'SELECT',
        sourceName: 'models/constants.sql',
        joinCount: 0,
        complexityScore: 1,
        nodes: [
          {
            id: 'table:constants',
            type: 'table',
            label: 'constants',
            qualifiedName: 'constants',
          },
          {
            id: 'column:constants.id',
            type: 'column',
            label: 'id',
            // Simulate a merged producer/consumer column node that picked up a
            // source qualifiedName from another statement.
            qualifiedName: 'constants.id',
          },
        ],
        edges: [
          {
            id: 'own:constants.id',
            from: 'table:constants',
            to: 'column:constants.id',
            type: 'ownership',
          },
        ],
      },
    ];

    const { nodes, edges } = buildScriptLevelGraph(toResult(statements), null, '', true);

    expect(nodes.find((node) => node.id === 'table:constants')).toBeDefined();
    expect(
      edges.find(
        (edge) => edge.source === 'script:models/constants.sql' && edge.target === 'table:constants'
      )
    ).toBeDefined();
    expect(
      edges.find(
        (edge) => edge.source === 'table:constants' && edge.target === 'script:models/constants.sql'
      )
    ).toBeUndefined();
  });

  it('includes dbt ephemeral model sinks as written relations in script graph mode', () => {
    const statements: TestStatement[] = [
      {
        statementIndex: 0,
        statementType: 'SELECT',
        sourceName: 'models/ephemeral_constants.sql',
        joinCount: 0,
        complexityScore: 1,
        nodes: [
          {
            id: 'cte:ephemeral_constants',
            type: 'cte',
            label: 'ephemeral_constants',
            qualifiedName: 'analytics.ephemeral_constants',
            metadata: { dbtModelSink: true },
          },
          {
            id: 'column:ephemeral_constants.id',
            type: 'column',
            label: 'id',
            qualifiedName: 'analytics.ephemeral_constants.id',
          },
        ],
        edges: [
          {
            id: 'own:ephemeral_constants.id',
            from: 'cte:ephemeral_constants',
            to: 'column:ephemeral_constants.id',
            type: 'ownership',
          },
        ],
      },
    ];

    const { nodes, edges } = buildScriptLevelGraph(toResult(statements), null, '', true);

    expect(nodes.find((node) => node.id === 'table:analytics.ephemeral_constants')).toBeDefined();
    expect(
      edges.find(
        (edge) =>
          edge.source === 'script:models/ephemeral_constants.sql' &&
          edge.target === 'table:analytics.ephemeral_constants'
      )
    ).toBeDefined();
    expect(
      edges.find(
        (edge) =>
          edge.source === 'table:analytics.ephemeral_constants' &&
          edge.target === 'script:models/ephemeral_constants.sql'
      )
    ).toBeUndefined();
  });
});
