import { describe, expect, it } from 'vitest';
import { FILTERED_GRAPH_SELECTION_ID, graphSelectionNodeIds, graphSelectionUnionNodeIds, relationBranchNodeIds } from '../../src/features/mitre-graph/category-branch';
import type { ObsidianGraphModel } from '../../src/features/mitre-graph/obsidian-graph';

const graph: ObsidianGraphModel = {
  nodes: [
    { id: 'category:a', kind: 'category', label: 'A', x: 0, y: 0, radius: 8 },
    { id: 'category:b', kind: 'category', label: 'B', x: 0, y: 0, radius: 8 },
    { id: 'finding:a', kind: 'finding', label: 'A finding', x: 0, y: 0, radius: 8 },
    { id: 'finding:b', kind: 'finding', label: 'B finding', x: 0, y: 0, radius: 8 },
    { id: 'finding:c', kind: 'finding', label: 'Direct tactic finding', x: 0, y: 0, radius: 8 },
    { id: 'technique:shared', kind: 'technique', label: 'Shared', x: 0, y: 0, radius: 8 },
    { id: 'target:a', kind: 'target', label: 'Target A', x: 0, y: 0, radius: 8 },
    { id: 'tactic:execution', kind: 'tactic', label: 'Execution', x: 0, y: 0, radius: 8 }
  ],
  edges: [
    { id: 'a-category', source: 'finding:a', target: 'category:a', kind: 'category' },
    { id: 'b-category', source: 'finding:b', target: 'category:b', kind: 'category' },
    { id: 'a-technique', source: 'finding:a', target: 'technique:shared', kind: 'technique' },
    { id: 'b-technique', source: 'finding:b', target: 'technique:shared', kind: 'technique' },
    { id: 'a-target', source: 'finding:a', target: 'target:a', kind: 'target' },
    { id: 'technique-tactic', source: 'technique:shared', target: 'tactic:execution', kind: 'technique' },
    { id: 'direct-tactic', source: 'finding:c', target: 'tactic:execution', kind: 'tactic' }
  ]
};

describe('category graph branch', () => {
  it('includes the selected category branch without escaping through shared MITRE nodes', () => {
    expect(relationBranchNodeIds(graph, 'category:a')).toEqual([
      'category:a', 'finding:a', 'technique:shared', 'target:a', 'tactic:execution'
    ]);
    expect(relationBranchNodeIds(graph, 'category:a')).not.toContain('finding:b');
    expect(relationBranchNodeIds(graph, 'category:a')).not.toContain('category:b');
  });

  it('supports shared MITRE technique and tactic roots without escaping recursively', () => {
    expect(relationBranchNodeIds(graph, 'technique:shared')).toEqual([
      'category:a', 'category:b', 'finding:a', 'finding:b', 'technique:shared', 'target:a', 'tactic:execution'
    ]);
    expect(relationBranchNodeIds(graph, 'tactic:execution')).toEqual([
      'category:a', 'category:b', 'finding:a', 'finding:b', 'finding:c', 'technique:shared', 'target:a', 'tactic:execution'
    ]);
  });

  it('keeps every same-kind endpoint attached to a matched finding', () => {
    const multiDefenseGraph: ObsidianGraphModel = {
      nodes: [
        { id: 'finding:a', kind: 'finding', label: 'Finding', x: 0, y: 0, radius: 8 },
        { id: 'defense:a', kind: 'defense', label: 'Defense A', x: 0, y: 0, radius: 8 },
        { id: 'defense:b', kind: 'defense', label: 'Defense B', x: 0, y: 0, radius: 8 },
        { id: 'target:a', kind: 'target', label: 'Target', x: 0, y: 0, radius: 8 }
      ],
      edges: [
        { id: 'defense-a', source: 'finding:a', target: 'defense:a', kind: 'defense' },
        { id: 'defense-b', source: 'finding:a', target: 'defense:b', kind: 'defense' },
        { id: 'target-a', source: 'finding:a', target: 'target:a', kind: 'target' }
      ]
    };

    expect(relationBranchNodeIds(multiDefenseGraph, 'defense:a')).toEqual([
      'finding:a', 'defense:a', 'defense:b', 'target:a'
    ]);
  });

  it('ignores finding roots', () => {
    expect(relationBranchNodeIds(graph, 'finding:a')).toEqual([]);
  });

  it('selects every node for filters without a dedicated relation node', () => {
    expect(graphSelectionNodeIds(graph, FILTERED_GRAPH_SELECTION_ID)).toEqual(graph.nodes.map(node => node.id));
  });

  it('combines multiple selected relation branches for extraction', () => {
    expect(graphSelectionUnionNodeIds(graph, ['category:a', 'category:b'])).toEqual([
      'category:a', 'category:b', 'finding:a', 'finding:b', 'technique:shared', 'target:a', 'tactic:execution'
    ]);
  });

  it('keeps a selected finding and its direct relations extractable', () => {
    expect(graphSelectionUnionNodeIds(graph, ['finding:a'])).toEqual([
      'category:a', 'finding:a', 'technique:shared', 'target:a'
    ]);
  });
});