spicebook/frontend/src/lib/notebook-store.ts

import { create } from 'zustand';
import type {
  Notebook,
  Cell,
  CellType,
  CellOutput,
  SimulationResponse,
} from './types';
import * as api from './api';

function generateId(): string {
  return crypto.randomUUID().slice(0, 8);
}

function makeCell(type: CellType): Cell {
  const templates: Record<CellType, string> = {
    markdown: '## New Section\n\nDescribe your circuit here.',
    spice: '* SPICE Netlist\n\nR1 in out 1k\nV1 in 0 DC 5\n\n.op\n.end',
    python: '# Python analysis\nimport numpy as np\n',
    schematic: '',
  };

  return {
    id: generateId(),
    type,
    source: templates[type],
    outputs: [],
  };
}

interface NotebookStore {
  // State
  notebook: Notebook | null;
  notebookId: string | null;
  activeCell: string | null;
  runningCells: Set<string>;
  dirty: boolean;
  loading: boolean;
  error: string | null;
  saving: boolean;

  // Actions
  loadNotebook: (id: string) => Promise<void>;
  saveNotebook: () => Promise<void>;
  updateCellSource: (cellId: string, source: string) => void;
  addCell: (type: CellType, afterId?: string) => void;
  deleteCell: (cellId: string) => void;
  moveCell: (cellId: string, direction: 'up' | 'down') => void;
  setActiveCell: (cellId: string | null) => void;
  runCell: (cellId: string) => Promise<void>;
  runAllCells: () => Promise<void>;
  setCellOutput: (cellId: string, output: CellOutput) => void;
  updateTitle: (title: string) => void;
  updateEngine: (engine: string) => void;
  clearError: () => void;
}

export const useNotebookStore = create<NotebookStore>((set, get) => ({
  notebook: null,
  notebookId: null,
  activeCell: null,
  runningCells: new Set(),
  dirty: false,
  loading: false,
  error: null,
  saving: false,

  loadNotebook: async (id: string) => {
    set({ loading: true, error: null });
    try {
      const notebook = await api.getNotebook(id);
      set({
        notebook,
        notebookId: id,
        loading: false,
        dirty: false,
        activeCell: notebook.cells.length > 0 ? notebook.cells[0].id : null,
      });
    } catch (err) {
      set({
        loading: false,
        error: err instanceof Error ? err.message : 'Failed to load notebook',
      });
    }
  },

  saveNotebook: async () => {
    const { notebook, notebookId } = get();
    if (!notebook || !notebookId) return;

    set({ saving: true });
    try {
      const updated = await api.updateNotebook(notebookId, notebook);
      set({ notebook: updated, dirty: false, saving: false });
    } catch (err) {
      set({
        saving: false,
        error: err instanceof Error ? err.message : 'Failed to save notebook',
      });
    }
  },

  updateCellSource: (cellId: string, source: string) => {
    const { notebook } = get();
    if (!notebook) return;

    set({
      notebook: {
        ...notebook,
        cells: notebook.cells.map((cell) =>
          cell.id === cellId ? { ...cell, source } : cell,
        ),
      },
      dirty: true,
    });
  },

  addCell: (type: CellType, afterId?: string) => {
    const { notebook } = get();
    if (!notebook) return;

    const newCell = makeCell(type);
    let cells: Cell[];

    if (afterId) {
      const idx = notebook.cells.findIndex((c) => c.id === afterId);
      cells = [
        ...notebook.cells.slice(0, idx + 1),
        newCell,
        ...notebook.cells.slice(idx + 1),
      ];
    } else {
      cells = [...notebook.cells, newCell];
    }

    set({
      notebook: { ...notebook, cells },
      activeCell: newCell.id,
      dirty: true,
    });
  },

  deleteCell: (cellId: string) => {
    const { notebook, activeCell } = get();
    if (!notebook) return;

    const cells = notebook.cells.filter((c) => c.id !== cellId);
    const newActive =
      activeCell === cellId
        ? cells.length > 0
          ? cells[0].id
          : null
        : activeCell;

    set({
      notebook: { ...notebook, cells },
      activeCell: newActive,
      dirty: true,
    });
  },

  moveCell: (cellId: string, direction: 'up' | 'down') => {
    const { notebook } = get();
    if (!notebook) return;

    const idx = notebook.cells.findIndex((c) => c.id === cellId);
    if (idx === -1) return;

    const targetIdx = direction === 'up' ? idx - 1 : idx + 1;
    if (targetIdx < 0 || targetIdx >= notebook.cells.length) return;

    const cells = [...notebook.cells];
    [cells[idx], cells[targetIdx]] = [cells[targetIdx], cells[idx]];

    set({
      notebook: { ...notebook, cells },
      dirty: true,
    });
  },

  setActiveCell: (cellId: string | null) => {
    set({ activeCell: cellId });
  },

  runCell: async (cellId: string) => {
    const { notebook, notebookId, runningCells } = get();
    if (!notebook || !notebookId) return;

    const cell = notebook.cells.find((c) => c.id === cellId);
    if (!cell || cell.type !== 'spice') return;

    const newRunning = new Set(runningCells);
    newRunning.add(cellId);
    set({ runningCells: newRunning });

    try {
      // Try the cell-specific endpoint first, fall back to raw simulation
      let result: SimulationResponse;
      try {
        result = await api.runCell(notebookId, cellId);
      } catch {
        result = await api.runSimulation(
          cell.source,
          notebook.metadata.engine,
        );
      }

      const output: CellOutput = {
        output_type: result.success ? 'simulation_result' : 'error',
        data: {
          success: result.success,
          waveform: result.waveform || null,
          log: result.log,
          error: result.error || null,
          elapsed_seconds: result.elapsed_seconds,
        },
        timestamp: new Date().toISOString(),
      };

      const updatedNotebook = get().notebook;
      if (updatedNotebook) {
        set({
          notebook: {
            ...updatedNotebook,
            cells: updatedNotebook.cells.map((c) =>
              c.id === cellId ? { ...c, outputs: [output] } : c,
            ),
          },
        });
      }
    } catch (err) {
      const output: CellOutput = {
        output_type: 'error',
        data: {
          success: false,
          log: '',
          error: err instanceof Error ? err.message : 'Simulation failed',
          elapsed_seconds: 0,
        },
        timestamp: new Date().toISOString(),
      };

      const updatedNotebook = get().notebook;
      if (updatedNotebook) {
        set({
          notebook: {
            ...updatedNotebook,
            cells: updatedNotebook.cells.map((c) =>
              c.id === cellId ? { ...c, outputs: [output] } : c,
            ),
          },
        });
      }
    } finally {
      const currentRunning = new Set(get().runningCells);
      currentRunning.delete(cellId);
      set({ runningCells: currentRunning });
    }
  },

  runAllCells: async () => {
    const { notebook } = get();
    if (!notebook) return;

    const spiceCells = notebook.cells.filter((c) => c.type === 'spice');
    for (const cell of spiceCells) {
      await get().runCell(cell.id);
    }
  },

  setCellOutput: (cellId: string, output: CellOutput) => {
    const { notebook } = get();
    if (!notebook) return;

    set({
      notebook: {
        ...notebook,
        cells: notebook.cells.map((c) =>
          c.id === cellId ? { ...c, outputs: [...c.outputs, output] } : c,
        ),
      },
    });
  },

  updateTitle: (title: string) => {
    const { notebook } = get();
    if (!notebook) return;

    set({
      notebook: {
        ...notebook,
        metadata: { ...notebook.metadata, title },
      },
      dirty: true,
    });
  },

  updateEngine: (engine: string) => {
    const { notebook } = get();
    if (!notebook) return;

    set({
      notebook: {
        ...notebook,
        metadata: { ...notebook.metadata, engine },
      },
      dirty: true,
    });
  },

  clearError: () => {
    set({ error: null });
  },
}));