import { useCallback, useRef, useState } from 'react'
import type { Node, Edge } from '@xyflow/react'

export type GraphState = { nodes: Node[]; edges: Edge[] }

function cloneState(state: GraphState): GraphState {
  return {
    nodes: state.nodes.map((n) => ({ ...n, data: n.data && typeof n.data === 'object' ? { ...n.data } : n.data })),
    edges: state.edges.map((e) => ({ ...e })),
  }
}

const MAX_HISTORY = 100

export function useGraphStateWithHistory(initialNodes: Node[], initialEdges: Edge[]) {
  const [nodes, setNodesState] = useState<Node[]>(initialNodes)
  const [edges, setEdgesState] = useState<Edge[]>(initialEdges)
  const [historySizes, setHistorySizes] = useState({ past: 0, future: 0 })

  const pastRef = useRef<GraphState[]>([])
  const futureRef = useRef<GraphState[]>([])
  const preDragRef = useRef<GraphState | null>(null)
  const nodesRef = useRef(nodes)
  const edgesRef = useRef(edges)
  nodesRef.current = nodes
  edgesRef.current = edges

  const pushToPast = useCallback((state: GraphState) => {
    pastRef.current = pastRef.current.slice(-(MAX_HISTORY - 1))
    pastRef.current.push(cloneState(state))
    futureRef.current = []
    setHistorySizes({ past: pastRef.current.length, future: 0 })
  }, [])

  const setNodes = useCallback((updater: Node[] | ((prev: Node[]) => Node[])) => {
    pushToPast({ nodes: nodesRef.current, edges: edgesRef.current })
    setNodesState(typeof updater === 'function' ? updater : () => updater)
  }, [pushToPast])

  const setEdges = useCallback((updater: Edge[] | ((prev: Edge[]) => Edge[])) => {
    pushToPast({ nodes: nodesRef.current, edges: edgesRef.current })
    setEdgesState(typeof updater === 'function' ? updater : () => updater)
  }, [pushToPast])

  const setNodesSilent = useCallback((updater: Node[] | ((prev: Node[]) => Node[])) => {
    setNodesState(typeof updater === 'function' ? updater : () => updater)
  }, [])

  const setEdgesSilent = useCallback((updater: Edge[] | ((prev: Edge[]) => Edge[])) => {
    setEdgesState(typeof updater === 'function' ? updater : () => updater)
  }, [])

  const applyGraph = useCallback((updater: (state: GraphState) => GraphState) => {
    pushToPast({ nodes: nodesRef.current, edges: edgesRef.current })
    const next = updater({ nodes: nodesRef.current, edges: edgesRef.current })
    setNodesState(next.nodes)
    setEdgesState(next.edges)
  }, [pushToPast])

  const saveForDragEnd = useCallback(() => {
    preDragRef.current = cloneState({ nodes: nodesRef.current, edges: edgesRef.current })
  }, [])

  const commitDragEnd = useCallback(() => {
    if (preDragRef.current) {
      pastRef.current = pastRef.current.slice(-(MAX_HISTORY - 1))
      pastRef.current.push(preDragRef.current)
      futureRef.current = []
      preDragRef.current = null
      setHistorySizes({ past: pastRef.current.length, future: 0 })
    }
  }, [])

  const undo = useCallback(() => {
    if (pastRef.current.length === 0) return
    const prev = pastRef.current.pop()!
    futureRef.current.push(cloneState({ nodes: nodesRef.current, edges: edgesRef.current }))
    setNodesState(prev.nodes)
    setEdgesState(prev.edges)
    setHistorySizes({ past: pastRef.current.length, future: futureRef.current.length })
  }, [])

  const redo = useCallback(() => {
    if (futureRef.current.length === 0) return
    const next = futureRef.current.pop()!
    pastRef.current.push(cloneState({ nodes: nodesRef.current, edges: edgesRef.current }))
    setNodesState(next.nodes)
    setEdgesState(next.edges)
    setHistorySizes({ past: pastRef.current.length, future: futureRef.current.length })
  }, [])

  const setStateImmediate = useCallback((state: GraphState) => {
    setNodesState(state.nodes)
    setEdgesState(state.edges)
    pastRef.current = []
    futureRef.current = []
    preDragRef.current = null
    setHistorySizes({ past: 0, future: 0 })
  }, [])

  return {
    nodes,
    edges,
    setNodes,
    setEdges,
    setNodesSilent,
    setEdgesSilent,
    applyGraph,
    saveForDragEnd,
    commitDragEnd,
    undo,
    redo,
    canUndo: historySizes.past > 0,
    canRedo: historySizes.future > 0,
    setStateImmediate,
  }
}