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

const nodeTypes = {
  llmNode: LLMNode,
};

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

  const reactFlowWrapper = useRef<HTMLDivElement>(null);
  const { project } = useReactFlow();

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

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

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

  const onPaneContextMenu = (event: React.MouseEvent) => {
     event.preventDefault();
     const bounds = reactFlowWrapper.current?.getBoundingClientRect();
     if (!bounds) return;
     
     const position = project({
       x: event.clientX - bounds.left,
       y: event.clientY - bounds.top
     });

     const id = `node_${Date.now()}`;
     addNode({
      id,
      type: 'llmNode',
      position,
      data: { 
        label: 'New Question', 
        model: 'gpt-4o', 
        temperature: 0.7,
        systemPrompt: '',
        userPrompt: '',
        mergeStrategy: 'smart',
        messages: [],
        traces: [],
        outgoingTraces: [],
        forkedTraces: [],
        response: '',
        status: 'idle',
        inputs: 1
      },
    });
  };

  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={onPaneContextMenu} // Use Right Click to add node for now
          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>
      </div>
      <Sidebar />
    </div>
  );
}

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