Create Builder.tsx

components/Builder.tsx

@@ -0,0 +1,560 @@
+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;