Atomizer/atomizer-dashboard/frontend/src/hooks/useCanvasStore.ts

import { create } from 'zustand';
import { Node, Edge, addEdge, applyNodeChanges, applyEdgeChanges, Connection, NodeChange, EdgeChange } from 'reactflow';
import { CanvasNodeData, NodeType } from '../lib/canvas/schema';
import { validateGraph, ValidationResult } from '../lib/canvas/validation';
import { serializeToIntent, OptimizationIntent } from '../lib/canvas/intent';

interface CanvasState {
  nodes: Node<CanvasNodeData>[];
  edges: Edge[];
  selectedNode: string | null;
  validation: ValidationResult;

  // Actions
  onNodesChange: (changes: NodeChange[]) => void;
  onEdgesChange: (changes: EdgeChange[]) => void;
  onConnect: (connection: Connection) => void;
  addNode: (type: NodeType, position: { x: number; y: number }) => void;
  updateNodeData: (nodeId: string, data: Partial<CanvasNodeData>) => void;
  selectNode: (nodeId: string | null) => void;
  deleteSelected: () => void;
  validate: () => ValidationResult;
  toIntent: () => OptimizationIntent;
  clear: () => void;
  loadFromIntent: (intent: OptimizationIntent) => void;
}

let nodeIdCounter = 0;
const getNodeId = () => `node_${++nodeIdCounter}`;

const getDefaultData = (type: NodeType): CanvasNodeData => {
  const base = { label: type.charAt(0).toUpperCase() + type.slice(1), configured: false };

  switch (type) {
    case 'model': return { ...base, type: 'model' };
    case 'solver': return { ...base, type: 'solver' };
    case 'designVar': return { ...base, type: 'designVar', label: 'Design Variable' };
    case 'extractor': return { ...base, type: 'extractor' };
    case 'objective': return { ...base, type: 'objective' };
    case 'constraint': return { ...base, type: 'constraint' };
    case 'algorithm': return { ...base, type: 'algorithm' };
    case 'surrogate': return { ...base, type: 'surrogate', enabled: false };
    default: return { ...base, type } as CanvasNodeData;
  }
};

export const useCanvasStore = create<CanvasState>((set, get) => ({
  nodes: [],
  edges: [],
  selectedNode: null,
  validation: { valid: false, errors: [], warnings: [] },

  onNodesChange: (changes) => {
    set({ nodes: applyNodeChanges(changes, get().nodes) });
  },

  onEdgesChange: (changes) => {
    set({ edges: applyEdgeChanges(changes, get().edges) });
  },

  onConnect: (connection) => {
    set({ edges: addEdge(connection, get().edges) });
  },

  addNode: (type, position) => {
    const newNode: Node<CanvasNodeData> = {
      id: getNodeId(),
      type,
      position,
      data: getDefaultData(type),
    };
    set({ nodes: [...get().nodes, newNode] });
  },

  updateNodeData: (nodeId, data) => {
    set({
      nodes: get().nodes.map((node) =>
        node.id === nodeId
          ? { ...node, data: { ...node.data, ...data } }
          : node
      ),
    });
  },

  selectNode: (nodeId) => {
    set({ selectedNode: nodeId });
  },

  deleteSelected: () => {
    const { selectedNode, nodes, edges } = get();
    if (!selectedNode) return;

    set({
      nodes: nodes.filter((n) => n.id !== selectedNode),
      edges: edges.filter((e) => e.source !== selectedNode && e.target !== selectedNode),
      selectedNode: null,
    });
  },

  validate: () => {
    const { nodes, edges } = get();
    const result = validateGraph(nodes, edges);
    set({ validation: result });
    return result;
  },

  toIntent: () => {
    const { nodes, edges } = get();
    return serializeToIntent(nodes, edges);
  },

  clear: () => {
    set({
      nodes: [],
      edges: [],
      selectedNode: null,
      validation: { valid: false, errors: [], warnings: [] },
    });
    nodeIdCounter = 0;
  },

  loadFromIntent: (intent) => {
    // TODO: Implement reverse serialization
    console.log('Loading intent:', intent);
  },
}));