import { useMemo, useRef, useState, type DragEvent } from "react";
import {
  addEdge,
  applyEdgeChanges,
  applyNodeChanges,
  Background,
  BackgroundVariant,
  ConnectionLineType,
  Controls,
  MiniMap,
  ReactFlow,
  reconnectEdge,
  type Connection,
  type Edge,
  type ReactFlowInstance
} from "@xyflow/react";
import { StatePanel } from "@govoplan/core-webui";
import { definitionConnectionError } from "@govoplan/core-webui/definition-graph";
import type {
  DataflowDiagnostic,
  NodeTypeDefinition,
  NodePreviewDiagnostic,
  PipelineGraph,
  PipelineGraphNode
} from "../../api/dataflow";
import DataflowNode, { type DataflowFlowNode } from "./DataflowNode";
import {
  closestProximityConnection,
  graphEdgesFromFlow,
  newGraphEdgeId,
  proximityPreviewEdge
} from "./graphInteractions";
import { FALLBACK_NODE_LIBRARY, newNode } from "./model";

const nodeTypes = { dataflow: DataflowNode };

type DataflowCanvasProps = {
  graph: PipelineGraph;
  diagnostics: DataflowDiagnostic[];
  nodeDiagnostics: NodePreviewDiagnostic[];
  nodeLibrary: NodeTypeDefinition[];
  selectedNodeId: string | null;
  readOnly: boolean;
  onGraphChange: (graph: PipelineGraph) => void;
  onSelectNode: (nodeId: string | null) => void;
};

export default function DataflowCanvas({
  graph,
  diagnostics,
  nodeDiagnostics,
  nodeLibrary,
  selectedNodeId,
  readOnly,
  onGraphChange,
  onSelectNode
}: DataflowCanvasProps) {
  const canvasRef = useRef<HTMLDivElement | null>(null);
  const [instance, setInstance] = useState<ReactFlowInstance<DataflowFlowNode, Edge> | null>(null);
  const [proximityEdge, setProximityEdge] = useState<Edge | null>(null);
  const [selectedEdgeId, setSelectedEdgeId] = useState<string | null>(null);
  const reconnectSuccessful = useRef(true);
  const reconnectingEdgeId = useRef<string | null>(null);
  const errorNodeIds = useMemo(
    () => new Set(diagnostics.filter((item) => item.severity === "error" && item.node_id).map((item) => item.node_id)),
    [diagnostics]
  );
  const rowCounts = useMemo(
    () => new Map(nodeDiagnostics.map((item) => [item.node_id, item.output_rows])),
    [nodeDiagnostics]
  );
  const definitions = useMemo(
    () => new Map(nodeLibrary.map((definition) => [definition.type, definition])),
    [nodeLibrary]
  );
  const nodes = useMemo<DataflowFlowNode[]>(
    () => graph.nodes.map((node) => {
      const definition = definitions.get(node.type)
        ?? FALLBACK_NODE_LIBRARY.find((item) => item.type === node.type)
        ?? FALLBACK_NODE_LIBRARY[0];
      return {
        id: node.id,
        type: "dataflow",
        position: node.position,
        initialWidth: 180,
        initialHeight: node.type.startsWith("combine.") ? 64 : 54,
        selected: node.id === selectedNodeId,
        data: {
          label: node.label,
          transformType: node.type,
          definition,
          config: node.config,
          hasError: errorNodeIds.has(node.id),
          outputRows: rowCounts.get(node.id)
        }
      };
    }),
    [definitions, errorNodeIds, graph.nodes, rowCounts, selectedNodeId]
  );
  const edges = useMemo<Edge[]>(
    () => graph.edges.map((edge) => ({
      id: edge.id,
      source: edge.source,
      target: edge.target,
      sourceHandle: edge.source_port ?? "output",
      targetHandle: edge.target_port ?? "input",
      type: "smoothstep",
      className: "dataflow-edge",
      selected: edge.id === selectedEdgeId
    })),
    [graph.edges, selectedEdgeId]
  );
  const displayedEdges = useMemo(
    () => proximityEdge ? [...edges, proximityEdge] : edges,
    [edges, proximityEdge]
  );

  const updateNodes = (nextNodes: DataflowFlowNode[]) => {
    const ids = new Set(nextNodes.map((node) => node.id));
    const nextGraphNodes = nextNodes.map((flowNode) => {
      const current = graph.nodes.find((node) => node.id === flowNode.id);
      if (!current) {
        throw new Error(`Unknown Dataflow node: ${flowNode.id}`);
      }
      return { ...current, position: flowNode.position };
    });
    onGraphChange({
      ...graph,
      nodes: nextGraphNodes,
      edges: graph.edges.filter((edge) => ids.has(edge.source) && ids.has(edge.target))
    });
  };

  const restoreFocusAfterRemoval = (removedNodeIds: Set<string>) => {
    const focusedNodeId = document.activeElement
      ?.closest<HTMLElement>(".react-flow__node")
      ?.dataset.id;
    if (!focusedNodeId || !removedNodeIds.has(focusedNodeId)) return;
    const removedIndex = graph.nodes.findIndex((node) => node.id === focusedNodeId);
    const remainingNodes = graph.nodes.filter((node) => !removedNodeIds.has(node.id));
    const nextNode = remainingNodes[Math.min(Math.max(0, removedIndex), remainingNodes.length - 1)];
    requestAnimationFrame(() => {
      const nextElement = nextNode
        ? canvasRef.current?.querySelector<HTMLElement>(
            `.react-flow__node[data-id="${CSS.escape(nextNode.id)}"]`
          )
        : null;
      (nextElement ?? canvasRef.current)?.focus();
    });
  };

  const updateEdges = (nextEdges: Edge[]) => {
    onGraphChange({
      ...graph,
      edges: graphEdgesFromFlow(nextEdges)
    });
  };

  const connectionError = (
    connection: Connection | Edge,
    ignoredEdgeId: string | null = reconnectingEdgeId.current
  ): string | null => definitionConnectionError(
    ignoredEdgeId
      ? {
          ...graph,
          edges: graph.edges.filter((edge) => edge.id !== ignoredEdgeId)
        }
      : graph,
    nodeLibrary,
    {
      source: connection.source,
      target: connection.target,
      sourcePort: connection.sourceHandle,
      targetPort: connection.targetHandle
    }
  );

  const isValidConnection = (connection: Connection | Edge): boolean => {
    if (readOnly || !connection.source || !connection.target) return false;
    return connectionError(connection) === null;
  };

  const onConnect = (connection: Connection) => {
    if (!isValidConnection(connection)) return;
    const next = addEdge(
      {
        ...connection,
        id: newGraphEdgeId(),
        type: "smoothstep"
      },
      edges
    );
    updateEdges(next);
  };

  const onReconnect = (oldEdge: Edge, connection: Connection) => {
    if (readOnly || connectionError(connection, oldEdge.id)) return;
    reconnectSuccessful.current = true;
    updateEdges(reconnectEdge(
      oldEdge,
      connection,
      edges,
      { shouldReplaceId: false }
    ));
  };

  const proximityConnection = (draggedNode: DataflowFlowNode) => (
    closestProximityConnection(
      draggedNode,
      nodes,
      graph,
      nodeLibrary
    )
  );

  const onDrop = (event: DragEvent<HTMLDivElement>) => {
    event.preventDefault();
    if (readOnly || !instance) return;
    const type = event.dataTransfer.getData("application/x-govoplan-dataflow-node");
    if (!type) return;
    const position = instance.screenToFlowPosition({ x: event.clientX, y: event.clientY });
    const node = newNode(type, position, nodeLibrary);
    onGraphChange({ ...graph, nodes: [...graph.nodes, node] });
    onSelectNode(node.id);
  };

  return (
    <div
      ref={canvasRef}
      className="definition-graph-canvas dataflow-canvas"
      tabIndex={0}
      aria-label="Dataflow graph canvas"
      onDragOver={(event) => {
        if (event.dataTransfer.types.includes("application/x-govoplan-dataflow-node")) {
          event.preventDefault();
          event.dataTransfer.dropEffect = "copy";
        }
      }}
      onDrop={onDrop}
    >
      <ReactFlow<DataflowFlowNode, Edge>
        nodes={nodes}
        edges={displayedEdges}
        nodeTypes={nodeTypes}
        onInit={setInstance}
        onNodesChange={(changes) => {
          if (readOnly) return;
          const graphChanges = changes.filter((change) => change.type !== "dimensions");
          if (graphChanges.length) {
            restoreFocusAfterRemoval(new Set(
              graphChanges
                .filter((change) => change.type === "remove")
                .map((change) => change.id)
            ));
            updateNodes(applyNodeChanges(graphChanges, nodes));
          }
        }}
        onEdgesChange={(changes) => {
          if (readOnly) return;
          const selectedChange = changes.find(
            (change) => change.type === "select" && change.selected
          );
          if (selectedChange?.type === "select") {
            setSelectedEdgeId(selectedChange.id);
          } else if (changes.some(
            (change) => (
              change.type === "select"
              && change.id === selectedEdgeId
              && !change.selected
            )
          )) {
            setSelectedEdgeId(null);
          }
          const graphChanges = changes.filter((change) => change.type !== "select");
          if (graphChanges.length) {
            if (graphChanges.some(
              (change) => change.type === "remove" && change.id === selectedEdgeId
            )) {
              setSelectedEdgeId(null);
            }
            updateEdges(applyEdgeChanges(graphChanges, edges));
          }
        }}
        onConnect={onConnect}
        onReconnect={onReconnect}
        onReconnectStart={(_event, edge) => {
          reconnectSuccessful.current = false;
          reconnectingEdgeId.current = edge.id;
          setProximityEdge(null);
        }}
        onReconnectEnd={(_event, edge) => {
          if (!reconnectSuccessful.current && !readOnly) {
            updateEdges(edges.filter((candidate) => candidate.id !== edge.id));
            setSelectedEdgeId(null);
          }
          reconnectSuccessful.current = true;
          reconnectingEdgeId.current = null;
        }}
        onNodeDrag={(_event, node) => {
          if (readOnly) return;
          const connection = proximityConnection(node);
          setProximityEdge(connection ? proximityPreviewEdge(connection) : null);
        }}
        onNodeDragStop={(_event, node) => {
          setProximityEdge(null);
          if (readOnly) return;
          const connection = proximityConnection(node);
          if (!connection) return;
          const nextEdges = addEdge(
            {
              ...connection,
              id: newGraphEdgeId(),
              type: "smoothstep",
              className: "dataflow-edge"
            },
            edges
          );
          onGraphChange({
            ...graph,
            nodes: graph.nodes.map((graphNode) => (
              graphNode.id === node.id
                ? { ...graphNode, position: node.position }
                : graphNode
            )),
            edges: graphEdgesFromFlow(nextEdges)
          });
        }}
        isValidConnection={isValidConnection}
        onNodeClick={(_event, node) => {
          setSelectedEdgeId(null);
          onSelectNode(node.id);
        }}
        onEdgeClick={(_event, edge) => {
          setSelectedEdgeId(edge.id);
          onSelectNode(null);
        }}
        onPaneClick={() => {
          setSelectedEdgeId(null);
          onSelectNode(null);
        }}
        nodesDraggable={!readOnly}
        nodesConnectable={!readOnly}
        edgesReconnectable={!readOnly}
        elementsSelectable
        deleteKeyCode={readOnly ? null : ["Backspace", "Delete"]}
        connectionLineType={ConnectionLineType.SmoothStep}
        connectionLineStyle={{ stroke: "var(--accent)", strokeWidth: 3 }}
        connectionRadius={32}
        fitView
        fitViewOptions={{ padding: 0.22, maxZoom: 1.25 }}
        minZoom={0.25}
        maxZoom={1.8}
      >
        <Background variant={BackgroundVariant.Dots} gap={20} size={1.3} />
        <MiniMap
          pannable
          zoomable
          nodeStrokeWidth={2}
          nodeColor={(node) => node.data.hasError ? "var(--danger)" : "var(--accent)"}
        />
        <Controls showInteractive={false} />
      </ReactFlow>
      {!graph.nodes.length ? (
        <StatePanel className="definition-graph-canvas-empty" size="fill" description="Drop a source here" />
      ) : null}
    </div>
  );
}

export function updateGraphNode(
  graph: PipelineGraph,
  updatedNode: PipelineGraphNode
): PipelineGraph {
  return {
    ...graph,
    nodes: graph.nodes.map((node) => node.id === updatedNode.id ? updatedNode : node)
  };
}