import { useCallback, useEffect, useRef, useState } from 'react';
import ReactFlow, { 
  Background, 
  Controls, 
  MiniMap, 
  ReactFlowProvider,
  Panel,
  useReactFlow,
  SelectionMode,
  type Node,
  type Edge
} from 'reactflow';
import 'reactflow/dist/style.css';
import useFlowStore from './store/flowStore';
import { useAuthStore } from './store/authStore';
import LLMNode from './components/nodes/LLMNode';
import MergedEdge from './components/edges/MergedEdge';
import Sidebar from './components/Sidebar';
import LeftSidebar from './components/LeftSidebar';
import { ContextMenu } from './components/ContextMenu';
import { Plus, Sun, Moon, LayoutGrid, Loader2 } from 'lucide-react';
import AuthPage from './pages/AuthPage';

const nodeTypes = {
  llmNode: LLMNode,
};

const edgeTypes = {
  merged: MergedEdge,
};

function Flow() {
  const { 
    nodes, 
    edges, 
    onNodesChange, 
    onEdgesChange, 
    onConnect, 
    addNode,
    deleteEdge,
    deleteNode,
    deleteBranch,
    deleteTrace,
    setSelectedNode,
    toggleNodeDisabled,
    archiveNode,
    createNodeFromArchive,
    toggleTraceDisabled,
    theme,
    toggleTheme,
    autoLayout,
    findNonOverlappingPosition,
    setLastViewport,
    saveCurrentBlueprint,
    currentBlueprintPath
  } = useFlowStore();

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

  // Get selected nodes
  const selectedNodes = nodes.filter(n => n.selected);

  const onPaneClick = useCallback(() => {
    setSelectedNode(null);
    setMenu(null);
  }, [setSelectedNode]);

  // Close menu on various interactions
  const closeMenu = useCallback(() => {
    setMenu(null);
  }, []);

  const handleNodeDragStart = useCallback(() => {
    setMenu(null);
  }, []);

  const handleMoveStart = useCallback(() => {
    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();
     // Check if multiple nodes are selected and the right-clicked node is one of them
     if (selectedNodes.length > 1 && selectedNodes.some(n => n.id === node.id)) {
       setMenu({ 
         x: event.clientX, 
         y: event.clientY, 
         type: 'multiselect', 
         selectedIds: selectedNodes.map(n => n.id) 
       });
     } else {
       setMenu({ x: event.clientX, y: event.clientY, type: 'node', id: node.id });
     }
  };

  // Batch operations for multi-select
  const handleBatchDelete = (nodeIds: string[]) => {
    nodeIds.forEach(id => deleteNode(id));
    setMenu(null);
  };

  const handleBatchDisable = (nodeIds: string[]) => {
    nodeIds.forEach(id => toggleNodeDisabled(id));
    setMenu(null);
  };

  const handleBatchArchive = (nodeIds: string[]) => {
    nodeIds.forEach(id => archiveNode(id));
    setMenu(null);
  };

  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()}`;
    
    // If no position provided, use viewport center
    let basePos = position;
    if (!basePos && reactFlowWrapper.current) {
      const { x, y, zoom } = getViewport();
      const rect = reactFlowWrapper.current.getBoundingClientRect();
      // Calculate center of viewport in flow coordinates
      basePos = {
        x: (-x + rect.width / 2) / zoom - 100, // offset by half node width
        y: (-y + rect.height / 2) / zoom - 40   // offset by half node height
      };
    }
    basePos = basePos || { x: 200, y: 200 };
    const pos = findNonOverlappingPosition(basePos.x, basePos.y);
    
    addNode({
      id,
      type: 'llmNode',
      position: pos,
      data: { 
        label: 'New Question', 
        model: 'gpt-5.1', 
        temperature: 0.7,
        systemPrompt: '',
        userPrompt: '',
        mergeStrategy: 'smart',
        reasoningEffort: 'medium',  // Default for reasoning models
        messages: [],
        traces: [],
        outgoingTraces: [],
        forkedTraces: [],
        mergedTraces: [],
        response: '',
        status: 'idle',
        inputs: 1,
        attachedFileIds: [],
        activeTraceIds: []
      },
    });
    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);
  };

  // Ctrl/Cmd + S manual save
  useEffect(() => {
    const handler = (e: KeyboardEvent) => {
      const isMac = navigator.platform.toLowerCase().includes('mac');
      if ((isMac ? e.metaKey : e.ctrlKey) && e.key.toLowerCase() === 's') {
        e.preventDefault();
        const path = currentBlueprintPath;
        if (path) {
          saveCurrentBlueprint(path, getViewport());
        }
      }
    };
    window.addEventListener('keydown', handler);
    return () => window.removeEventListener('keydown', handler);
  }, [currentBlueprintPath, saveCurrentBlueprint, getViewport]);

  return (
    <div className={`w-screen h-screen flex ${theme === 'dark' ? 'dark bg-gray-900' : 'bg-white'}`}>
      <LeftSidebar isOpen={isLeftOpen} onToggle={() => setIsLeftOpen(!isLeftOpen)} />
      
      <div 
        className={`flex-1 h-full relative ${theme === 'dark' ? 'bg-slate-900' : 'bg-slate-50'}`}
        ref={reactFlowWrapper}
        onDragOver={onDragOver}
        onDrop={onDrop}
      >
        <ReactFlow
          nodes={nodes}
          edges={edges}
          onNodesChange={onNodesChange}
          onEdgesChange={onEdgesChange}
          onConnect={onConnect}
          onMoveEnd={(_, viewport) => setLastViewport(viewport)}
          nodeTypes={nodeTypes}
          edgeTypes={edgeTypes}
          defaultEdgeOptions={{ type: 'merged' }}
          onNodeClick={(e, node) => { closeMenu(); onNodeClick(e, node); }}
          onPaneClick={onPaneClick}
          onPaneContextMenu={handlePaneContextMenu}
          onNodeContextMenu={handleNodeContextMenu}
          onEdgeContextMenu={handleEdgeContextMenu}
          onNodeDragStart={handleNodeDragStart}
          onMoveStart={handleMoveStart}
          onSelectionStart={closeMenu}
          fitView
          panOnDrag
          selectionOnDrag
          selectionKeyCode="Shift"
          multiSelectionKeyCode="Shift"
          selectionMode={SelectionMode.Partial}
        >
          <Background color={theme === 'dark' ? '#374151' : '#aaa'} gap={16} />
          <Controls className={theme === 'dark' ? 'dark-controls' : ''} />
          <MiniMap 
            nodeColor={theme === 'dark' ? '#4b5563' : '#e5e7eb'}
            maskColor={theme === 'dark' ? 'rgba(0,0,0,0.6)' : 'rgba(255,255,255,0.6)'}
          />
          <Panel position="top-left">
            <div className="flex gap-2">
              <button 
                onClick={() => handleAddNode()}
                className={`px-4 py-2 rounded-md shadow-md font-medium flex items-center gap-2 transition-colors ${
                  theme === 'dark' 
                    ? 'bg-gray-800 text-gray-200 hover:bg-gray-700 border border-gray-600' 
                    : 'bg-white text-gray-700 hover:bg-gray-50'
                }`}
              >
                <Plus size={16} /> Add Node
              </button>
              <button 
                onClick={autoLayout}
                className={`px-3 py-2 rounded-md shadow-md font-medium flex items-center gap-2 transition-colors ${
                  theme === 'dark' 
                    ? 'bg-gray-800 text-gray-200 hover:bg-gray-700 border border-gray-600' 
                    : 'bg-white text-gray-700 hover:bg-gray-50'
                }`}
                title="Auto Layout"
              >
                <LayoutGrid size={16} />
              </button>
              <button 
                onClick={toggleTheme}
                className={`px-3 py-2 rounded-md shadow-md font-medium flex items-center gap-2 transition-colors ${
                  theme === 'dark' 
                    ? 'bg-gray-800 text-gray-200 hover:bg-gray-700 border border-gray-600' 
                    : 'bg-white text-gray-700 hover:bg-gray-50'
                }`}
                title={theme === 'dark' ? 'Switch to Light Mode' : 'Switch to Dark Mode'}
              >
                {theme === 'dark' ? <Sun size={16} /> : <Moon size={16} />}
              </button>
            </div>
          </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 === 'multiselect' ? (() => {
                // Multi-select menu - batch operations
                const ids = menu.selectedIds || [];
                const allDisabled = ids.every(id => nodes.find(n => n.id === id)?.data?.disabled);
                
                return [
                  {
                    label: allDisabled ? `Enable ${ids.length} Nodes` : `Disable ${ids.length} Nodes`,
                    onClick: () => handleBatchDisable(ids)
                  },
                  {
                    label: `Archive ${ids.length} Nodes`,
                    onClick: () => handleBatchArchive(ids)
                  },
                  {
                    label: `Delete ${ids.length} Nodes`,
                    danger: true,
                    onClick: () => handleBatchDelete(ids)
                  }
                ];
              })() : 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)
                  }
                ];
              })() : menu.type === 'edge' ? (() => {
                // 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) 
                  },
                  { 
                    label: 'Delete Trace', 
                    danger: true,
                    onClick: () => menu.id && deleteTrace(menu.id) 
                  }
                ];
              })() : []
            }
          />
        )}
      </div>
      <Sidebar isOpen={isRightOpen} onToggle={() => setIsRightOpen(!isRightOpen)} onInteract={closeMenu} />
    </div>
  );
}

function AuthWrapper() {
  const { isAuthenticated, checkAuth, user } = useAuthStore();
  const { refreshProjectTree, refreshFiles, loadArchivedNodes, clearBlueprint } = useFlowStore();
  const [checking, setChecking] = useState(true);
  const [initializing, setInitializing] = useState(false);
  const prevUserRef = useRef<string | null>(null);

  useEffect(() => {
    checkAuth().finally(() => setChecking(false));
  }, [checkAuth]);

  // When user changes (login/logout), refresh all user-specific data
  useEffect(() => {
    const currentUsername = user?.username || null;
    
    // If user changed, reload everything
    if (isAuthenticated && currentUsername && currentUsername !== prevUserRef.current) {
      setInitializing(true);
      prevUserRef.current = currentUsername;
      
      // Clear old data and load new user's data
      clearBlueprint();
      Promise.all([
        refreshProjectTree(),
        refreshFiles(),
        loadArchivedNodes()
      ]).finally(() => setInitializing(false));
    } else if (!isAuthenticated) {
      prevUserRef.current = null;
    }
  }, [isAuthenticated, user, refreshProjectTree, refreshFiles, loadArchivedNodes, clearBlueprint]);

  if (checking || initializing) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-gray-900">
        <div className="flex flex-col items-center gap-4">
          <Loader2 className="animate-spin text-blue-500" size={48} />
          <p className="text-gray-400">{initializing ? 'Loading workspace...' : 'Loading...'}</p>
        </div>
      </div>
    );
  }

  if (!isAuthenticated) {
    return <AuthPage />;
  }

  return (
    <ReactFlowProvider>
      <Flow />
    </ReactFlowProvider>
  );
}

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