Update App.tsx

App.tsx

@@ -1,566 +1,22 @@
-
-import React, { useCallback, useRef, useState, useMemo, useEffect } from 'react';
-import ReactFlow, {
-  ReactFlowProvider,
-  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 './components/Sidebar';
-import PropertiesPanel from './components/PropertiesPanel';
-import CustomNode from './components/CustomNode';
-import CodeViewer from './components/CodeViewer';
-import AIBuilderModal from './components/AIBuilderModal';
-import SuggestionsModal from './components/SuggestionsModal';
-import ApiKeyModal from './components/ApiKeyModal';
-import FixerModal from './components/FixerModal';
+import React, { useState, Suspense } from 'react';
+import { ReactFlowProvider } from 'reactflow';
+import { Sparkles } from 'lucide-react';
 import LandingPage from './components/LandingPage';
-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;
-      })
-    );
-  };
+// Lazy load the heavy builder component to improve initial load performance
+const Builder = React.lazy(() => import('./components/Builder'));
 
-  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>
-  );
-};
+// Loading screen fallback for Suspense
+const LoadingScreen = () => (
+  <div className="fixed inset-0 bg-slate-950 flex items-center justify-center text-white z-50">
+     <div className="flex flex-col items-center gap-4">
+        <Sparkles className="animate-spin text-blue-500" size={32} />
+        <p className="text-slate-400 text-sm animate-pulse font-medium">Loading Builder Environment...</p>
+     </div>
+  </div>
+);
 
-const Main = () => {
+const App = () => {
   const [showLanding, setShowLanding] = useState(true);
 
   if (showLanding) {
@@ -569,9 +25,11 @@ const Main = () => {
 
   return (
     <ReactFlowProvider>
-      <Builder onBackToHome={() => setShowLanding(true)} />
+      <Suspense fallback={<LoadingScreen />}>
+        <Builder onBackToHome={() => setShowLanding(true)} />
+      </Suspense>
     </ReactFlowProvider>
   );
 };
 
-export default Main;
+export default App;