Atomizer/atomizer-dashboard/frontend/src/components/canvas/panels/CodeEditorPanel.tsx

/**
 * CodeEditorPanel - Monaco editor for custom extractor Python code
 *
 * Features:
 * - Python syntax highlighting
 * - Auto-completion for common patterns
 * - Error display
 * - Claude AI code generation button
 * - Preview of extracted outputs
 */

import { useState, useCallback, useRef } from 'react';
import Editor, { OnMount, OnChange } from '@monaco-editor/react';
import {
  Play,
  Wand2,
  Copy,
  Check,
  AlertCircle,
  RefreshCw,
  X,
  ChevronDown,
  ChevronRight,
  FileCode,
  Sparkles,
} from 'lucide-react';

// Monaco editor types
type Monaco = Parameters<OnMount>[1];
type EditorInstance = Parameters<OnMount>[0];

interface CodeEditorPanelProps {
  /** Initial code content */
  initialCode?: string;
  /** Callback when code changes */
  onChange?: (code: string) => void;
  /** Callback when user requests Claude generation */
  onRequestGeneration?: (prompt: string) => Promise<string>;
  /** Whether the panel is read-only */
  readOnly?: boolean;
  /** Extractor name for context */
  extractorName?: string;
  /** Output variable names */
  outputs?: string[];
  /** Optional height (default: 300px) */
  height?: number | string;
  /** Show/hide header */
  showHeader?: boolean;
  /** Callback when running code (validation) */
  onRun?: (code: string) => Promise<{ success: boolean; error?: string; outputs?: Record<string, unknown> }>;
  /** Close button callback */
  onClose?: () => void;
}

// Default Python template for custom extractors
const DEFAULT_EXTRACTOR_TEMPLATE = `"""
Custom Extractor Function

This function is called after FEA simulation completes.
It receives the results and should return extracted values.

Available inputs:
- op2_path: Path to the .op2 results file
- fem_path: Path to the .fem file  
- params: Dict of current design variable values
- subcase_id: Current subcase being analyzed (optional)

Return a dict with your extracted values.
"""

from pyNastran.op2.op2 import OP2
import numpy as np

def extract(op2_path: str, fem_path: str, params: dict, subcase_id: int = 1) -> dict:
    """
    Extract physics from FEA results.
    
    Args:
        op2_path: Path to OP2 results file
        fem_path: Path to FEM file
        params: Current design variable values
        subcase_id: Subcase ID to analyze
        
    Returns:
        Dict with extracted values, e.g. {'max_stress': 150.5, 'mass': 2.3}
    """
    # Load OP2 results
    op2 = OP2()
    op2.read_op2(op2_path)
    
    # Example: Extract max displacement
    if subcase_id in op2.displacements:
        disp = op2.displacements[subcase_id]
        # Get magnitude of displacement vectors
        magnitudes = np.sqrt(np.sum(disp.data[0, :, 1:4]**2, axis=1))
        max_disp = float(np.max(magnitudes))
    else:
        max_disp = 0.0
    
    return {
        'max_displacement': max_disp,
        # Add more outputs as needed
    }
`;

export function CodeEditorPanel({
  initialCode = DEFAULT_EXTRACTOR_TEMPLATE,
  onChange,
  onRequestGeneration,
  readOnly = false,
  extractorName = 'custom_extractor',
  outputs = [],
  height = 400,
  showHeader = true,
  onRun,
  onClose,
}: CodeEditorPanelProps) {
  const [code, setCode] = useState(initialCode);
  const [isGenerating, setIsGenerating] = useState(false);
  const [isRunning, setIsRunning] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [runResult, setRunResult] = useState<{ success: boolean; outputs?: Record<string, unknown> } | null>(null);
  const [copied, setCopied] = useState(false);
  const [showPromptInput, setShowPromptInput] = useState(false);
  const [generationPrompt, setGenerationPrompt] = useState('');
  const [showOutputs, setShowOutputs] = useState(true);

  const editorRef = useRef<EditorInstance | null>(null);
  const monacoRef = useRef<Monaco | null>(null);

  // Handle editor mount
  const handleEditorMount: OnMount = (editor, monaco) => {
    editorRef.current = editor;
    monacoRef.current = monaco;

    // Configure Python language
    monaco.languages.registerCompletionItemProvider('python', {
      provideCompletionItems: (model: Parameters<typeof monaco.editor.createModel>[0], position: { lineNumber: number; column: number }) => {
        const word = model.getWordUntilPosition(position);
        const range = {
          startLineNumber: position.lineNumber,
          endLineNumber: position.lineNumber,
          startColumn: word.startColumn,
          endColumn: word.endColumn,
        };

        const suggestions = [
          {
            label: 'op2.read_op2',
            kind: monaco.languages.CompletionItemKind.Method,
            insertText: 'op2.read_op2(op2_path)',
            documentation: 'Read OP2 results file',
            range,
          },
          {
            label: 'op2.displacements',
            kind: monaco.languages.CompletionItemKind.Property,
            insertText: 'op2.displacements[subcase_id]',
            documentation: 'Access displacement results for a subcase',
            range,
          },
          {
            label: 'op2.eigenvectors',
            kind: monaco.languages.CompletionItemKind.Property,
            insertText: 'op2.eigenvectors[subcase_id]',
            documentation: 'Access eigenvector results for modal analysis',
            range,
          },
          {
            label: 'np.max',
            kind: monaco.languages.CompletionItemKind.Function,
            insertText: 'np.max(${1:array})',
            insertTextRules: monaco.languages.CompletionItemInsertTextRule.InsertAsSnippet,
            documentation: 'Get maximum value from array',
            range,
          },
          {
            label: 'np.sqrt',
            kind: monaco.languages.CompletionItemKind.Function,
            insertText: 'np.sqrt(${1:array})',
            insertTextRules: monaco.languages.CompletionItemInsertTextRule.InsertAsSnippet,
            documentation: 'Square root of array elements',
            range,
          },
          {
            label: 'extract_function',
            kind: monaco.languages.CompletionItemKind.Snippet,
            insertText: `def extract(op2_path: str, fem_path: str, params: dict, subcase_id: int = 1) -> dict:
    """Extract physics from FEA results."""
    op2 = OP2()
    op2.read_op2(op2_path)
    
    # Your extraction logic here
    
    return {
        '\${1:output_name}': \${2:value},
    }`,
            insertTextRules: monaco.languages.CompletionItemInsertTextRule.InsertAsSnippet,
            documentation: 'Insert a complete extract function template',
            range,
          },
        ];

        return { suggestions };
      },
    });
  };

  // Handle code change
  const handleCodeChange: OnChange = (value) => {
    const newCode = value || '';
    setCode(newCode);
    setError(null);
    setRunResult(null);
    onChange?.(newCode);
  };

  // Copy code to clipboard
  const handleCopy = useCallback(() => {
    navigator.clipboard.writeText(code);
    setCopied(true);
    setTimeout(() => setCopied(false), 2000);
  }, [code]);

  // Request Claude generation
  const handleGenerate = useCallback(async () => {
    if (!onRequestGeneration || !generationPrompt.trim()) return;

    setIsGenerating(true);
    setError(null);

    try {
      const generatedCode = await onRequestGeneration(generationPrompt);
      setCode(generatedCode);
      onChange?.(generatedCode);
      setShowPromptInput(false);
      setGenerationPrompt('');
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Generation failed');
    } finally {
      setIsGenerating(false);
    }
  }, [onRequestGeneration, generationPrompt, onChange]);

  // Run/validate code
  const handleRun = useCallback(async () => {
    if (!onRun) return;

    setIsRunning(true);
    setError(null);
    setRunResult(null);

    try {
      const result = await onRun(code);
      setRunResult(result);
      if (!result.success && result.error) {
        setError(result.error);
      }
    } catch (err) {
      setError(err instanceof Error ? err.message : 'Validation failed');
    } finally {
      setIsRunning(false);
    }
  }, [code, onRun]);

  return (
    <div className="flex flex-col h-full bg-dark-850 border-l border-dark-700">
      {/* Header */}
      {showHeader && (
        <div className="flex items-center justify-between px-4 py-2 border-b border-dark-700">
          <div className="flex items-center gap-2">
            <FileCode size={16} className="text-emerald-400" />
            <span className="font-medium text-white text-sm">{extractorName}</span>
            <span className="text-xs text-dark-500">.py</span>
          </div>
          <div className="flex items-center gap-2">
            {/* Claude Generate Button */}
            {onRequestGeneration && (
              <button
                onClick={() => setShowPromptInput(!showPromptInput)}
                className="p-1.5 rounded text-violet-400 hover:bg-violet-500/20 transition-colors"
                title="Generate with Claude"
              >
                <Sparkles size={16} />
              </button>
            )}

            {/* Copy Button */}
            <button
              onClick={handleCopy}
              className="p-1.5 rounded text-dark-400 hover:text-white hover:bg-dark-700 transition-colors"
              title="Copy code"
            >
              {copied ? <Check size={16} className="text-emerald-400" /> : <Copy size={16} />}
            </button>

            {/* Run Button */}
            {onRun && (
              <button
                onClick={handleRun}
                disabled={isRunning}
                className="p-1.5 rounded text-emerald-400 hover:bg-emerald-500/20 transition-colors disabled:opacity-50"
                title="Validate code"
              >
                {isRunning ? (
                  <RefreshCw size={16} className="animate-spin" />
                ) : (
                  <Play size={16} />
                )}
              </button>
            )}

            {/* Close Button */}
            {onClose && (
              <button
                onClick={onClose}
                className="p-1.5 rounded text-dark-400 hover:text-white hover:bg-dark-700 transition-colors"
              >
                <X size={16} />
              </button>
            )}
          </div>
        </div>
      )}

      {/* Claude Prompt Input */}
      {showPromptInput && (
        <div className="px-4 py-3 border-b border-dark-700 bg-violet-500/5">
          <div className="flex items-center gap-2 mb-2">
            <Wand2 size={14} className="text-violet-400" />
            <span className="text-xs text-violet-400 font-medium">Generate with Claude</span>
          </div>
          <textarea
            value={generationPrompt}
            onChange={(e) => setGenerationPrompt(e.target.value)}
            placeholder="Describe what you want to extract... e.g., 'Extract maximum von Mises stress and total mass from the model'"
            className="w-full px-3 py-2 bg-dark-800 border border-dark-600 rounded-lg text-sm text-white placeholder-dark-500 resize-none focus:outline-none focus:border-violet-500"
            rows={2}
          />
          <div className="flex justify-end gap-2 mt-2">
            <button
              onClick={() => setShowPromptInput(false)}
              className="px-3 py-1.5 text-xs text-dark-400 hover:text-white transition-colors"
            >
              Cancel
            </button>
            <button
              onClick={handleGenerate}
              disabled={isGenerating || !generationPrompt.trim()}
              className="px-3 py-1.5 text-xs bg-violet-600 text-white rounded hover:bg-violet-500 disabled:opacity-50 disabled:cursor-not-allowed transition-colors flex items-center gap-1.5"
            >
              {isGenerating ? (
                <>
                  <RefreshCw size={12} className="animate-spin" />
                  Generating...
                </>
              ) : (
                <>
                  <Sparkles size={12} />
                  Generate
                </>
              )}
            </button>
          </div>
        </div>
      )}

      {/* Error Display */}
      {error && (
        <div className="px-4 py-2 bg-red-500/10 border-b border-red-500/30 flex items-center gap-2">
          <AlertCircle size={14} className="text-red-400 flex-shrink-0" />
          <span className="text-xs text-red-400 font-mono">{error}</span>
        </div>
      )}

      {/* Monaco Editor */}
      <div className="flex-1 min-h-0">
        <Editor
          height={height}
          language="python"
          theme="vs-dark"
          value={code}
          onChange={handleCodeChange}
          onMount={handleEditorMount}
          options={{
            readOnly,
            minimap: { enabled: false },
            fontSize: 13,
            lineNumbers: 'on',
            scrollBeyondLastLine: false,
            wordWrap: 'on',
            automaticLayout: true,
            tabSize: 4,
            insertSpaces: true,
            padding: { top: 8, bottom: 8 },
            scrollbar: {
              vertical: 'auto',
              horizontal: 'auto',
            },
          }}
        />
      </div>

      {/* Outputs Preview */}
      {(outputs.length > 0 || runResult?.outputs) && (
        <div className="border-t border-dark-700">
          <button
            onClick={() => setShowOutputs(!showOutputs)}
            className="w-full px-4 py-2 flex items-center gap-2 text-xs text-dark-400 hover:text-white hover:bg-dark-800/50 transition-colors"
          >
            {showOutputs ? <ChevronDown size={12} /> : <ChevronRight size={12} />}
            <span>Outputs</span>
            <span className="ml-auto text-dark-500">
              {runResult?.outputs
                ? Object.keys(runResult.outputs).length
                : outputs.length}
            </span>
          </button>
          {showOutputs && (
            <div className="px-4 pb-3 space-y-1">
              {runResult?.outputs ? (
                Object.entries(runResult.outputs).map(([key, value]) => (
                  <div
                    key={key}
                    className="flex items-center justify-between px-2 py-1 bg-dark-800 rounded text-xs"
                  >
                    <span className="text-emerald-400 font-mono">{key}</span>
                    <span className="text-dark-300">{String(value)}</span>
                  </div>
                ))
              ) : (
                outputs.map((output) => (
                  <div
                    key={output}
                    className="flex items-center px-2 py-1 bg-dark-800 rounded text-xs"
                  >
                    <span className="text-dark-400 font-mono">{output}</span>
                  </div>
                ))
              )}
            </div>
          )}
        </div>
      )}
    </div>
  );
}

export default CodeEditorPanel;