import { useCallback, useEffect, useMemo, useRef, useState, type MouseEvent as ReactMouseEvent } from 'react'
import { addEdge, useEdgesState, useNodesState, type Connection, type Edge, type Node, type ReactFlowInstance } from '@xyflow/react'
import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
import { toast } from 'sonner'

import {
  dispatchWorkflowDraft,
  getNodeDefinitions,
  getWorkflowRunComparison,
  getWorkflowRuns,
  preflightWorkflowDraft,
  type WorkflowConfig,
  type WorkflowDefinition,
  type WorkflowExecutionMode,
  type WorkflowNodeDefinition,
  type WorkflowParams,
  type WorkflowPreflightResponse,
} from '../../api/workflows'
import {
  applyAutoLayout,
  buildCurrentWorkflowConfig,
  inferNodeLabel,
  inferNodeType,
  inferStepFromNodeType,
  normalizeWorkflowParams,
  resolveParamsForStepChange,
  type WorkflowCanvasNodeData,
  validateWorkflowDraft,
  workflowToGraph,
} from './workflowGraphDraft'
import {
  inferWorkflowFamily,
} from './workflowBlueprints'
import {
  GRAPH_FAMILY_LABELS,
  isDefinitionAllowedForGraphFamily,
} from './workflowNodeLibrary'
import type { WorkflowUtilityTab } from './WorkflowCanvasUtilitySidebar'

export type NodeMenuAnchor = {
  clientX: number
  clientY: number
  flowPosition: { x: number; y: number }
}

function buildNodeData(
  step: string,
  params: WorkflowParams = {},
  definition?: WorkflowNodeDefinition,
  overrides?: Partial<WorkflowCanvasNodeData>,
): WorkflowCanvasNodeData {
  return {
    label: overrides?.label ?? definition?.label ?? inferNodeLabel(step),
    params: normalizeWorkflowParams(params),
    step,
    description: overrides?.description ?? definition?.description,
    icon: overrides?.icon ?? definition?.icon,
    category: overrides?.category ?? definition?.category,
  }
}

type UseWorkflowCanvasControllerArgs = {
  workflow: WorkflowDefinition
  onSave: (config: WorkflowConfig) => void
}

export function useWorkflowCanvasController({ workflow, onSave }: UseWorkflowCanvasControllerArgs) {
  const queryClient = useQueryClient()
  const { data: nodeDefinitionsData } = useQuery({
    queryKey: ['workflow-node-definitions'],
    queryFn: getNodeDefinitions,
    staleTime: 5 * 60 * 1000,
  })
  const nodeDefinitions = nodeDefinitionsData?.definitions ?? []
  const nodeDefinitionsByStep = Object.fromEntries(nodeDefinitions.map(definition => [definition.step, definition]))

  const { nodes: initNodes, edges: initEdges } = workflowToGraph(workflow.config, nodeDefinitionsByStep)
  const [nodes, setNodes, onNodesChange] = useNodesState(initNodes)
  const [edges, setEdges, onEdgesChange] = useEdgesState(initEdges)
  const [selectedNodeId, setSelectedNodeId] = useState<string | null>(null)
  const [selectedRunId, setSelectedRunId] = useState<string | null>(null)
  const [dispatchContextId, setDispatchContextId] = useState('')
  const [preflightResult, setPreflightResult] = useState<WorkflowPreflightResponse | null>(null)
  const [executionMode, setExecutionMode] = useState<WorkflowExecutionMode>(workflow.config.ui?.execution_mode ?? 'legacy')
  const [nodeMenuAnchor, setNodeMenuAnchor] = useState<NodeMenuAnchor | null>(null)
  const [activeUtilityTab, setActiveUtilityTab] = useState<WorkflowUtilityTab>('library')
  const reactFlowWrapper = useRef<HTMLDivElement>(null)
  const [reactFlowInstance, setReactFlowInstance] = useState<ReactFlowInstance<Node, Edge> | null>(null)

  const validation = validateWorkflowDraft(nodes, edges, nodeDefinitionsByStep, nodeDefinitions.length > 0)
  const selectedEdgeIds = useMemo(
    () => edges.filter(edge => Boolean((edge as Edge & { selected?: boolean }).selected)).map(edge => edge.id),
    [edges],
  )
  const graphFamily = useMemo(
    () =>
      inferWorkflowFamily(
        buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode),
        nodeDefinitionsByStep,
      ),
    [edges, executionMode, nodeDefinitionsByStep, nodes, workflow],
  )

  const { data: workflowRuns = [] } = useQuery({
    queryKey: ['workflow-runs', workflow.id],
    queryFn: () => getWorkflowRuns(workflow.id),
    refetchInterval: 5000,
  })

  const selectedRun = workflowRuns.find(run => run.id === selectedRunId) ?? workflowRuns[0] ?? null

  const { data: selectedRunComparison, isFetching: isComparisonLoading } = useQuery({
    queryKey: ['workflow-run-comparison', selectedRun?.id],
    queryFn: () => getWorkflowRunComparison(selectedRun!.id),
    enabled: Boolean(selectedRun?.id && selectedRun.execution_mode === 'shadow'),
    refetchInterval: selectedRun?.status === 'pending' || selectedRun?.status === 'running' ? 5000 : false,
  })

  const dispatchMutation = useMutation({
    mutationFn: ({ contextId, config }: { contextId: string; config: WorkflowConfig }) =>
      dispatchWorkflowDraft({
        workflow_id: workflow.id,
        context_id: contextId,
        config,
      }),
    onSuccess: result => {
      queryClient.invalidateQueries({ queryKey: ['workflow-runs', workflow.id] })
      setSelectedRunId(result.workflow_run.id)
      toast.success(`Graph run dispatched: ${result.dispatched} task${result.dispatched === 1 ? '' : 's'}`)
    },
    onError: (error: any) => {
      toast.error(error?.response?.data?.detail || 'Failed to dispatch workflow')
    },
  })

  const preflightMutation = useMutation({
    mutationFn: ({ contextId, config }: { contextId: string; config: WorkflowConfig }) =>
      preflightWorkflowDraft({
        workflow_id: workflow.id,
        context_id: contextId,
        config,
      }),
    onSuccess: result => {
      setPreflightResult(result)
      if (result.graph_dispatch_allowed) {
        toast.success(result.summary)
      } else {
        toast.error(result.summary)
      }
    },
    onError: (error: any) => {
      setPreflightResult(null)
      toast.error(error?.response?.data?.detail || 'Failed to preflight workflow')
    },
  })

  useEffect(() => {
    const graph = workflowToGraph(workflow.config, nodeDefinitionsByStep)
    setNodes(graph.nodes)
    setEdges(graph.edges)
    setSelectedNodeId(null)
    setSelectedRunId(null)
    setNodeMenuAnchor(null)
    setPreflightResult(null)
    setExecutionMode(workflow.config.ui?.execution_mode ?? 'legacy')
    setActiveUtilityTab('library')
  }, [nodeDefinitionsData, setEdges, setNodes, workflow.config])

  useEffect(() => {
    if (!selectedRunId && workflowRuns.length > 0) {
      setSelectedRunId(workflowRuns[0].id)
      return
    }
    if (selectedRunId && !workflowRuns.some(run => run.id === selectedRunId)) {
      setSelectedRunId(workflowRuns[0]?.id ?? null)
    }
  }, [selectedRunId, workflowRuns])

  const onConnect = useCallback(
    (connection: Connection) => setEdges(currentEdges => addEdge(connection, currentEdges)),
    [setEdges],
  )

  const onNodeClick = useCallback((_: ReactMouseEvent, node: Node) => {
    setNodeMenuAnchor(null)
    setSelectedNodeId(node.id)
    setActiveUtilityTab('inspector')
  }, [])

  const onEdgeClick = useCallback((_: ReactMouseEvent, edge: Edge) => {
    setNodeMenuAnchor(null)
    setSelectedNodeId(null)
    setEdges(currentEdges =>
      currentEdges.map(currentEdge => ({
        ...currentEdge,
        selected: currentEdge.id === edge.id,
      })),
    )
  }, [setEdges])

  const onPaneClick = useCallback(() => {
    setNodeMenuAnchor(null)
    setSelectedNodeId(null)
    setEdges(currentEdges =>
      currentEdges.map(edge => ({
        ...edge,
        selected: false,
      })),
    )
  }, [setEdges])

  const handleParamsChange = useCallback(
    (newParams: WorkflowParams) => {
      setNodes(currentNodes =>
        currentNodes.map(node => {
          if (node.id === selectedNodeId) {
            return { ...node, data: { ...node.data, params: normalizeWorkflowParams(newParams) } }
          }
          return node
        }),
      )
    },
    [selectedNodeId, setNodes],
  )

  const handlePipelineStepChange = useCallback(
    (stepName: string) => {
      const definition = nodeDefinitionsByStep[stepName]
      if (definition && !isDefinitionAllowedForGraphFamily(definition, graphFamily)) {
        toast.error(`${definition.label} does not belong to the ${GRAPH_FAMILY_LABELS[graphFamily]} family.`)
        return
      }

      setNodes(currentNodes =>
        currentNodes.map(node => {
          if (node.id !== selectedNodeId) return node

          const currentData = (node.data as WorkflowCanvasNodeData | undefined) ?? buildNodeData(stepName)
          return {
            ...node,
            type: definition?.node_type ?? inferNodeType(stepName),
            data: {
              ...buildNodeData(
                stepName || inferStepFromNodeType(node.type),
                resolveParamsForStepChange(definition, currentData.params),
                definition,
              ),
              step: stepName || inferStepFromNodeType(node.type),
            },
          }
        }),
      )
    },
    [graphFamily, nodeDefinitionsByStep, selectedNodeId, setNodes],
  )

  const openNodeMenu = useCallback(
    (clientX: number, clientY: number) => {
      if (!reactFlowInstance) return
      setNodeMenuAnchor({
        clientX,
        clientY,
        flowPosition: reactFlowInstance.screenToFlowPosition({ x: clientX, y: clientY }),
      })
    },
    [reactFlowInstance],
  )

  const handlePaneContextMenu = useCallback(
    (event: MouseEvent | ReactMouseEvent) => {
      event.preventDefault()
      setSelectedNodeId(null)
      openNodeMenu(event.clientX, event.clientY)
    },
    [openNodeMenu],
  )

  const handleNodeContextMenu = useCallback(
    (event: ReactMouseEvent, node: Node) => {
      event.preventDefault()
      setSelectedNodeId(node.id)
      openNodeMenu(event.clientX, event.clientY)
    },
    [openNodeMenu],
  )

  const insertNode = useCallback(
    (step: string, preferredPosition?: { x: number; y: number }) => {
      const definition = nodeDefinitionsByStep[step]
      if (definition && !isDefinitionAllowedForGraphFamily(definition, graphFamily)) {
        toast.error(`${definition.label} cannot be added to a ${GRAPH_FAMILY_LABELS[graphFamily]} workflow.`)
        return
      }

      const type = definition?.node_type ?? inferNodeType(step)
      const fallbackX = nodes.length > 0 ? Math.max(...nodes.map(node => node.position.x)) + 220 : 120
      const fallbackY = nodes.length > 0 ? Math.max(...nodes.map(node => node.position.y)) + 40 : 120
      const newNode: Node = {
        id: `${step}_${Date.now()}`,
        type,
        position: preferredPosition ?? { x: fallbackX, y: fallbackY },
        data: buildNodeData(step, definition?.defaults ?? {}, definition),
      }

      setNodes(currentNodes => [...currentNodes, newNode])
      setSelectedNodeId(newNode.id)
      setNodeMenuAnchor(null)
      setActiveUtilityTab('inspector')
    },
    [graphFamily, nodeDefinitionsByStep, nodes, setNodes],
  )

  const handleOpenToolbarNodeMenu = useCallback(() => {
    if (!reactFlowWrapper.current || !reactFlowInstance) return
    const bounds = reactFlowWrapper.current.getBoundingClientRect()
    openNodeMenu(bounds.left + 36, bounds.top + 36)
  }, [openNodeMenu, reactFlowInstance])

  const handleAutoLayout = useCallback(() => {
    setNodes(currentNodes => applyAutoLayout(currentNodes, edges))
    setNodeMenuAnchor(null)
    window.requestAnimationFrame(() => {
      reactFlowInstance?.fitView({ padding: 0.2, duration: 250 })
    })
  }, [edges, reactFlowInstance, setNodes])

  const deleteEdgesById = useCallback((edgeIds: string[]) => {
    if (edgeIds.length === 0) return
    setEdges(currentEdges => currentEdges.filter(edge => !edgeIds.includes(edge.id)))
    setSelectedNodeId(null)
    setNodeMenuAnchor(null)
    toast.success(edgeIds.length === 1 ? 'Connection deleted' : `${edgeIds.length} connections deleted`)
  }, [setEdges])

  const handleDeleteSelectedEdges = useCallback(() => {
    deleteEdgesById(selectedEdgeIds)
  }, [deleteEdgesById, selectedEdgeIds])

  const onEdgeContextMenu = useCallback((event: ReactMouseEvent, edge: Edge) => {
    event.preventDefault()
    event.stopPropagation()
    deleteEdgesById([edge.id])
  }, [deleteEdgesById])

  const onEdgeDoubleClick = useCallback((event: ReactMouseEvent, edge: Edge) => {
    event.preventDefault()
    event.stopPropagation()
    deleteEdgesById([edge.id])
  }, [deleteEdgesById])

  useEffect(() => {
    const handleKeyDown = (event: KeyboardEvent) => {
      const target = event.target as HTMLElement | null
      const isEditingField =
        target instanceof HTMLInputElement ||
        target instanceof HTMLTextAreaElement ||
        target instanceof HTMLSelectElement ||
        Boolean(target?.closest('[contenteditable="true"]'))

      if (isEditingField) return
      if ((event.key === 'Delete' || event.key === 'Backspace') && selectedEdgeIds.length > 0) {
        event.preventDefault()
        deleteEdgesById(selectedEdgeIds)
      }
    }

    window.addEventListener('keydown', handleKeyDown)
    return () => window.removeEventListener('keydown', handleKeyDown)
  }, [deleteEdgesById, selectedEdgeIds])

  useEffect(() => {
    if (selectedNodeId) {
      setActiveUtilityTab('inspector')
      return
    }
    if (activeUtilityTab === 'inspector') {
      setActiveUtilityTab('library')
    }
  }, [activeUtilityTab, selectedNodeId])

  const handleSave = useCallback(() => {
    if (validation.errors.length > 0) {
      toast.error('Resolve workflow validation errors before saving.')
      return
    }
    onSave(buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode))
  }, [edges, executionMode, nodes, onSave, validation.errors.length, workflow])

  const handleDispatch = useCallback(() => {
    if (!dispatchContextId.trim()) {
      toast.error('Context ID is required for a graph test run.')
      return
    }
    if (validation.errors.length > 0) {
      toast.error('Resolve workflow validation errors before dispatching.')
      return
    }
    dispatchMutation.mutate({
      contextId: dispatchContextId.trim(),
      config: buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode),
    })
  }, [dispatchContextId, dispatchMutation, edges, executionMode, nodes, validation.errors.length, workflow])

  const handlePreflight = useCallback(() => {
    if (!dispatchContextId.trim()) {
      toast.error('Context ID is required for a graph preflight.')
      return
    }
    if (validation.errors.length > 0) {
      toast.error('Resolve workflow validation errors before running preflight.')
      return
    }
    preflightMutation.mutate({
      contextId: dispatchContextId.trim(),
      config: buildCurrentWorkflowConfig(workflow, nodes, edges, executionMode),
    })
  }, [dispatchContextId, edges, executionMode, nodes, preflightMutation, validation.errors.length, workflow])

  const selectedNode = useMemo(
    () => nodes.find(node => node.id === selectedNodeId),
    [nodes, selectedNodeId],
  )

  return {
    reactFlowWrapper,
    nodeDefinitions,
    nodeDefinitionsByStep,
    nodes,
    edges,
    onNodesChange,
    onEdgesChange,
    selectedEdgeIds,
    selectedNode,
    workflowRuns,
    selectedRun,
    selectedRunComparison,
    isComparisonLoading,
    dispatchMutation,
    preflightMutation,
    dispatchContextId,
    setDispatchContextId,
    preflightResult,
    executionMode,
    setExecutionMode,
    nodeMenuAnchor,
    setNodeMenuAnchor,
    activeUtilityTab,
    setActiveUtilityTab,
    validation,
    graphFamily,
    onConnect,
    onNodeClick,
    onEdgeClick,
    onPaneClick,
    handleParamsChange,
    handlePipelineStepChange,
    handlePaneContextMenu,
    handleNodeContextMenu,
    insertNode,
    handleOpenToolbarNodeMenu,
    handleAutoLayout,
    handleDeleteSelectedEdges,
    onEdgeContextMenu,
    onEdgeDoubleClick,
    handleSave,
    handleDispatch,
    handlePreflight,
    setReactFlowInstance,
    setSelectedRunId,
  }
}