import { useCallback, useRef, useState } from 'react';
import ReactFlow, { 
  Background, 
  Controls, 
  MiniMap, 
  ReactFlowProvider,
  Panel,
  useReactFlow,
  type Node,
  type Edge
} from 'reactflow';
import 'reactflow/dist/style.css';
import useFlowStore from './store/flowStore';
import LLMNode from './components/nodes/LLMNode';
import Sidebar from './components/Sidebar';
import { ContextMenu } from './components/ContextMenu';
import { Plus } from 'lucide-react';

const nodeTypes = {
  llmNode: LLMNode,
};

function Flow() {
  const { 
    nodes, 
    edges, 
    onNodesChange, 
    onEdgesChange, 
    onConnect, 
    addNode,
    deleteEdge,
    deleteNode,
    deleteBranch,
    setSelectedNode 
  } = useFlowStore();

  const reactFlowWrapper = useRef<HTMLDivElement>(null);
  const { project } = useReactFlow();
  const [menu, setMenu] = useState<{ x: number; y: number; type: 'pane' | 'node' | 'edge'; id?: string } | null>(null);

  const onPaneClick = () => {
    setSelectedNode(null);
    setMenu(null);
  };

  const handlePaneContextMenu = (event: React.MouseEvent) => {
     event.preventDefault();
     setMenu({ x: event.clientX, y: event.clientY, type: 'pane' });
  };

  const handleNodeContextMenu = (event: React.MouseEvent, node: Node) => {
     event.preventDefault();
     setMenu({ x: event.clientX, y: event.clientY, type: 'node', id: node.id });
  };

  const handleEdgeContextMenu = (event: React.MouseEvent, edge: Edge) => {
     event.preventDefault();
     setMenu({ x: event.clientX, y: event.clientY, type: 'edge', id: edge.id });
  };

  const handleAddNode = (position?: { x: number, y: number }) => {
    const id = `node_${Date.now()}`;
    const pos = position || { x: Math.random() * 400, y: Math.random() * 400 };
    
    addNode({
      id,
      type: 'llmNode',
      position: pos,
      data: { 
        label: 'New Question', 
        model: 'gpt-4o', 
        temperature: 0.7,
        systemPrompt: '',
        userPrompt: '',
        mergeStrategy: 'smart',
        messages: [],
        traces: [],
        outgoingTraces: [],
        forkedTraces: [],
        response: '',
        status: 'idle',
        inputs: 1
      },
    });
    setMenu(null);
  };

  const onNodeClick = (_: any, node: Node) => {
    setSelectedNode(node.id);
  };

  return (
    <div style={{ width: '100vw', height: '100vh', display: 'flex' }}>
      <div style={{ flex: 1, height: '100%' }} ref={reactFlowWrapper}>
        <ReactFlow
          nodes={nodes}
          edges={edges}
          onNodesChange={onNodesChange}
          onEdgesChange={onEdgesChange}
          onConnect={onConnect}
          nodeTypes={nodeTypes}
          onNodeClick={onNodeClick}
          onPaneClick={onPaneClick}
          onPaneContextMenu={handlePaneContextMenu}
          onNodeContextMenu={handleNodeContextMenu}
          onEdgeContextMenu={handleEdgeContextMenu}
          fitView
        >
          <Background color="#aaa" gap={16} />
          <Controls />
          <MiniMap />
          <Panel position="top-left">
            <button 
              onClick={() => handleAddNode()}
              className="bg-white px-4 py-2 rounded-md shadow-md font-medium text-gray-700 hover:bg-gray-50 flex items-center gap-2"
            >
              <Plus size={16} /> Add Block
            </button>
          </Panel>
        </ReactFlow>
        
        {menu && (
          <ContextMenu 
            x={menu.x} 
            y={menu.y} 
            onClose={() => setMenu(null)}
            items={
              menu.type === 'pane' ? [
                { 
                  label: 'Add Node', 
                  onClick: () => {
                    const bounds = reactFlowWrapper.current?.getBoundingClientRect();
                    if (bounds) {
                      const position = project({
                        x: menu.x - bounds.left,
                        y: menu.y - bounds.top
                      });
                      handleAddNode(position);
                    }
                  } 
                }
              ] : menu.type === 'node' ? [
                { 
                  label: 'Delete Node (Cascade)', 
                  danger: true,
                  onClick: () => menu.id && deleteBranch(menu.id)
                }
              ] : [
                { 
                  label: 'Disconnect', 
                  onClick: () => menu.id && deleteEdge(menu.id) 
                },
                { 
                  label: 'Delete Branch', 
                  danger: true,
                  onClick: () => menu.id && deleteBranch(undefined, menu.id) 
                }
              ]
            }
          />
        )}
      </div>
      <Sidebar />
    </div>
  );
}

export default function App() {
  return (
    <ReactFlowProvider>
      <Flow />
    </ReactFlowProvider>
  );
}