import type { GraphNode, GraphLink, Domain } from "./domain-colors";
import { DOMAIN_COLORS } from "./domain-colors";

// ─── Layout Constants ───────────────────────────────────────────────────────

const LEVEL_Y_SPACING = 120;
const DOMAIN_X_SPACING = 200;
const NODE_Z_JITTER = 40;

// Domain → X position offset for clustering
const DOMAIN_X_OFFSETS: Partial<Record<Domain, number>> = {
  INPUT: 0,
  SAH: -300,
  ALLOCATION: -100,
  RULES: 100,
  CHARGEABILITY: 300,
  BUDGET: 500,
  ESTIMATE: -200,
  COMMERCIAL: 0,
  EXPERIENCE: 200,
  EFFORT: -400,
  SPREAD: 400,
};

// ─── Position Calculator ────────────────────────────────────────────────────

export interface PositionedNode extends GraphNode {
  fx: number;
  fy: number;
  fz: number;
  color: string;
}

export interface ForceGraphData {
  nodes: PositionedNode[];
  links: GraphLink[];
}

/**
 * Assigns 3D positions to nodes based on their level (Y) and domain (X),
 * with slight Z jitter to prevent overlap.
 */
export function buildForceGraphData(
  nodes: GraphNode[],
  links: GraphLink[],
): ForceGraphData {
  // Group nodes by domain+level to spread within each cluster
  const groups = new Map<string, GraphNode[]>();
  for (const node of nodes) {
    const key = `${node.domain}:${node.level}`;
    const arr = groups.get(key) ?? [];
    arr.push(node);
    groups.set(key, arr);
  }

  const positionedNodes: PositionedNode[] = nodes.map((node) => {
    const key = `${node.domain}:${node.level}`;
    const siblings = groups.get(key) ?? [node];
    const idx = siblings.indexOf(node);
    const count = siblings.length;

    // Center siblings around the domain X offset
    const baseX = DOMAIN_X_OFFSETS[node.domain] ?? 0;
    const spreadX = count > 1 ? (idx - (count - 1) / 2) * 80 : 0;

    return {
      ...node,
      fx: baseX + spreadX,
      fy: node.level * LEVEL_Y_SPACING,
      fz: (idx % 2 === 0 ? 1 : -1) * NODE_Z_JITTER * (Math.floor(idx / 2) + 1) * 0.5,
      color: DOMAIN_COLORS[node.domain],
    };
  });

  // Filter links to only reference existing node IDs
  const nodeIds = new Set(positionedNodes.map((n) => n.id));
  const validLinks = links.filter((link) => nodeIds.has(link.source) && nodeIds.has(link.target));

  return { nodes: positionedNodes, links: validLinks };
}

// ─── Highlight Helpers ──────────────────────────────────────────────────────

/**
 * Given a clicked node, returns the set of node IDs in its
 * upstream (ancestors) and downstream (descendants) path.
 */
export function getConnectedNodeIds(
  nodeId: string,
  links: GraphLink[],
): Set<string> {
  const connected = new Set<string>([nodeId]);

  // BFS upstream (nodes that feed into this one)
  const queue = [nodeId];
  while (queue.length > 0) {
    const current = queue.shift()!;
    for (const link of links) {
      const targetId = typeof link.target === "object" ? (link.target as { id: string }).id : link.target;
      const sourceId = typeof link.source === "object" ? (link.source as { id: string }).id : link.source;
      if (targetId === current && !connected.has(sourceId)) {
        connected.add(sourceId);
        queue.push(sourceId);
      }
    }
  }

  // BFS downstream (nodes this one feeds into)
  const queue2 = [nodeId];
  const visited = new Set<string>([nodeId]);
  while (queue2.length > 0) {
    const current = queue2.shift()!;
    for (const link of links) {
      const targetId = typeof link.target === "object" ? (link.target as { id: string }).id : link.target;
      const sourceId = typeof link.source === "object" ? (link.source as { id: string }).id : link.source;
      if (sourceId === current && !visited.has(targetId)) {
        connected.add(targetId);
        visited.add(targetId);
        queue2.push(targetId);
      }
    }
  }

  return connected;
}