import { create } from "zustand";
import {
  applyNodeChanges,
  applyEdgeChanges,
} from "@xyflow/react";
import type {
  Node,
  Edge,
  OnNodesChange,
  OnEdgesChange,
  Viewport,
} from "@xyflow/react";

import type { LayoutDirection, EdgeRouting } from "../lib/elk-layout";
import type { GraphData } from "../types/graph";
import { graphToFlow } from "../lib/graph-converter";

type ProposalStatus = "idle" | "pending" | "accepted" | "rejected";

interface GraphState {
  nodes: Node[];
  edges: Edge[];
  viewport: Viewport;
  nodeCount: number;
  zoomLevel: number;
  layoutDirection: LayoutDirection;
  edgeRouting: EdgeRouting;
  isLayouting: boolean;
  highlightedNodeId: string | null;
  selectedNodeIds: string[];
  layoutRequestId: number;
  proposedPatch: GraphData | null;
  previousGraphSnapshot: { nodes: Node[]; edges: Edge[] } | null;
  proposalStatus: ProposalStatus;
  lastProposalOutcome: "accepted" | "rejected" | null;
  prefillChat: { nodeId: string; text: string } | null;
  fitViewRequested: number;
  focusNodeId: string | null;
  setNodes: (nodes: Node[]) => void;
  setEdges: (edges: Edge[]) => void;
  onNodesChange: OnNodesChange;
  onEdgesChange: OnEdgesChange;
  onViewportChange: (viewport: Viewport) => void;
  setLayoutDirection: (direction: LayoutDirection) => void;
  setEdgeRouting: (routing: EdgeRouting) => void;
  setIsLayouting: (isLayouting: boolean) => void;
  setHighlightedNodeId: (id: string | null) => void;
  setSelectedNodeIds: (ids: string[]) => void;
  requestLayout: () => void;
  initializeFromGraphData: (nodes: Node[], edges: Edge[]) => void;
  proposeChanges: (graphData: GraphData) => void;
  acceptProposal: () => void;
  rejectProposal: () => void;
  clearProposal: () => void;
  setPrefillChat: (nodeId: string, text: string) => void;
  clearPrefillChat: () => void;
  requestFitView: () => void;
  setFocusNodeId: (id: string | null) => void;
  reset: () => void;
}

export const useGraphStore = create<GraphState>((set, get) => ({
  nodes: [],
  edges: [],
  viewport: { x: 0, y: 0, zoom: 1 },
  nodeCount: 0,
  zoomLevel: 100,
  layoutDirection: "DOWN",
  edgeRouting: "ORTHOGONAL",
  isLayouting: false,
  highlightedNodeId: null,
  selectedNodeIds: [],
  layoutRequestId: 0,
  proposedPatch: null,
  previousGraphSnapshot: null,
  proposalStatus: "idle",
  lastProposalOutcome: null,
  prefillChat: null,
  fitViewRequested: 0,
  focusNodeId: null,

  setNodes: (nodes) => set({ nodes, nodeCount: nodes.length }),
  setEdges: (edges) => set({ edges }),

  onNodesChange: (changes) => {
    const updatedNodes = applyNodeChanges(changes, get().nodes);
    set({ nodes: updatedNodes, nodeCount: updatedNodes.length });
  },

  onEdgesChange: (changes) => {
    set({ edges: applyEdgeChanges(changes, get().edges) });
  },

  onViewportChange: (viewport) => {
    set({ viewport, zoomLevel: Math.round(viewport.zoom * 100) });
  },

  setLayoutDirection: (layoutDirection) => set({ layoutDirection }),
  setEdgeRouting: (edgeRouting) => set({ edgeRouting }),
  setIsLayouting: (isLayouting) => set({ isLayouting }),
  setHighlightedNodeId: (highlightedNodeId) => set({ highlightedNodeId }),
  setSelectedNodeIds: (selectedNodeIds) => set({ selectedNodeIds }),
  requestLayout: () => set((s) => ({ layoutRequestId: s.layoutRequestId + 1 })),

  initializeFromGraphData: (nodes, edges) => {
    set({ nodes, edges, nodeCount: nodes.length });
  },

  proposeChanges: (graphData) => {
    const { nodes, edges } = get();

    // Snapshot current state for revert
    set({
      previousGraphSnapshot: { nodes: [...nodes], edges: [...edges] },
      proposedPatch: graphData,
      proposalStatus: "pending",
      lastProposalOutcome: null,
      highlightedNodeId: null,
    });

    // Convert proposed graph to flow format
    const proposed = graphToFlow(graphData);

    // Compute diff by comparing node IDs
    const currentIds = new Set(nodes.map((n) => n.id));
    const proposedIds = new Set(proposed.nodes.map((n) => n.id));

    const currentNodeMap = new Map(nodes.map((n) => [n.id, n]));
    const isDifferentNode = (p: Node, c: Node | undefined): boolean => {
      if (!c) return false;
      const pData = p.data as Record<string, unknown>;
      const cData = c.data as Record<string, unknown>;
      if (pData.label !== cData.label || pData.type !== cData.type) return true;
      // Deep compare additional properties (columns, tag, etc.)
      if (JSON.stringify(pData.columns) !== JSON.stringify(cData.columns)) return true;
      if (pData.tag !== cData.tag) return true;
      return false;
    };

    // Merge: proposed nodes with diff classes + removed nodes with remove class
    const mergedNodes = [
      ...proposed.nodes.map((n) => ({
        ...n,
        className: !currentIds.has(n.id)
          ? "ai-diff-add"
          : isDifferentNode(n, currentNodeMap.get(n.id))
            ? "ai-diff-modified"
            : undefined,
      })),
      ...nodes
        .filter((n) => !proposedIds.has(n.id))
        .map((n) => ({ ...n, className: "ai-diff-remove" })),
    ];

    // Same for edges
    const currentEdgeIds = new Set(edges.map((e) => e.id));
    const proposedEdgeIds = new Set(proposed.edges.map((e) => e.id));
    const currentEdgeMap = new Map(edges.map((e) => [e.id, e]));

    const isDifferentEdge = (p: Edge, c: Edge | undefined): boolean => {
      if (!c) return false;
      return p.label !== c.label || p.type !== c.type;
    };

    const mergedEdges = [
      ...proposed.edges.map((e) => ({
        ...e,
        className: !currentEdgeIds.has(e.id)
          ? "ai-diff-add"
          : isDifferentEdge(e, currentEdgeMap.get(e.id))
            ? "ai-diff-modified"
            : undefined,
      })),
      ...edges
        .filter((e) => !proposedEdgeIds.has(e.id))
        .map((e) => ({ ...e, className: "ai-diff-remove" })),
    ];

    set({ nodes: mergedNodes, edges: mergedEdges, nodeCount: mergedNodes.length });
  },

  acceptProposal: () => {
    const { proposedPatch } = get();
    if (!proposedPatch) return;

    const { nodes, edges } = graphToFlow(proposedPatch);

    set({
      nodes,
      edges,
      nodeCount: nodes.length,
      proposedPatch: null,
      previousGraphSnapshot: null,
      proposalStatus: "idle",
      lastProposalOutcome: "accepted",
      ...(proposedPatch.meta?.layoutDirection && {
        layoutDirection: proposedPatch.meta.layoutDirection,
      }),
      ...(proposedPatch.meta?.edgeRouting && {
        edgeRouting: proposedPatch.meta.edgeRouting,
      }),
    });

    get().requestLayout();
  },

  rejectProposal: () => {
    const { previousGraphSnapshot } = get();
    if (!previousGraphSnapshot) return;

    set({
      nodes: previousGraphSnapshot.nodes,
      edges: previousGraphSnapshot.edges,
      nodeCount: previousGraphSnapshot.nodes.length,
      proposedPatch: null,
      previousGraphSnapshot: null,
      proposalStatus: "idle",
      lastProposalOutcome: "rejected",
    });
  },

  clearProposal: () => {
    set({
      proposedPatch: null,
      previousGraphSnapshot: null,
      proposalStatus: "idle",
      lastProposalOutcome: null,
    });
  },

  setPrefillChat: (nodeId, text) => set({ prefillChat: { nodeId, text } }),
  clearPrefillChat: () => set({ prefillChat: null }),
  requestFitView: () =>
    set((s) => ({ fitViewRequested: s.fitViewRequested + 1 })),
  setFocusNodeId: (id) => set({ focusNodeId: id }),

  reset: () => {
    set({
      nodes: [],
      edges: [],
      viewport: { x: 0, y: 0, zoom: 1 },
      nodeCount: 0,
      zoomLevel: 100,
      layoutDirection: "DOWN",
      edgeRouting: "ORTHOGONAL",
      isLayouting: false,
      highlightedNodeId: null,
      selectedNodeIds: [],
      layoutRequestId: 0,
      proposedPatch: null,
      previousGraphSnapshot: null,
      proposalStatus: "idle",
      lastProposalOutcome: null,
      prefillChat: null,
      fitViewRequested: 0,
      focusNodeId: null,
    });
  },
}));
