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";

interface GraphState {
  nodes: Node[];
  edges: Edge[];
  viewport: Viewport;
  nodeCount: number;
  zoomLevel: number;
  layoutDirection: LayoutDirection;
  edgeRouting: EdgeRouting;
  isLayouting: boolean;
  highlightedNodeId: 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;
  initializeFromGraphData: (nodes: Node[], edges: Edge[]) => 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,

  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 }),

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

  reset: () => {
    set({
      nodes: [],
      edges: [],
      viewport: { x: 0, y: 0, zoom: 1 },
      nodeCount: 0,
      zoomLevel: 100,
      layoutDirection: "DOWN",
      edgeRouting: "ORTHOGONAL",
      isLayouting: false,
      highlightedNodeId: null,
    });
  },
}));
