components/Builder.tsx

import React, { useCallback, useRef, useState, useMemo, useEffect } from 'react';
import ReactFlow, {
  addEdge,
  useNodesState,
  useEdgesState,
  Controls,
  Background,
  Connection,
  Edge,
  Node,
  BackgroundVariant,
  Panel,
  ConnectionMode
} from 'reactflow';
import { FileText, CheckCircle, AlertTriangle, X, Copy, Check, Lightbulb, Key, Wrench, Menu, Activity } from 'lucide-react';
import Sidebar from './Sidebar';
import PropertiesPanel from './PropertiesPanel';
import CustomNode from './CustomNode';
import CodeViewer from './CodeViewer';
import AIBuilderModal from './AIBuilderModal';
import SuggestionsModal from './SuggestionsModal';
import ApiKeyModal from './ApiKeyModal';
import FixerModal from './FixerModal';
import { INITIAL_NODES, INITIAL_EDGES, LAYER_DEFINITIONS, TEMPLATES } from '../constants';
import { generateRefinedPrompt, validateArchitecture, getArchitectureSuggestions, getUserApiKey, implementArchitectureSuggestions } from '../services/geminiService';
import { NodeData, LayerType, LogEntry } from '../types';

let id = 1000;
const getId = () => `${id++}`;

// --- MAIN BUILDER CONTENT ---
interface BuilderProps {
  onBackToHome: () => void;
}

const Builder: React.FC<BuilderProps> = ({ onBackToHome }) => {
  const reactFlowWrapper = useRef<HTMLDivElement>(null);
  const [nodes, setNodes, onNodesChange] = useNodesState(INITIAL_NODES);
  const [edges, setEdges, onEdgesChange] = useEdgesState(INITIAL_EDGES);
  const [reactFlowInstance, setReactFlowInstance] = useState<any>(null);
  const [selectedNodeId, setSelectedNodeId] = useState<string | null>(null);
  
  const [generatedPrompt, setGeneratedPrompt] = useState('');
  const [isPromptViewerOpen, setIsPromptViewerOpen] = useState(false);
  const [isGeneratingPrompt, setIsGeneratingPrompt] = useState(false);
  
  const [validationMsg, setValidationMsg] = useState<string | null>(null);
  const [layoutCopied, setLayoutCopied] = useState(false);

  // Responsive Layout State
  const [isSidebarOpen, setIsSidebarOpen] = useState(true);
  const [isPropertiesOpen, setIsPropertiesOpen] = useState(false);

  // Suggestions State
  const [isSuggestionsOpen, setIsSuggestionsOpen] = useState(false);
  const [suggestionsText, setSuggestionsText] = useState('');
  const [isSuggestionsLoading, setIsSuggestionsLoading] = useState(false);
  const [isImplementingSuggestions, setIsImplementingSuggestions] = useState(false);

  // Fixer State
  const [isFixerOpen, setIsFixerOpen] = useState(false);

  // AI Builder and Template State
  const [isAIBuilderOpen, setIsAIBuilderOpen] = useState(false);
  const [isTemplateMenuOpen, setIsTemplateMenuOpen] = useState(false);

  // API Key State
  const [isApiKeyModalOpen, setIsApiKeyModalOpen] = useState(false);
  const [isConnected, setIsConnected] = useState(false);

  // Activity Log State
  const [logs, setLogs] = useState<LogEntry[]>([]);

  const addLog = useCallback((message: string, type: 'info' | 'success' | 'warning' | 'error' = 'info') => {
    setLogs(prev => [{
      id: Date.now().toString() + Math.random(),
      timestamp: new Date(),
      message,
      type
    }, ...prev]);
  }, []);

  // Initialize responsiveness
  useEffect(() => {
    const handleResize = () => {
      if (window.innerWidth < 768) {
        setIsSidebarOpen(false);
      } else {
        setIsSidebarOpen(true);
      }
    };
    
    // Set initial state
    handleResize();
    
    window.addEventListener('resize', handleResize);
    return () => window.removeEventListener('resize', handleResize);
  }, []);

  // Load connection status on mount
  useEffect(() => {
    setIsConnected(!!getUserApiKey());
    addLog('System initialized. Ready to build.', 'info');
  }, [addLog]);

  // Define custom node types
  const nodeTypes = useMemo(() => ({ custom: CustomNode }), []);

  const onConnect = useCallback(
    (params: Connection) => setEdges((eds) => addEdge({ ...params, animated: true, style: { stroke: '#94a3b8' } }, eds)),
    [setEdges]
  );

  const onDragOver = useCallback((event: React.DragEvent) => {
    event.preventDefault();
    event.dataTransfer.dropEffect = 'move';
  }, []);

  const onDrop = useCallback(
    (event: React.DragEvent) => {
      event.preventDefault();

      if (!reactFlowWrapper.current || !reactFlowInstance) return;

      const type = event.dataTransfer.getData('application/reactflow') as LayerType;
      
      if (!type) return;

      const position = reactFlowInstance.screenToFlowPosition({
        x: event.clientX,
        y: event.clientY,
      });

      const label = LAYER_DEFINITIONS[type].label;
      const newNode: Node<NodeData> = {
        id: getId(),
        type: 'custom',
        position,
        data: { 
          label: label, 
          type: type, 
          params: LAYER_DEFINITIONS[type].parameters.reduce((acc, p) => ({...acc, [p.name]: p.default}), {}) 
        },
      };

      setNodes((nds) => nds.concat(newNode));
      setSelectedNodeId(newNode.id);
      setIsPropertiesOpen(true); // Open properties panel on drop
      
      // On mobile, close sidebar after dropping
      if (window.innerWidth < 768) {
        setIsSidebarOpen(false);
      }

      addLog(`Added layer: ${label}`, 'info');
    },
    [reactFlowInstance, setNodes, addLog]
  );

  const onNodeClick = useCallback((_: React.MouseEvent, node: Node) => {
    setSelectedNodeId(node.id);
    setIsPropertiesOpen(true);
  }, []);

  const onPaneClick = useCallback(() => {
    setSelectedNodeId(null);
    setIsPropertiesOpen(false);
  }, []);

  // Update node data (params or label) from Properties Panel
  const updateNodeData = (id: string, newData: Partial<NodeData>) => {
    setNodes((nds) =>
      nds.map((node) => {
        if (node.id === id) {
          return {
            ...node,
            data: { ...node.data, ...newData }
          };
        }
        return node;
      })
    );
  };

  const deleteNode = (id: string) => {
    const node = nodes.find(n => n.id === id);
    setNodes((nds) => nds.filter((node) => node.id !== id));
    setEdges((eds) => eds.filter((edge) => edge.source !== id && edge.target !== id));
    setSelectedNodeId(null);
    setIsPropertiesOpen(false);
    if (node) addLog(`Deleted layer: ${node.data.label}`, 'info');
  };

  const handleAPIError = (error: any) => {
      const msg = error.message || String(error);
      addLog(`API Error: ${msg}`, 'error');
      
      // Force disconnect if authentication failed to allow user to re-enter key
      if (msg.includes('API Key') || msg.includes('403') || msg.includes('401') || msg.includes('400')) {
          setIsConnected(false);
          setIsApiKeyModalOpen(true);
          addLog("Disconnected due to invalid API key. Please check your key.", 'warning');
      }
  }

  const handleGeneratePrompt = async () => {
    if (!isConnected) { setIsApiKeyModalOpen(true); return; }
    setIsGeneratingPrompt(true);
    setIsPromptViewerOpen(true);
    addLog("Generating PyTorch code prompt...", 'info');
    try {
        const promptText = await generateRefinedPrompt(nodes, edges);
        setGeneratedPrompt(promptText);
        addLog("Code prompt generated successfully.", 'success');
    } catch (e) {
        setGeneratedPrompt("Failed to generate prompt. See logs for details.");
        handleAPIError(e);
    } finally {
        setIsGeneratingPrompt(false);
    }
  };

  const handleValidate = async () => {
    if (!isConnected) { setIsApiKeyModalOpen(true); return; }
    setValidationMsg("Validating...");
    addLog("Running architecture validation...", 'info');
    try {
        const result = await validateArchitecture(nodes, edges);
        setValidationMsg(result);
        if (result.toLowerCase().includes("valid") && !result.toLowerCase().includes("invalid") && result.length < 50) {
            setTimeout(() => setValidationMsg(null), 5000);
            addLog("Validation passed: Architecture is valid.", 'success');
        } else {
            addLog("Validation issues found.", 'warning');
        }
    } catch (e) {
        setValidationMsg("Validation failed due to API error.");
        handleAPIError(e);
    }
  };

  const handleGetSuggestions = async () => {
    if (!isConnected) { setIsApiKeyModalOpen(true); return; }
    setIsSuggestionsOpen(true);
    setIsSuggestionsLoading(true);
    try {
      const suggestions = await getArchitectureSuggestions(nodes, edges);
      setSuggestionsText(suggestions);
    } catch (error) {
      setSuggestionsText("Failed to get suggestions. API Error.");
      handleAPIError(error);
    } finally {
      setIsSuggestionsLoading(false);
    }
  };

  const handleImplementSuggestions = async () => {
    if (!suggestionsText || isImplementingSuggestions) return;
    setIsImplementingSuggestions(true);
    addLog("AI Agents implementing suggestions...", 'info');
    try {
      const result = await implementArchitectureSuggestions(nodes, edges, suggestionsText);
      if (result && result.nodes) {
         // Fix types and defaults safely
         const processedNodes = result.nodes.map((n: any) => {
            const data = n.data || {};
            const type = data.type || n.type || 'Identity'; // Fallback type
            return {
              ...n,
              type: 'custom',
              data: { 
                ...data, 
                type: type,
                label: data.label || n.label || type, // Safe fallback for label
                params: data.params || {}
              }
            };
         });
         
         const processedEdges = result.edges.map((e: any) => ({
            ...e, animated: true, style: { stroke: '#94a3b8' }
         }));
         
         setNodes(processedNodes);
         setEdges(processedEdges);
         setIsSuggestionsOpen(false); // Close modal on success
         addLog("Implemented AI improvements.", 'success');
      }
    } catch (e) {
      handleAPIError(e);
      addLog("Failed to implement suggestions.", 'error');
    } finally {
      setIsImplementingSuggestions(false);
    }
  };

  const handleCopyLayout = () => {
    const layout = {
      nodes,
      edges
    };
    navigator.clipboard.writeText(JSON.stringify(layout, null, 2));
    setLayoutCopied(true);
    setTimeout(() => setLayoutCopied(false), 2000);
    addLog("Layout JSON copied to clipboard.", 'info');
  };

  const handleApplyAIBuild = (newNodes: any[], newEdges: any[], logMsg?: string) => {
    setNodes(newNodes);
    setEdges(newEdges);
    if (logMsg) {
      addLog(logMsg, 'success');
    }
  };

  const loadTemplate = (templateKey: string) => {
    const template = TEMPLATES[templateKey];
    if (template) {
      // Need to deep clone to avoid reference issues if loaded multiple times
      const clonedNodes = JSON.parse(JSON.stringify(template.nodes));
      const clonedEdges = JSON.parse(JSON.stringify(template.edges));
      
      // Update styling and functional props that aren't in JSON
      const hydratedNodes = clonedNodes.map((n: any) => ({
        ...n,
        data: {
          ...n.data,
          label: LAYER_DEFINITIONS[n.data.type as LayerType]?.label || n.data.label
        }
      }));

      const hydratedEdges = clonedEdges.map((e: any) => ({
        ...e,
        animated: true,
        style: { stroke: '#94a3b8' }
      }));

      setNodes(hydratedNodes);
      setEdges(hydratedEdges);
      setIsTemplateMenuOpen(false);
      
      // Auto open properties to see details of first node
      if (hydratedNodes.length > 0) {
        setSelectedNodeId(hydratedNodes[0].id);
        setIsPropertiesOpen(true);
      }
      addLog(`Loaded template: ${template.name}`, 'info');
    }
  };

  const selectedNode = useMemo(() => 
    nodes.find((n) => n.id === selectedNodeId) || null, 
    [nodes, selectedNodeId]
  );

  return (
    <div className="flex h-screen w-screen bg-slate-950 overflow-hidden">
      <Sidebar 
        onOpenAIBuilder={() => isConnected ? setIsAIBuilderOpen(true) : setIsApiKeyModalOpen(true)}
        onSelectTemplate={(id) => id === 'menu' ? setIsTemplateMenuOpen(true) : loadTemplate(id)}
        onBackToHome={onBackToHome}
        isConnected={isConnected}
        isOpen={isSidebarOpen}
        onToggle={() => setIsSidebarOpen(!isSidebarOpen)}
      />
      
      <div className="flex-1 h-full relative" ref={reactFlowWrapper}>
        <ReactFlow
          nodes={nodes}
          edges={edges}
          onNodesChange={onNodesChange}
          onEdgesChange={onEdgesChange}
          onConnect={onConnect}
          onInit={setReactFlowInstance}
          onDrop={onDrop}
          onDragOver={onDragOver}
          onNodeClick={onNodeClick}
          onPaneClick={onPaneClick}
          nodeTypes={nodeTypes}
          connectionMode={ConnectionMode.Loose}
          fitView
          className="bg-slate-950"
        >
          <Background color="#1e293b" variant={BackgroundVariant.Dots} gap={20} size={1} />
          <Controls className="!bg-slate-800 !border-slate-700 [&>button]:!fill-slate-300 [&>button:hover]:!bg-slate-700 md:bottom-4 md:left-4 bottom-24 left-4" />
          
          {/* Mobile Menu Trigger */}
          <Panel position="top-left" className="md:hidden">
             <button 
               onClick={() => setIsSidebarOpen(true)}
               className="p-2 bg-slate-800/80 backdrop-blur border border-slate-700 rounded-lg text-slate-300 shadow-lg"
             >
               <Menu size={20} />
             </button>
          </Panel>

          {/* Action Toolbar */}
          <Panel position="top-right" className="flex flex-wrap justify-end gap-2 max-w-[80vw]">
            <button 
              onClick={handleCopyLayout}
              className="flex items-center gap-2 px-3 py-2 bg-slate-800/80 backdrop-blur hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-xs md:text-sm font-medium transition-colors shadow-lg"
              title="Copy Layout JSON"
            >
              {layoutCopied ? <Check size={16} className="text-emerald-500" /> : <Copy size={16} className="text-slate-400" />}
              <span className="hidden md:inline">{layoutCopied ? 'Copied' : 'Copy'}</span>
            </button>
            <button
               onClick={() => setIsApiKeyModalOpen(true)}
               className="flex items-center gap-2 px-3 py-2 bg-slate-800/80 backdrop-blur hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-xs md:text-sm font-medium transition-colors shadow-lg"
               title="API Key Settings"
            >
              <Key size={16} className={isConnected ? "text-emerald-500" : "text-slate-400"} />
              <span className="hidden md:inline">{isConnected ? 'Key Active' : 'Add Key'}</span>
            </button>
            <button 
              onClick={handleGetSuggestions}
              className="flex items-center gap-2 px-3 py-2 bg-slate-800/80 backdrop-blur hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-xs md:text-sm font-medium transition-colors shadow-lg"
              title="Get AI Suggestions"
            >
              <Lightbulb size={16} className="text-amber-400" />
              <span className="hidden md:inline">Tips</span>
            </button>
            <button 
              onClick={handleValidate}
              className="flex items-center gap-2 px-3 py-2 bg-slate-800/80 backdrop-blur hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 text-xs md:text-sm font-medium transition-colors shadow-lg"
              title="Validate Architecture"
            >
              <CheckCircle size={16} className="text-emerald-500" />
              <span className="hidden md:inline">Validate</span>
            </button>
            <button 
              onClick={handleGeneratePrompt}
              className="flex items-center gap-2 px-3 py-2 bg-blue-600/90 hover:bg-blue-500 backdrop-blur rounded-lg text-white text-xs md:text-sm font-medium transition-colors shadow-lg shadow-blue-900/20"
              title="Generate Code Prompt"
            >
              <FileText size={16} />
              <span className="hidden md:inline">Code</span>
            </button>
             {/* Mobile Activity Toggle */}
             <button 
              onClick={() => { setSelectedNodeId(null); setIsPropertiesOpen(true); }}
              className="md:hidden flex items-center gap-2 px-3 py-2 bg-slate-800/80 backdrop-blur hover:bg-slate-700 border border-slate-700 rounded-lg text-slate-200 transition-colors shadow-lg"
              title="System Logs"
            >
              <Activity size={16} className="text-blue-400" />
            </button>
          </Panel>

          {validationMsg && (
            <Panel position="bottom-center" className="mb-20 md:mb-8 w-full flex justify-center px-4">
              <div className="bg-slate-900/95 backdrop-blur border border-slate-700 text-slate-200 px-6 py-4 rounded-xl shadow-2xl max-w-2xl w-full flex flex-col gap-3 animate-in slide-in-from-bottom-5">
                <div className="flex items-start gap-3">
                  <AlertTriangle className="text-amber-500 shrink-0 mt-0.5" size={20} />
                  <div className="flex-1 min-w-0">
                     <h4 className="font-bold text-sm text-slate-300 mb-1">Validation Report</h4>
                     <p className="text-sm whitespace-pre-wrap leading-relaxed text-slate-400 max-h-[30vh] overflow-y-auto scrollbar-thin scrollbar-thumb-slate-700 pr-2">{validationMsg}</p>
                  </div>
                </div>
                
                <div className="flex justify-end gap-3 mt-2">
                   <button 
                    onClick={() => setValidationMsg(null)}
                    className="px-4 py-1.5 text-xs font-medium text-slate-400 hover:text-slate-200 bg-slate-800 hover:bg-slate-700 rounded transition-colors"
                  >
                    Dismiss
                  </button>
                  {!validationMsg.toLowerCase().includes("architecture is valid") && (
                    <button 
                      onClick={() => setIsFixerOpen(true)}
                      className="flex items-center gap-2 px-4 py-1.5 bg-red-600 hover:bg-red-500 text-white rounded text-xs font-medium transition-colors shadow-lg shadow-red-900/20"
                    >
                      <Wrench size={14} />
                      Fix Errors
                    </button>
                  )}
                </div>
              </div>
            </Panel>
          )}

          {/* Template Selection Modal Overlay */}
          {isTemplateMenuOpen && (
             <div className="absolute inset-0 z-50 flex items-center justify-center bg-black/60 backdrop-blur-sm p-4">
                <div className="bg-slate-900 border border-slate-700 rounded-xl shadow-2xl w-96 max-w-full flex flex-col max-h-[80vh] overflow-hidden animate-in zoom-in-95 duration-200">
                  <div className="flex justify-between items-center p-6 border-b border-slate-800 shrink-0 bg-slate-900">
                     <h3 className="text-lg font-bold text-white">Select Template</h3>
                     <button onClick={() => setIsTemplateMenuOpen(false)}><X size={20} className="text-slate-500 hover:text-white"/></button>
                  </div>
                  <div className="overflow-y-auto p-6 scrollbar-thin scrollbar-thumb-slate-700">
                    <div className="grid gap-3">
                      {Object.entries(TEMPLATES).map(([key, t]) => (
                        <button 
                          key={key}
                          onClick={() => loadTemplate(key)}
                          className="text-left p-3 rounded bg-slate-800 hover:bg-slate-750 border border-slate-700 hover:border-blue-500/50 transition-all group"
                        >
                           <div className="font-semibold text-slate-200 group-hover:text-blue-400">{t.name}</div>
                           <div className="text-xs text-slate-500">{t.description}</div>
                        </button>
                      ))}
                    </div>
                  </div>
                </div>
             </div>
          )}

        </ReactFlow>
      </div>

      <PropertiesPanel 
        selectedNode={selectedNode} 
        onChange={updateNodeData}
        onDelete={deleteNode}
        onClose={() => { setSelectedNodeId(null); setIsPropertiesOpen(false); }}
        logs={logs}
        isOpen={isPropertiesOpen}
      />

      <CodeViewer 
        isOpen={isPromptViewerOpen} 
        onClose={() => setIsPromptViewerOpen(false)} 
        code={generatedPrompt}
        isLoading={isGeneratingPrompt}
      />
      
      <SuggestionsModal
        isOpen={isSuggestionsOpen}
        onClose={() => setIsSuggestionsOpen(false)}
        suggestions={suggestionsText}
        isLoading={isSuggestionsLoading}
        onImplement={handleImplementSuggestions}
        isImplementing={isImplementingSuggestions}
      />

      <AIBuilderModal 
        isOpen={isAIBuilderOpen}
        onClose={() => setIsAIBuilderOpen(false)}
        onApply={handleApplyAIBuild}
        currentNodes={nodes}
      />

      <FixerModal
        isOpen={isFixerOpen}
        onClose={() => setIsFixerOpen(false)}
        onApply={handleApplyAIBuild}
        errorMsg={validationMsg || ''}
        nodes={nodes}
        edges={edges}
      />

      <ApiKeyModal 
        isOpen={isApiKeyModalOpen}
        onClose={() => setIsApiKeyModalOpen(false)}
        onSuccess={() => setIsConnected(true)}
      />
    </div>
  );
};

export default Builder;