Atomizer/atomizer-dashboard/frontend/src/pages/CanvasView.tsx

import { useState, useEffect, useCallback } from 'react';
import { useNavigate, useParams, useSearchParams } from 'react-router-dom';
import { ClipboardList, Download, Trash2, Layers, Home, ChevronRight, Save, RefreshCw, Zap, MessageSquare, X, Folder, SlidersHorizontal } from 'lucide-react';
import { AtomizerCanvas } from '../components/canvas/AtomizerCanvas';
import { SpecRenderer } from '../components/canvas/SpecRenderer';
import { NodePalette } from '../components/canvas/palette/NodePalette';
import { FileStructurePanel } from '../components/canvas/panels/FileStructurePanel';
import { TemplateSelector } from '../components/canvas/panels/TemplateSelector';
import { ConfigImporter } from '../components/canvas/panels/ConfigImporter';
import { NodeConfigPanel } from '../components/canvas/panels/NodeConfigPanel';
import { NodeConfigPanelV2 } from '../components/canvas/panels/NodeConfigPanelV2';
import { ChatPanel } from '../components/canvas/panels/ChatPanel';
import { useCanvasStore } from '../hooks/useCanvasStore';
import { useSpecStore, useSpec, useSpecLoading, useSpecIsDirty, useSelectedNodeId } from '../hooks/useSpecStore';
import { useStudy } from '../context/StudyContext';
import { useChat } from '../hooks/useChat';
import { CanvasTemplate } from '../lib/canvas/templates';

export function CanvasView() {
  const [showTemplates, setShowTemplates] = useState(false);
  const [showImporter, setShowImporter] = useState(false);
  const [showChat, setShowChat] = useState(true);
  const [chatPowerMode, setChatPowerMode] = useState(false);
  const [notification, setNotification] = useState<string | null>(null);
  const [isSaving, setIsSaving] = useState(false);
  const [isLoading, setIsLoading] = useState(false);
  const [hasUnsavedChanges, setHasUnsavedChanges] = useState(false);
  const [paletteCollapsed, setPaletteCollapsed] = useState(false);
  const [leftSidebarTab, setLeftSidebarTab] = useState<'components' | 'files'>('components');
  const navigate = useNavigate();
  const [searchParams] = useSearchParams();

  // Spec mode is the default (AtomizerSpec v2.0)
  // Legacy mode can be enabled via:
  //   1. VITE_USE_LEGACY_CANVAS=true environment variable
  //   2. ?mode=legacy query param (for emergency fallback)
  const legacyEnvEnabled = import.meta.env.VITE_USE_LEGACY_CANVAS === 'true';
  const legacyQueryParam = searchParams.get('mode') === 'legacy';
  const useSpecMode = !legacyEnvEnabled && !legacyQueryParam;

  // Get study ID from URL params (supports nested paths like M1_Mirror/study_name)
  const { '*': urlStudyId } = useParams<{ '*': string }>();

  // Legacy canvas store (for backwards compatibility)
  const { nodes, edges, clear, loadFromConfig, toIntent } = useCanvasStore();

  // New spec store (AtomizerSpec v2.0)
  const spec = useSpec();
  const specLoading = useSpecLoading();
  const specIsDirty = useSpecIsDirty();
  const selectedNodeId = useSelectedNodeId();
  const { loadSpec, saveSpec, reloadSpec } = useSpecStore();

  const { setSelectedStudy, studies } = useStudy();
  const { clearSpec, setSpecFromWebSocket } = useSpecStore();

  // Active study ID comes ONLY from URL - don't auto-load from context
  // This ensures /canvas shows empty canvas, /canvas/{id} shows the study
  const activeStudyId = urlStudyId;

  // Chat hook for assistant panel
  const { messages, isThinking, isConnected, sendMessage, notifyCanvasEdit } = useChat({
    studyId: activeStudyId,
    mode: chatPowerMode ? 'power' : 'user',
    useWebSocket: true,
    onCanvasModification: chatPowerMode ? (modification) => {
      // Handle canvas modifications from Claude in power mode (legacy)
      console.log('Canvas modification from Claude:', modification);
      showNotification(`Claude: ${modification.action} ${modification.nodeType || modification.nodeId || ''}`);
      // The actual modification is handled by the MCP tools on the backend
      // which update atomizer_spec.json, then the canvas reloads via WebSocket
      reloadSpec();
    } : undefined,
    onSpecUpdated: useSpecMode ? (newSpec) => {
      // Direct spec update from Claude via WebSocket - no HTTP reload needed
      console.log('Spec updated from Claude via WebSocket:', newSpec.meta?.study_name);
      setSpecFromWebSocket(newSpec, activeStudyId);
      showNotification('Canvas synced with Claude');
    } : undefined,
  });

  // Load or clear spec based on URL study ID
  useEffect(() => {
    if (urlStudyId) {
      if (useSpecMode) {
        // Try to load spec first, fall back to legacy config
        loadSpec(urlStudyId).catch(() => {
          // If spec doesn't exist, try legacy config
          loadStudyConfig(urlStudyId);
        });
      } else {
        loadStudyConfig(urlStudyId);
      }
    } else {
      // No study ID in URL - clear spec for empty canvas (new study creation)
      if (useSpecMode) {
        clearSpec();
      } else {
        clear();
      }
    }
  }, [urlStudyId, useSpecMode]);

  // Notify Claude when user edits the spec (bi-directional sync)
  // This sends the updated spec to Claude so it knows what the user changed
  useEffect(() => {
    if (useSpecMode && spec && specIsDirty && chatPowerMode) {
      // User made changes - notify Claude via WebSocket
      notifyCanvasEdit(spec);
    }
  }, [spec, specIsDirty, useSpecMode, chatPowerMode, notifyCanvasEdit]);

  // Track unsaved changes (legacy mode only)
  useEffect(() => {
    if (!useSpecMode && activeStudyId && nodes.length > 0) {
      setHasUnsavedChanges(true);
    }
  }, [nodes, edges, useSpecMode]);

  const loadStudyConfig = async (studyId: string) => {
    setIsLoading(true);
    try {
      const response = await fetch(`/api/optimization/studies/${encodeURIComponent(studyId)}/config`);
      if (!response.ok) {
        throw new Error(`Failed to load study: ${response.status}`);
      }
      const data = await response.json();
      loadFromConfig(data.config);
      setHasUnsavedChanges(false);

      // Also select the study in context
      const study = studies.find(s => s.id === studyId);
      if (study) {
        setSelectedStudy(study);
      }

      showNotification(`Loaded: ${studyId}`);
    } catch (error) {
      console.error('Failed to load study config:', error);
      showNotification('Failed to load study config');
    } finally {
      setIsLoading(false);
    }
  };

  const saveToConfig = async () => {
    if (!activeStudyId) {
      showNotification('No study selected to save to');
      return;
    }

    setIsSaving(true);
    try {
      if (useSpecMode && spec) {
        // Save spec using new API
        await saveSpec(spec);
        showNotification('Saved to atomizer_spec.json');
      } else {
        // Legacy save
        const intent = toIntent();

        const response = await fetch(`/api/optimization/studies/${encodeURIComponent(activeStudyId)}/config`, {
          method: 'PUT',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({ intent }),
        });

        if (!response.ok) {
          const error = await response.json();
          throw new Error(error.detail || 'Failed to save');
        }

        setHasUnsavedChanges(false);
        showNotification('Saved to optimization_config.json');
      }
    } catch (error) {
      console.error('Failed to save:', error);
      showNotification('Failed to save: ' + (error instanceof Error ? error.message : 'Unknown error'));
    } finally {
      setIsSaving(false);
    }
  };

  const handleTemplateSelect = (template: CanvasTemplate) => {
    setHasUnsavedChanges(true);
    showNotification(`Loaded template: ${template.name}`);
  };

  const handleImport = (source: string) => {
    setHasUnsavedChanges(true);
    showNotification(`Imported from ${source}`);
  };

  const handleClear = () => {
    if (useSpecMode) {
      // In spec mode, clearing is not typically needed since changes sync automatically
      showNotification('Use Reload to reset to saved state');
      return;
    }

    if (nodes.length === 0 || window.confirm('Clear all nodes from the canvas?')) {
      clear();
      setHasUnsavedChanges(true);
      showNotification('Canvas cleared');
    }
  };

  const handleReload = () => {
    if (activeStudyId) {
      const hasChanges = useSpecMode ? specIsDirty : hasUnsavedChanges;
      if (hasChanges && !window.confirm('Reload will discard unsaved changes. Continue?')) {
        return;
      }

      if (useSpecMode) {
        reloadSpec();
        showNotification('Reloaded from atomizer_spec.json');
      } else {
        loadStudyConfig(activeStudyId);
      }
    }
  };

  const showNotification = (message: string) => {
    setNotification(message);
    setTimeout(() => setNotification(null), 3000);
  };

  // Navigate to canvas with study ID
  const navigateToStudy = useCallback((studyId: string) => {
    navigate(`/canvas/${studyId}`);
  }, [navigate]);

  return (
    <div className="h-screen flex flex-col bg-dark-900">
      {/* Minimal Header */}
      <header className="flex-shrink-0 h-12 bg-dark-850 border-b border-dark-700 px-4 flex items-center justify-between">
        <div className="flex items-center gap-3">
          {/* Home button */}
          <button
            onClick={() => navigate('/')}
            className="p-1.5 rounded-lg text-dark-400 hover:text-white hover:bg-dark-700 transition-colors"
            title="Back to Home"
          >
            <Home size={18} />
          </button>

          {/* Breadcrumb */}
          <div className="flex items-center gap-2">
            <Layers size={18} className="text-primary-400" />
            <span className="text-sm font-medium text-white">Canvas Builder</span>
            {activeStudyId && (
              <>
                <ChevronRight size={14} className="text-dark-500" />
                <span className="text-sm text-primary-400 font-medium">
                  {activeStudyId}
                </span>
                {hasUnsavedChanges && (
                  <span className="text-xs text-amber-400 ml-1" title="Unsaved changes">•</span>
                )}
              </>
            )}
          </div>

          {/* Stats */}
          {useSpecMode && spec ? (
            <span className="text-xs text-dark-500 tabular-nums ml-2">
              {spec.design_variables.length} vars • {spec.extractors.length} ext • {spec.objectives.length} obj
            </span>
          ) : (
            <span className="text-xs text-dark-500 tabular-nums ml-2">
              {nodes.length} node{nodes.length !== 1 ? 's' : ''} • {edges.length} edge{edges.length !== 1 ? 's' : ''}
            </span>
          )}

          {/* Mode indicator */}
          {useSpecMode && (
            <span className="ml-2 px-1.5 py-0.5 text-xs bg-primary-900/50 text-primary-400 rounded border border-primary-800 flex items-center gap-1">
              <Zap size={10} />
              v2.0
            </span>
          )}

          {(isLoading || specLoading) && (
            <RefreshCw size={14} className="text-primary-400 animate-spin ml-2" />
          )}
        </div>

        {/* Action Buttons */}
        <div className="flex items-center gap-2">
          {/* Save Button - only show when there's a study and changes */}
          {activeStudyId && (
            <button
              onClick={saveToConfig}
              disabled={isSaving || (useSpecMode ? !specIsDirty : !hasUnsavedChanges)}
              className={`px-3 py-1.5 text-sm rounded-lg transition-colors flex items-center gap-1.5 ${
                (useSpecMode ? specIsDirty : hasUnsavedChanges)
                  ? 'bg-green-600 hover:bg-green-500 text-white'
                  : 'bg-dark-700 text-dark-400 cursor-not-allowed border border-dark-600'
              }`}
              title={(useSpecMode ? specIsDirty : hasUnsavedChanges) ? `Save changes to ${useSpecMode ? 'atomizer_spec.json' : 'optimization_config.json'}` : 'No changes to save'}
            >
              <Save size={14} />
              {isSaving ? 'Saving...' : 'Save'}
            </button>
          )}

          {/* Reload Button */}
          {activeStudyId && (
            <button
              onClick={handleReload}
              disabled={isLoading || specLoading}
              className="px-3 py-1.5 bg-dark-700 text-dark-200 hover:bg-dark-600 hover:text-white text-sm rounded-lg transition-colors flex items-center gap-1.5 border border-dark-600"
              title={`Reload from ${useSpecMode ? 'atomizer_spec.json' : 'optimization_config.json'}`}
            >
              <RefreshCw size={14} className={(isLoading || specLoading) ? 'animate-spin' : ''} />
              Reload
            </button>
          )}

          <button
            onClick={() => setShowTemplates(true)}
            className="px-3 py-1.5 bg-primary-600 hover:bg-primary-500 text-white text-sm rounded-lg transition-colors flex items-center gap-1.5"
          >
            <ClipboardList size={14} />
            Templates
          </button>
          <button
            onClick={() => setShowImporter(true)}
            className="px-3 py-1.5 bg-dark-700 text-dark-200 hover:bg-dark-600 hover:text-white text-sm rounded-lg transition-colors flex items-center gap-1.5 border border-dark-600"
          >
            <Download size={14} />
            Import
          </button>
          <button
            onClick={handleClear}
            className="px-3 py-1.5 bg-dark-700 text-dark-200 hover:bg-red-900/50 hover:text-red-400 text-sm rounded-lg transition-colors flex items-center gap-1.5 border border-dark-600"
          >
            <Trash2 size={14} />
            Clear
          </button>

          {/* Divider */}
          <div className="w-px h-6 bg-dark-600" />

          {/* Chat Toggle */}
          <button
            onClick={() => setShowChat(!showChat)}
            className={`px-3 py-1.5 text-sm rounded-lg transition-colors flex items-center gap-1.5 border ${
              showChat
                ? 'bg-primary-600 text-white border-primary-500'
                : 'bg-dark-700 text-dark-200 hover:bg-dark-600 hover:text-white border-dark-600'
            }`}
            title={showChat ? 'Hide Assistant' : 'Show Assistant'}
          >
            <MessageSquare size={14} />
            Assistant
          </button>
        </div>
      </header>

      {/* Main Canvas */}
      <main className="flex-1 overflow-hidden flex">
        {/* Left Sidebar with tabs (spec mode only - AtomizerCanvas has its own) */}
        {useSpecMode && (
          <div className={`${paletteCollapsed ? 'w-14' : 'w-60'} bg-dark-850 border-r border-dark-700 flex flex-col transition-all duration-200`}>
            {/* Tab buttons (only show when expanded) */}
            {!paletteCollapsed && (
              <div className="flex border-b border-dark-700">
                <button
                  onClick={() => setLeftSidebarTab('components')}
                  className={`flex-1 flex items-center justify-center gap-2 px-3 py-2.5 text-xs font-medium transition-colors
                    ${leftSidebarTab === 'components' 
                      ? 'text-primary-400 border-b-2 border-primary-400 -mb-px bg-dark-800/50' 
                      : 'text-dark-400 hover:text-white'}`}
                >
                  <SlidersHorizontal size={14} />
                  Components
                </button>
                <button
                  onClick={() => setLeftSidebarTab('files')}
                  className={`flex-1 flex items-center justify-center gap-2 px-3 py-2.5 text-xs font-medium transition-colors
                    ${leftSidebarTab === 'files' 
                      ? 'text-primary-400 border-b-2 border-primary-400 -mb-px bg-dark-800/50' 
                      : 'text-dark-400 hover:text-white'}`}
                >
                  <Folder size={14} />
                  Files
                </button>
              </div>
            )}

            {/* Tab content */}
            <div className="flex-1 overflow-hidden">
              {leftSidebarTab === 'components' || paletteCollapsed ? (
                <NodePalette
                  collapsed={paletteCollapsed}
                  onToggleCollapse={() => setPaletteCollapsed(!paletteCollapsed)}
                  showToggle={true}
                />
              ) : (
                <FileStructurePanel
                  studyId={activeStudyId || null}
                  selectedModelPath={spec?.model?.sim?.path}
                  onModelSelect={(path, _type) => {
                    // TODO: Update model path in spec
                    showNotification(`Selected: ${path.split(/[/\\]/).pop()}`);
                  }}
                />
              )}
            </div>
          </div>
        )}

        {/* Canvas area - must have explicit height for ReactFlow */}
        <div className="flex-1 h-full">
          {useSpecMode ? (
            <SpecRenderer
              studyId={activeStudyId}
              onStudyChange={navigateToStudy}
              enableWebSocket={true}
              showConnectionStatus={true}
              editable={true}
            />
          ) : (
            <AtomizerCanvas
              studyId={activeStudyId}
              onStudyChange={navigateToStudy}
            />
          )}
        </div>

        {/* Config Panel - use V2 for spec mode, legacy for AtomizerCanvas */}
        {selectedNodeId && !showChat && (
          useSpecMode ? (
            <NodeConfigPanelV2 onClose={() => useSpecStore.getState().clearSelection()} />
          ) : (
            <div className="w-80 border-l border-dark-700 bg-dark-850 overflow-y-auto">
              <NodeConfigPanel nodeId={selectedNodeId} />
            </div>
          )
        )}

        {/* Chat/Assistant Panel */}
        {showChat && (
          <div className="w-96 border-l border-dark-700 bg-dark-850 flex flex-col">
            {/* Chat Header */}
            <div className="flex items-center justify-between px-4 py-3 border-b border-dark-700">
              <div className="flex items-center gap-2">
                <MessageSquare size={16} className="text-primary-400" />
                <span className="font-medium text-white">Assistant</span>
                {isConnected && (
                  <span className="w-2 h-2 rounded-full bg-green-400" title="Connected" />
                )}
              </div>
              <div className="flex items-center gap-2">
                {/* Power Mode Toggle */}
                <button
                  onClick={() => setChatPowerMode(!chatPowerMode)}
                  className={`px-2 py-1 text-xs rounded transition-colors flex items-center gap-1 ${
                    chatPowerMode
                      ? 'bg-amber-600 text-white'
                      : 'bg-dark-700 text-dark-400 hover:text-white'
                  }`}
                  title={chatPowerMode ? 'Power Mode: Claude can modify the canvas' : 'User Mode: Read-only assistant'}
                >
                  <Zap size={12} />
                  {chatPowerMode ? 'Power' : 'User'}
                </button>
                <button
                  onClick={() => setShowChat(false)}
                  className="p-1 rounded hover:bg-dark-700 text-dark-400 hover:text-white transition-colors"
                >
                  <X size={16} />
                </button>
              </div>
            </div>
            {/* Chat Content */}
            <ChatPanel
              messages={messages}
              isThinking={isThinking}
              onSendMessage={sendMessage}
              isConnected={isConnected}
            />
          </div>
        )}
      </main>

      {/* Template Selector Modal */}
      <TemplateSelector
        isOpen={showTemplates}
        onClose={() => setShowTemplates(false)}
        onSelect={handleTemplateSelect}
      />

      {/* Config Importer Modal */}
      <ConfigImporter
        isOpen={showImporter}
        onClose={() => setShowImporter(false)}
        onImport={handleImport}
      />

      {/* Notification Toast */}
      {notification && (
        <div
          className="fixed bottom-4 left-1/2 transform -translate-x-1/2 px-4 py-2 bg-dark-800 text-white rounded-lg shadow-lg z-50 border border-dark-600"
          style={{ animation: 'slideUp 0.3s ease-out' }}
        >
          {notification}
        </div>
      )}

      <style>{`
        @keyframes slideUp {
          from { opacity: 0; transform: translate(-50%, 20px); }
          to { opacity: 1; transform: translate(-50%, 0); }
        }
      `}</style>
    </div>
  );
}

export default CanvasView;