Atomizer/atomizer-dashboard/frontend/src/components/canvas/SpecRenderer.tsx

/**
 * SpecRenderer - ReactFlow canvas that renders from AtomizerSpec v2.0
 *
 * This component replaces the legacy canvas approach with a spec-driven architecture:
 * - Reads from useSpecStore instead of useCanvasStore
 * - Converts spec to ReactFlow nodes/edges using spec converters
 * - All changes flow through the spec store and sync with backend
 * - Supports WebSocket real-time updates
 *
 * P2.7-P2.10: SpecRenderer component with node/edge/selection handling
 */

import { useCallback, useRef, useEffect, useMemo, useState, DragEvent } from 'react';
import ReactFlow, {
  Background,
  Controls,
  MiniMap,
  ReactFlowProvider,
  ReactFlowInstance,
  Edge,
  Node,
  NodeChange,
  EdgeChange,
  Connection,
  applyNodeChanges,
} from 'reactflow';
import 'reactflow/dist/style.css';

import { nodeTypes } from './nodes';
import { specToNodes, specToEdges } from '../../lib/spec';
import {
  useSpecStore,
  useSpec,
  useSpecLoading,
  useSpecError,
  useSelectedNodeId,
  useSelectedEdgeId,
} from '../../hooks/useSpecStore';
import { useSpecWebSocket } from '../../hooks/useSpecWebSocket';
import { ConnectionStatusIndicator } from './ConnectionStatusIndicator';
import { CanvasNodeData } from '../../lib/canvas/schema';

// ============================================================================
// Drag-Drop Helpers
// ============================================================================

/** Addable node types via drag-drop */
const ADDABLE_NODE_TYPES = ['designVar', 'extractor', 'objective', 'constraint'] as const;
type AddableNodeType = typeof ADDABLE_NODE_TYPES[number];

function isAddableNodeType(type: string): type is AddableNodeType {
  return ADDABLE_NODE_TYPES.includes(type as AddableNodeType);
}

/** Maps canvas NodeType to spec API type */
function mapNodeTypeToSpecType(type: AddableNodeType): 'designVar' | 'extractor' | 'objective' | 'constraint' {
  return type;
}

/** Creates default data for a new node of the given type */
function getDefaultNodeData(type: AddableNodeType, position: { x: number; y: number }): Record<string, unknown> {
  const timestamp = Date.now();
  
  switch (type) {
    case 'designVar':
      return {
        name: `variable_${timestamp}`,
        expression_name: `expr_${timestamp}`,
        type: 'continuous',
        bounds: { min: 0, max: 1 },
        baseline: 0.5,
        enabled: true,
        canvas_position: position,
      };
    case 'extractor':
      return {
        name: `extractor_${timestamp}`,
        type: 'custom_function',  // Must be valid ExtractorType
        builtin: false,
        enabled: true,
        // Custom function extractors need a function definition
        function: {
          name: 'extract',
          source_code: `def extract(op2_path: str, config: dict = None) -> dict:
    """
    Custom extractor function.
    
    Args:
        op2_path: Path to the OP2 results file
        config: Optional configuration dict
    
    Returns:
        Dictionary with extracted values
    """
    # TODO: Implement extraction logic
    return {'value': 0.0}
`,
        },
        outputs: [{ name: 'value', metric: 'custom' }],
        canvas_position: position,
      };
    case 'objective':
      return {
        name: `objective_${timestamp}`,
        direction: 'minimize',
        weight: 1.0,
        source_extractor_id: null,
        source_output: null,
        canvas_position: position,
      };
    case 'constraint':
      return {
        name: `constraint_${timestamp}`,
        constraint_type: 'hard',  // Must be 'hard' or 'soft'
        operator: '<=',
        limit: 1.0,
        source_extractor_id: null,
        source_output: null,
        enabled: true,
        canvas_position: position,
      };
  }
}

// ============================================================================
// Component Props
// ============================================================================

interface SpecRendererProps {
  /**
   * Optional study ID to load on mount.
   * If not provided, assumes spec is already loaded in the store.
   */
  studyId?: string;

  /**
   * Callback when study changes (for URL updates)
   */
  onStudyChange?: (studyId: string) => void;

  /**
   * Show loading overlay while spec is loading
   */
  showLoadingOverlay?: boolean;

  /**
   * Enable/disable editing (drag, connect, delete)
   */
  editable?: boolean;

  /**
   * Enable real-time WebSocket sync (default: true)
   */
  enableWebSocket?: boolean;

  /**
   * Show connection status indicator (default: true when WebSocket enabled)
   */
  showConnectionStatus?: boolean;
}

function SpecRendererInner({
  studyId,
  onStudyChange,
  showLoadingOverlay = true,
  editable = true,
  enableWebSocket = true,
  showConnectionStatus = true,
}: SpecRendererProps) {
  const reactFlowWrapper = useRef<HTMLDivElement>(null);
  const reactFlowInstance = useRef<ReactFlowInstance | null>(null);

  // Spec store state and actions
  const spec = useSpec();
  const isLoading = useSpecLoading();
  const error = useSpecError();
  const selectedNodeId = useSelectedNodeId();
  const selectedEdgeId = useSelectedEdgeId();

  const {
    loadSpec,
    selectNode,
    selectEdge,
    clearSelection,
    updateNodePosition,
    addNode,
    addEdge,
    removeEdge,
    removeNode,
    setError,
  } = useSpecStore();

  // WebSocket for real-time sync
  const storeStudyId = useSpecStore((s) => s.studyId);
  const wsStudyId = enableWebSocket ? storeStudyId : null;
  const { status: wsStatus } = useSpecWebSocket(wsStudyId);

  // Load spec on mount if studyId provided
  useEffect(() => {
    if (studyId) {
      loadSpec(studyId).then(() => {
        if (onStudyChange) {
          onStudyChange(studyId);
        }
      });
    }
  }, [studyId, loadSpec, onStudyChange]);

  // Convert spec to ReactFlow nodes
  const nodes = useMemo(() => {
    return specToNodes(spec);
  }, [spec]);

  // Convert spec to ReactFlow edges with selection styling
  const edges = useMemo(() => {
    const baseEdges = specToEdges(spec);
    return baseEdges.map((edge) => ({
      ...edge,
      style: {
        stroke: edge.id === selectedEdgeId ? '#60a5fa' : '#6b7280',
        strokeWidth: edge.id === selectedEdgeId ? 3 : 2,
      },
      animated: edge.id === selectedEdgeId,
    }));
  }, [spec, selectedEdgeId]);

  // Track node positions for change handling
  const nodesRef = useRef<Node<CanvasNodeData>[]>(nodes);
  useEffect(() => {
    nodesRef.current = nodes;
  }, [nodes]);

  // Track local node state for smooth dragging
  const [localNodes, setLocalNodes] = useState(nodes);
  
  // Sync local nodes with spec-derived nodes when spec changes
  useEffect(() => {
    setLocalNodes(nodes);
  }, [nodes]);

  // Handle node position changes
  const onNodesChange = useCallback(
    (changes: NodeChange[]) => {
      if (!editable) return;

      // Apply changes to local state for smooth dragging
      setLocalNodes((nds) => applyNodeChanges(changes, nds));

      // Handle position changes - save to spec when drag ends
      for (const change of changes) {
        if (change.type === 'position' && change.position && change.dragging === false) {
          // Dragging ended - update spec
          updateNodePosition(change.id, {
            x: change.position.x,
            y: change.position.y,
          });
        }
      }
    },
    [editable, updateNodePosition]
  );

  // Handle edge changes (deletion)
  const onEdgesChange = useCallback(
    (changes: EdgeChange[]) => {
      if (!editable) return;

      for (const change of changes) {
        if (change.type === 'remove') {
          // Find the edge being removed
          const edge = edges.find((e) => e.id === change.id);
          if (edge) {
            removeEdge(edge.source, edge.target).catch((err) => {
              console.error('Failed to remove edge:', err);
              setError(err.message);
            });
          }
        }
      }
    },
    [editable, edges, removeEdge, setError]
  );

  // Handle new connections
  const onConnect = useCallback(
    (connection: Connection) => {
      if (!editable) return;
      if (!connection.source || !connection.target) return;

      addEdge(connection.source, connection.target).catch((err) => {
        console.error('Failed to add edge:', err);
        setError(err.message);
      });
    },
    [editable, addEdge, setError]
  );

  // Handle node clicks for selection
  const onNodeClick = useCallback(
    (_: React.MouseEvent, node: { id: string }) => {
      selectNode(node.id);
    },
    [selectNode]
  );

  // Handle edge clicks for selection
  const onEdgeClick = useCallback(
    (_: React.MouseEvent, edge: Edge) => {
      selectEdge(edge.id);
    },
    [selectEdge]
  );

  // Handle pane clicks to clear selection
  const onPaneClick = useCallback(() => {
    clearSelection();
  }, [clearSelection]);

  // Keyboard handler for Delete/Backspace
  useEffect(() => {
    if (!editable) return;

    const handleKeyDown = (event: KeyboardEvent) => {
      if (event.key === 'Delete' || event.key === 'Backspace') {
        const target = event.target as HTMLElement;
        if (
          target.tagName === 'INPUT' ||
          target.tagName === 'TEXTAREA' ||
          target.isContentEditable
        ) {
          return;
        }

        // Delete selected edge first
        if (selectedEdgeId) {
          const edge = edges.find((e) => e.id === selectedEdgeId);
          if (edge) {
            removeEdge(edge.source, edge.target).catch((err) => {
              console.error('Failed to delete edge:', err);
              setError(err.message);
            });
          }
          return;
        }

        // Delete selected node
        if (selectedNodeId) {
          // Don't allow deleting synthetic nodes (model, solver, optimization)
          if (['model', 'solver', 'optimization', 'surrogate'].includes(selectedNodeId)) {
            return;
          }

          removeNode(selectedNodeId).catch((err) => {
            console.error('Failed to delete node:', err);
            setError(err.message);
          });
        }
      }
    };

    window.addEventListener('keydown', handleKeyDown);
    return () => window.removeEventListener('keydown', handleKeyDown);
  }, [editable, selectedNodeId, selectedEdgeId, edges, removeNode, removeEdge, setError]);

  // =========================================================================
  // Drag-Drop Handlers
  // =========================================================================

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

  const onDrop = useCallback(
    async (event: DragEvent) => {
      if (!editable || !reactFlowInstance.current) return;
      event.preventDefault();

      const type = event.dataTransfer.getData('application/reactflow');
      if (!type || !isAddableNodeType(type)) {
        console.warn('Invalid or non-addable node type dropped:', type);
        return;
      }

      // Convert screen position to flow position
      const position = reactFlowInstance.current.screenToFlowPosition({
        x: event.clientX,
        y: event.clientY,
      });

      // Create default data for the node
      const nodeData = getDefaultNodeData(type, position);
      const specType = mapNodeTypeToSpecType(type);

      try {
        const nodeId = await addNode(specType, nodeData);
        // Select the newly created node
        selectNode(nodeId);
      } catch (err) {
        console.error('Failed to add node:', err);
        setError(err instanceof Error ? err.message : 'Failed to add node');
      }
    },
    [editable, addNode, selectNode, setError]
  );

  // Loading state
  if (showLoadingOverlay && isLoading && !spec) {
    return (
      <div className="flex-1 flex items-center justify-center bg-dark-900">
        <div className="text-center">
          <div className="animate-spin w-8 h-8 border-2 border-primary-500 border-t-transparent rounded-full mx-auto mb-4" />
          <p className="text-dark-400">Loading spec...</p>
        </div>
      </div>
    );
  }

  // Error state
  if (error && !spec) {
    return (
      <div className="flex-1 flex items-center justify-center bg-dark-900">
        <div className="text-center max-w-md">
          <div className="w-12 h-12 rounded-full bg-red-900/50 flex items-center justify-center mx-auto mb-4">
            <svg className="w-6 h-6 text-red-400" fill="none" viewBox="0 0 24 24" stroke="currentColor">
              <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M12 8v4m0 4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z" />
            </svg>
          </div>
          <h3 className="text-lg font-medium text-white mb-2">Failed to load spec</h3>
          <p className="text-dark-400 mb-4">{error}</p>
          {studyId && (
            <button
              onClick={() => loadSpec(studyId)}
              className="px-4 py-2 bg-primary-600 text-white rounded-lg hover:bg-primary-500 transition-colors"
            >
              Retry
            </button>
          )}
        </div>
      </div>
    );
  }

  // Empty state
  if (!spec) {
    return (
      <div className="flex-1 flex items-center justify-center bg-dark-900">
        <div className="text-center">
          <p className="text-dark-400">No spec loaded</p>
          <p className="text-dark-500 text-sm mt-2">Load a study to see its optimization configuration</p>
        </div>
      </div>
    );
  }

  return (
    <div className="w-full h-full relative" ref={reactFlowWrapper}>
      {/* Status indicators (overlay) */}
      <div className="absolute top-4 right-4 z-20 flex items-center gap-2">
        {/* WebSocket connection status */}
        {enableWebSocket && showConnectionStatus && (
          <div className="px-3 py-1.5 bg-dark-800/90 backdrop-blur rounded-lg border border-dark-600">
            <ConnectionStatusIndicator status={wsStatus} />
          </div>
        )}

        {/* Loading indicator */}
        {isLoading && (
          <div className="flex items-center gap-2 px-3 py-1.5 bg-dark-800/90 backdrop-blur rounded-lg border border-dark-600">
            <div className="animate-spin w-4 h-4 border-2 border-primary-500 border-t-transparent rounded-full" />
            <span className="text-xs text-dark-300">Syncing...</span>
          </div>
        )}
      </div>

      {/* Error banner (overlay) */}
      {error && (
        <div className="absolute top-4 left-4 right-20 z-20 px-4 py-2 bg-red-900/90 backdrop-blur border border-red-700 text-red-200 rounded-lg text-sm flex justify-between items-center">
          <span>{error}</span>
          <button
            onClick={() => setError(null)}
            className="text-red-400 hover:text-red-200 ml-2"
          >
            ×
          </button>
        </div>
      )}

      <ReactFlow
        nodes={localNodes}
        edges={edges}
        onNodesChange={onNodesChange}
        onEdgesChange={onEdgesChange}
        onConnect={onConnect}
        onInit={(instance) => {
          reactFlowInstance.current = instance;
        }}
        onDragOver={onDragOver}
        onDrop={onDrop}
        onNodeClick={onNodeClick}
        onEdgeClick={onEdgeClick}
        onPaneClick={onPaneClick}
        nodeTypes={nodeTypes}
        fitView
        deleteKeyCode={null} // We handle delete ourselves
        nodesDraggable={editable}
        nodesConnectable={editable}
        elementsSelectable={true}
        className="bg-dark-900"
      >
        <Background color="#374151" gap={20} />
        <Controls className="!bg-dark-800 !border-dark-600 !rounded-lg [&>button]:!bg-dark-700 [&>button]:!border-dark-600 [&>button]:!fill-dark-300 [&>button:hover]:!bg-dark-600" />
        <MiniMap
          className="!bg-dark-800 !border-dark-600 !rounded-lg"
          nodeColor="#4B5563"
          maskColor="rgba(0, 0, 0, 0.5)"
        />
      </ReactFlow>

      {/* Study name badge */}
      <div className="absolute bottom-4 left-4 z-10 px-3 py-1.5 bg-dark-800/90 backdrop-blur rounded-lg border border-dark-600">
        <span className="text-sm text-dark-300">{spec.meta.study_name}</span>
      </div>
    </div>
  );
}

/**
 * SpecRenderer with ReactFlowProvider wrapper.
 *
 * Usage:
 * ```tsx
 * // Load spec on mount
 * <SpecRenderer studyId="M1_Mirror/m1_mirror_flatback" />
 *
 * // Use with already-loaded spec
 * const { loadSpec } = useSpecStore();
 * await loadSpec('M1_Mirror/m1_mirror_flatback');
 * <SpecRenderer />
 * ```
 */
export function SpecRenderer(props: SpecRendererProps) {
  return (
    <ReactFlowProvider>
      <SpecRendererInner {...props} />
    </ReactFlowProvider>
  );
}

export default SpecRenderer;