import type { Edge, Node } from "@xyflow/react";
import ELK from "elkjs/lib/elk.bundled.js";
import type {
  DagNodeData,
  GraphEdge,
  GraphNode,
  JobGraphNode,
  PlanGraph,
  TaskGraphNode,
} from "../types/graph-types.ts";
import type { LayoutResult } from "../types/layout-result.ts";
import {
  getEdgeStyle,
  LATERAL_BOTTOM,
  LATERAL_BOTTOM_OUT,
  LATERAL_EDGE_STYLE,
  LATERAL_TOP,
  LATERAL_TOP_OUT,
} from "../utils/diff-state-styles.ts";

/** Lazily instantiate ELK — deferred to avoid Worker creation at import time (breaks Bun test runner).
 *  No crash recovery: if the Worker dies the layout call rejects and the UI shows an error state,
 *  which is acceptable for a visualization tool (just reload). */
const getElk = (() => {
  let instance: InstanceType<typeof ELK> | undefined;
  return (): InstanceType<typeof ELK> => {
    if (!instance) {
      instance = new ELK();
    }
    return instance;
  };
})();

export const NODE_WIDTH = 220;
const NODE_HEIGHT_TASK = 56;
const NODE_HEIGHT_RESOURCE = 56;
const NODE_HEIGHT_GROUP = 56;

const JOB_PADDING_TOP = 40;
const JOB_PADDING_SIDE = 20;
const JOB_PADDING_BOTTOM = 20;

type JobGroup = {
  readonly job: JobGraphNode;
  readonly tasks: readonly TaskGraphNode[];
};

/** Group graph nodes by job, pairing each job with its child tasks. */
export const groupNodesByJob = (nodes: readonly GraphNode[]): readonly JobGroup[] => {
  const jobs = nodes.filter((node): node is JobGraphNode => node.nodeKind === "job");
  const tasksByResource = Map.groupBy(
    nodes.filter((node): node is TaskGraphNode => node.nodeKind === "task"),
    (task) => task.resourceKey,
  );
  return jobs.map((job) => ({
    job,
    tasks: tasksByResource.get(job.resourceKey) ?? [],
  }));
};

/** Extract the parent job ID from a node ID (task IDs contain "::", job IDs don't). */
const parentJobId = (nodeId: string): string => {
  const separator = nodeId.indexOf("::");
  return separator === -1 ? nodeId : nodeId.substring(0, separator);
};

/** Collect cross-hierarchy edges (source and target in different jobs) mapped to job-level for ELK. */
const collectCrossJobEdges = (
  edges: PlanGraph["edges"],
): { id: string; sources: string[]; targets: string[] }[] => {
  const seen = new Set<string>();
  return edges.flatMap((edge) => {
    const sourceJob = parentJobId(edge.source);
    const targetJob = parentJobId(edge.target);
    if (sourceJob === targetJob) return [];
    const pairKey = `${sourceJob}→${targetJob}`;
    if (seen.has(pairKey)) return [];
    seen.add(pairKey);
    return [{ id: `elk-${pairKey}`, sources: [sourceJob], targets: [targetJob] }];
  });
};

/** Topologically sort tasks within a job so ELK model order matches dependency flow.
 *  Uses Kahn's algorithm with local mutable state (adjacency map, in-degree map, queue)
 *  as a deliberate exception to immutability-by-default — the algorithm requires it. */
export const topologicalSortTasks = (
  tasks: readonly GraphNode[],
  edges: readonly { readonly source: string; readonly target: string }[],
): readonly GraphNode[] => {
  const taskIds = new Set(tasks.map((task) => task.id));
  const intraEdges = edges.filter((edge) => taskIds.has(edge.source) && taskIds.has(edge.target));

  const adjacency = new Map<string, string[]>();
  const inDegree = new Map<string, number>();

  for (const task of tasks) {
    adjacency.set(task.id, []);
    inDegree.set(task.id, 0);
  }

  for (const edge of intraEdges) {
    adjacency.get(edge.source)?.push(edge.target);
    inDegree.set(edge.target, (inDegree.get(edge.target) ?? 0) + 1);
  }

  const queue: string[] = [];
  for (const task of tasks) {
    if ((inDegree.get(task.id) ?? 0) === 0) {
      queue.push(task.id);
    }
  }

  const sortedIds: string[] = [];
  let front = 0;
  while (front < queue.length) {
    // bounds check above guarantees this element exists
    const current = queue[front++] as string;
    sortedIds.push(current);
    for (const neighbor of adjacency.get(current) ?? []) {
      const newDegree = (inDegree.get(neighbor) ?? 1) - 1;
      inDegree.set(neighbor, newDegree);
      if (newDegree === 0) {
        queue.push(neighbor);
      }
    }
  }

  const taskById = new Map(tasks.map((task) => [task.id, task]));
  const sorted = sortedIds.map((id) => taskById.get(id)).filter((task) => task !== undefined);

  // Append any disconnected tasks not reached by BFS in original order
  const sortedSet = new Set(sortedIds);
  for (const task of tasks) {
    if (!sortedSet.has(task.id)) {
      sorted.push(task);
    }
  }

  return sorted;
};

/** Build an ELK compound graph with jobs as parents and tasks as children. */
export const buildElkCompoundGraph = (groups: readonly JobGroup[], edges: PlanGraph["edges"]) => ({
  id: "root",
  layoutOptions: {
    "elk.algorithm": "layered",
    "elk.direction": "RIGHT",
    "elk.separateConnectedComponents": "false",
    "elk.spacing.nodeNode": "60",
    "elk.layered.spacing.nodeNodeBetweenLayers": "80",
    "elk.layered.considerModelOrder.strategy": "NODES_AND_EDGES",
    "elk.layered.crossingMinimization.forceNodeModelOrder": "true",
  },
  children: [...groups]
    .sort((a, b) => a.job.label.localeCompare(b.job.label))
    .map((group) => {
      const taskIds = new Set(group.tasks.map((task) => task.id));
      return {
        id: group.job.id,
        layoutOptions: {
          "elk.algorithm": "layered",
          "elk.direction": "RIGHT",
          "elk.hierarchyHandling": "SEPARATE_CHILDREN",
          "elk.spacing.nodeNode": "40",
          "elk.layered.spacing.nodeNodeBetweenLayers": "60",
          "elk.spacing.edgeNode": "20",
          "elk.layered.spacing.edgeNodeBetweenLayers": "20",
          "elk.layered.nodePlacement.strategy": "BRANDES_KOEPF",
          "elk.layered.nodePlacement.bk.fixedAlignment": "BALANCED",
          "elk.layered.considerModelOrder.portModelOrder": "true",
          "elk.layered.considerModelOrder.strategy": "NODES_AND_EDGES",
          "elk.layered.crossingMinimization.forceNodeModelOrder": "true",
          "elk.layered.compaction.connectedComponents": "true",
          "elk.layered.compaction.postCompaction.strategy": "EDGE_LENGTH",
          "elk.layered.cycleBreaking.strategy": "MODEL_ORDER",
          "elk.padding": `[top=${JOB_PADDING_TOP},left=${JOB_PADDING_SIDE},bottom=${JOB_PADDING_BOTTOM},right=${JOB_PADDING_SIDE}]`,
        },
        children: topologicalSortTasks(group.tasks, edges).map((task) => ({
          id: task.id,
          width: NODE_WIDTH,
          height: NODE_HEIGHT_TASK,
        })),
        edges: edges
          .filter((edge) => taskIds.has(edge.source) && taskIds.has(edge.target))
          .map((edge) => ({
            id: edge.id,
            sources: [edge.source],
            targets: [edge.target],
          })),
      };
    }),
  edges: collectCrossJobEdges(edges),
});

type ElkOutput = {
  readonly children?: ReadonlyArray<{
    readonly id: string;
    readonly x?: number;
    readonly y?: number;
    readonly width?: number;
    readonly height?: number;
    readonly children?: ReadonlyArray<{
      readonly id: string;
      readonly x?: number;
      readonly y?: number;
    }>;
  }>;
};

/** Extract layout positions and job dimensions from ELK result. */
export const extractLayoutData = (
  elkResult: ElkOutput,
): {
  readonly positions: ReadonlyMap<string, { readonly x: number; readonly y: number }>;
  readonly dimensions: ReadonlyMap<string, { readonly width: number; readonly height: number }>;
} => {
  const positions = new Map<string, { readonly x: number; readonly y: number }>();
  const dimensions = new Map<string, { readonly width: number; readonly height: number }>();

  for (const jobElk of elkResult.children ?? []) {
    positions.set(jobElk.id, { x: jobElk.x ?? 0, y: jobElk.y ?? 0 });
    dimensions.set(jobElk.id, {
      width: jobElk.width ?? 0,
      height: jobElk.height ?? 0,
    });

    for (const taskElk of jobElk.children ?? []) {
      positions.set(taskElk.id, { x: taskElk.x ?? 0, y: taskElk.y ?? 0 });
    }
  }

  return { positions, dimensions };
};

/** Build a descriptive aria-label for screen readers, e.g. "added task: ingest_raw_data". */
const buildAriaLabel = (node: GraphNode): string =>
  `${node.diffState} ${node.nodeKind}: ${node.label}`;

/** Strip the `id` field from a GraphNode, producing the data payload for React Flow nodes. */
const toNodeData = (node: GraphNode): DagNodeData => {
  const { id: _, ...data } = node;
  return data;
};

/** Convert a job GraphNode to a React Flow container node. */
const toJobFlowNode = (
  node: GraphNode,
  position: { readonly x: number; readonly y: number },
  dimension: { readonly width: number; readonly height: number },
): Node => ({
  id: node.id,
  type: node.nodeKind,
  position: { x: position.x, y: position.y },
  style: { width: dimension.width, height: dimension.height },
  data: toNodeData(node),
  ariaLabel: buildAriaLabel(node),
});

/** Convert a task GraphNode to a React Flow child node inside its job. */
const toTaskFlowNode = (
  node: GraphNode,
  position: { readonly x: number; readonly y: number },
): Node => ({
  id: node.id,
  type: node.nodeKind,
  position: { x: position.x, y: position.y },
  parentId: node.resourceKey,
  extent: "parent" as const,
  data: { ...toNodeData(node), taskChangeSummary: undefined },
  ariaLabel: buildAriaLabel(node),
});

/** Assemble React Flow nodes from layout data, with jobs before their children. */
export const assembleFlowNodes = (
  groups: readonly JobGroup[],
  positions: ReadonlyMap<string, { readonly x: number; readonly y: number }>,
  dimensions: ReadonlyMap<string, { readonly width: number; readonly height: number }>,
): readonly Node[] => {
  const nodes: Node[] = [];
  for (const group of groups) {
    const jobPosition = positions.get(group.job.id) ?? { x: 0, y: 0 };
    const jobDimension = dimensions.get(group.job.id) ?? { width: 0, height: 0 };
    nodes.push(toJobFlowNode(group.job, jobPosition, jobDimension));

    for (const task of group.tasks) {
      const taskPosition = positions.get(task.id) ?? { x: 0, y: 0 };
      nodes.push(toTaskFlowNode(task, taskPosition));
    }
  }
  return nodes;
};

/** Convert graph edges to React Flow edges with diff-state-derived colors. */
export const toFlowEdges = (edges: PlanGraph["edges"]): readonly Edge[] =>
  edges.map((edge) => {
    const style = getEdgeStyle(edge.diffState);
    return {
      id: edge.id,
      source: edge.source,
      target: edge.target,
      label: edge.label,
      type: "bezier",
      style: {
        stroke: style.stroke,
        opacity: style.opacity,
        strokeDasharray: style.strokeDasharray,
      },
    };
  });

/** Convert a PlanGraph to React Flow nodes and edges with ELK compound layout. */
export const toReactFlowElements = async (
  graph: PlanGraph,
): Promise<{ readonly nodes: readonly Node[]; readonly edges: readonly Edge[] }> => {
  if (graph.nodes.length === 0) {
    return { nodes: [], edges: [] };
  }

  const groups = groupNodesByJob(graph.nodes);
  const elkGraph = buildElkCompoundGraph(groups, graph.edges);
  const elkResult = await getElk().layout(elkGraph);
  const { positions, dimensions } = extractLayoutData(elkResult);

  return {
    nodes: assembleFlowNodes(groups, positions, dimensions),
    edges: toFlowEdges(graph.edges),
  };
};

/** Convert a GraphNode to a React Flow node at the given position (flat, no parent). */
const toFlatFlowNode = (
  node: GraphNode,
  position: { readonly x: number; readonly y: number },
): Node => ({
  id: node.id,
  type: node.nodeKind,
  position: { x: position.x, y: position.y },
  data: toNodeData(node),
  ariaLabel: buildAriaLabel(node),
});

type Position2D = { readonly x: number; readonly y: number };

/** Pick the top/bottom lateral handle pair that produces the shortest straight-line connection. */
const chooseLateralHandles = (
  sourcePos: Position2D,
  targetPos: Position2D,
  sourceHeight: number,
  targetHeight: number,
): { readonly sourceHandle: string; readonly targetHandle: string } => {
  const anchors = [
    { src: LATERAL_TOP_OUT, tgt: LATERAL_TOP, sy: 0, ty: 0 },
    { src: LATERAL_TOP_OUT, tgt: LATERAL_BOTTOM, sy: 0, ty: targetHeight },
    { src: LATERAL_BOTTOM_OUT, tgt: LATERAL_TOP, sy: sourceHeight, ty: 0 },
    { src: LATERAL_BOTTOM_OUT, tgt: LATERAL_BOTTOM, sy: sourceHeight, ty: targetHeight },
  ] as const;

  let best: (typeof anchors)[number] = anchors[0];
  let bestDist = Number.POSITIVE_INFINITY;
  for (const anchor of anchors) {
    const dx = sourcePos.x - targetPos.x;
    const dy = sourcePos.y + anchor.sy - (targetPos.y + anchor.ty);
    const dist = dx * dx + dy * dy;
    if (dist < bestDist) {
      bestDist = dist;
      best = anchor;
    }
  }
  return { sourceHandle: best.src, targetHandle: best.tgt };
};

type NodeLayout = { readonly position: Position2D; readonly height: number };

/** Convert lateral GraphEdges to React Flow edges with the lateral visual style.
 *  Routes each edge through the handle pair that produces the shortest straight line. */
export const toLateralFlowEdges = (
  edges: readonly GraphEdge[],
  nodeLayouts: ReadonlyMap<string, NodeLayout>,
): readonly Edge[] =>
  edges.flatMap((edge) => {
    const srcLayout = nodeLayouts.get(edge.source);
    const tgtLayout = nodeLayouts.get(edge.target);
    if (srcLayout === undefined || tgtLayout === undefined) return [];

    const handles = chooseLateralHandles(
      srcLayout.position,
      tgtLayout.position,
      srcLayout.height,
      tgtLayout.height,
    );

    return [
      {
        id: edge.id,
        source: edge.source,
        target: edge.target,
        sourceHandle: handles.sourceHandle,
        targetHandle: handles.targetHandle,
        type: "straight",
        zIndex: 10,
        style: {
          stroke: LATERAL_EDGE_STYLE.stroke,
          opacity: LATERAL_EDGE_STYLE.opacity,
          strokeDasharray: LATERAL_EDGE_STYLE.strokeDasharray,
          strokeWidth: 2.5,
        },
      },
    ];
  });

/** Flat ELK layout for resource graphs (left-to-right, no compound hierarchy). */
export const layoutResourceGraph = async (
  graph: PlanGraph & { readonly lateralEdges?: readonly GraphEdge[] },
): Promise<LayoutResult> => {
  if (graph.nodes.length === 0) {
    return { nodes: [], edges: [] };
  }

  const targetNodeIds = new Set(graph.edges.map((edge) => edge.target));

  const elkGraph = {
    id: "root",
    layoutOptions: {
      "elk.algorithm": "layered",
      "elk.direction": "RIGHT",
      "elk.separateConnectedComponents": "false",
      "elk.spacing.nodeNode": "40",
      "elk.layered.spacing.nodeNodeBetweenLayers": "80",
      "elk.layered.nodePlacement.strategy": "BRANDES_KOEPF",
      "elk.layered.nodePlacement.bk.fixedAlignment": "BALANCED",
    },
    children: graph.nodes.map((node) => ({
      id: node.id,
      width: NODE_WIDTH,
      height:
        node.nodeKind === "root" || node.nodeKind === "phantom"
          ? NODE_HEIGHT_GROUP
          : NODE_HEIGHT_RESOURCE,
      ...(targetNodeIds.has(node.id)
        ? {}
        : { layoutOptions: { "elk.layered.layering.layerConstraint": "FIRST" } }),
    })),
    edges: graph.edges.map((edge) => ({
      id: edge.id,
      sources: [edge.source],
      targets: [edge.target],
      layoutOptions: { "elk.layered.priority.shortness": "10" },
    })),
  };

  const elkResult = await getElk().layout(elkGraph);

  const nodeById = new Map(graph.nodes.map((node) => [node.id, node]));
  const flowNodes: Node[] = [];
  const nodeLayouts = new Map<string, NodeLayout>();

  for (const child of elkResult.children ?? []) {
    const graphNode = nodeById.get(child.id);
    if (graphNode !== undefined) {
      const pos = { x: child.x ?? 0, y: child.y ?? 0 };
      const height =
        graphNode.nodeKind === "root" || graphNode.nodeKind === "phantom"
          ? NODE_HEIGHT_GROUP
          : NODE_HEIGHT_RESOURCE;
      flowNodes.push(toFlatFlowNode(graphNode, pos));
      nodeLayouts.set(child.id, { position: pos, height });
    }
  }

  const lateralEdges =
    graph.lateralEdges !== undefined && graph.lateralEdges.length > 0
      ? toLateralFlowEdges(graph.lateralEdges, nodeLayouts)
      : undefined;

  return {
    nodes: flowNodes,
    edges: toFlowEdges(graph.edges),
    ...(lateralEdges !== undefined ? { lateralEdges } : {}),
  };
};
