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 LeftSidebar from './components/LeftSidebar';
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,
    toggleNodeDisabled,
    archiveNode,
    createNodeFromArchive,
    toggleTraceDisabled
  } = 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 [isLeftOpen, setIsLeftOpen] = useState(true);
  const [isRightOpen, setIsRightOpen] = useState(true);

  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',
        reasoningEffort: 'medium',  // Default for reasoning models
        messages: [],
        traces: [],
        outgoingTraces: [],
        forkedTraces: [],
        response: '',
        status: 'idle',
        inputs: 1
      },
    });
    setMenu(null);
  };

  const onNodeClick = (_: any, node: Node) => {
    // Don't select disabled nodes
    const nodeData = node.data as any;
    if (nodeData?.disabled) return;
    setSelectedNode(node.id);
  };

  const onDragOver = (event: React.DragEvent) => {
    event.preventDefault();
    event.dataTransfer.dropEffect = 'copy';
  };

  const onDrop = (event: React.DragEvent) => {
    event.preventDefault();
    
    const archiveId = event.dataTransfer.getData('archiveId');
    if (!archiveId) return;
    
    const bounds = reactFlowWrapper.current?.getBoundingClientRect();
    if (!bounds) return;
    
    const position = project({
      x: event.clientX - bounds.left,
      y: event.clientY - bounds.top
    });
    
    createNodeFromArchive(archiveId, position);
  };

  return (
    <div style={{ width: '100vw', height: '100vh', display: 'flex' }}>
      <LeftSidebar isOpen={isLeftOpen} onToggle={() => setIsLeftOpen(!isLeftOpen)} />
      
      <div 
        style={{ flex: 1, height: '100%', position: 'relative' }} 
        ref={reactFlowWrapper}
        onDragOver={onDragOver}
        onDrop={onDrop}
      >
        <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' ? (() => {
                const targetNode = nodes.find(n => n.id === menu.id);
                const isDisabled = targetNode?.data?.disabled;
                
                // If disabled, only show Enable option
                if (isDisabled) {
                  return [
                    { 
                      label: 'Enable Node',
                      onClick: () => menu.id && toggleNodeDisabled(menu.id)
                    }
                  ];
                }
                
                // Normal node menu
                return [
                  { 
                    label: 'Disable Node',
                    onClick: () => menu.id && toggleNodeDisabled(menu.id)
                  },
                  {
                    label: 'Add to Archive',
                    onClick: () => menu.id && archiveNode(menu.id)
                  },
                  { 
                    label: 'Delete Node (Cascade)', 
                    danger: true,
                    onClick: () => menu.id && deleteBranch(menu.id)
                  }
                ];
              })() : (() => {
                // Check if any node connected to this edge is disabled
                const targetEdge = edges.find(e => e.id === menu.id);
                const sourceNode = nodes.find(n => n.id === targetEdge?.source);
                const targetNode = nodes.find(n => n.id === targetEdge?.target);
                const isTraceDisabled = sourceNode?.data?.disabled || targetNode?.data?.disabled;
                
                return [
                  {
                    label: isTraceDisabled ? 'Enable Trace' : 'Disable Trace',
                    onClick: () => menu.id && toggleTraceDisabled(menu.id)
                  },
                  { 
                    label: 'Disconnect', 
                    onClick: () => menu.id && deleteEdge(menu.id) 
                  },
                  { 
                    label: 'Delete Branch', 
                    danger: true,
                    onClick: () => menu.id && deleteBranch(undefined, menu.id) 
                  }
                ];
              })()
            }
          />
        )}
      </div>
      <Sidebar isOpen={isRightOpen} onToggle={() => setIsRightOpen(!isRightOpen)} />
    </div>
  );
}

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