/* Copyright 2026 Marimo. All rights reserved. */

import { historyField } from "@codemirror/commands";
import { dequal as isEqual } from "dequal";
import { type Atom, atom, useAtom, useAtomValue } from "jotai";
import { atomFamily, selectAtom, splitAtom } from "jotai/utils";
import { createRef, type ReducerWithoutAction } from "react";
import type { CellHandle } from "@/components/editor/notebook-cell";
import type { CollapsibleTree } from "@/utils/id-tree";
import {
  type CellColumnId,
  type CellIndex,
  MultiColumn,
} from "@/utils/id-tree";
import { invariant } from "@/utils/invariant";
import { Logger } from "@/utils/Logger";
import { clamp } from "@/utils/math";
import { Objects } from "@/utils/objects";
import { extractAllTracebackInfo, type TracebackInfo } from "@/utils/traceback";
import { createReducerAndAtoms } from "../../utils/createReducer";
import { foldAllBulk, unfoldAllBulk } from "../codemirror/editing/commands";
import {
  splitEditor,
  updateEditorCodeFromPython,
} from "../codemirror/language/utils";
import { findCollapseRange, mergeOutlines } from "../dom/outline";
import type { CellMessage } from "../kernel/messages";
import { isErrorMime } from "../mime";
import type { CellConfig } from "../network/types";
import { isRtcEnabled } from "../rtc/state";
import { createDeepEqualAtom, store } from "../state/jotai";
import { prepareCellForExecution, transitionCell } from "./cell";
import { documentTransactionMiddleware } from "./document-changes";
import { CellId, SCRATCH_CELL_ID, SETUP_CELL_ID } from "./ids";
import { type CellLog, getCellLogsForMessage } from "./logs";
import {
  focusAndScrollCellIntoView,
  scrollToBottom,
  scrollToTop,
} from "./scrollCellIntoView";
import {
  type CellData,
  type CellRuntimeState,
  createCell,
  createCellConfig,
  createCellRuntimeState,
} from "./types";
import {
  canUndoDeletes,
  disabledCellIds,
  enabledCellIds,
  notebookIsRunning,
  notebookNeedsRun,
  notebookQueueOrRunningCount,
} from "./utils";

/**
 * The state of the notebook.
 */
export interface NotebookState {
  /**
   * Order of cells on the page.
   */
  cellIds: MultiColumn<CellId>;
  /**
   * Map of cells to their view state
   */
  cellData: Record<CellId, CellData>;
  /**
   * Map of cells to their runtime state
   */
  cellRuntime: Record<CellId, CellRuntimeState>;
  /**
   * Cell handlers
   */
  cellHandles: Record<CellId, React.RefObject<CellHandle | null>>;
  /**
   * Array of deleted cells (with their data and index) so that cell deletion can be undone
   *
   * (CodeMirror types the serialized config as any.)
   */
  history: {
    name: string;
    // oxlint-disable-next-line typescript/no-explicit-any
    serializedEditorState: any;
    column: CellColumnId;
    index: CellIndex;
    isSetupCell: boolean;
    config: CellConfig;
  }[];
  /**
   * Key of cell to scroll to; typically set by actions that re-order the cell
   * array. Call the SCROLL_TO_TARGET action to scroll to the specified cell
   * and clear this field.
   */
  scrollKey: CellId | null;
  /**
   * Logs of all cell messages
   */
  cellLogs: CellLog[];
  /**
   * Set of cells that have been created that are initialized with `hide_code`.
   *
   * These start temporarily open until their first blur event.
   */
  untouchedNewCells: Set<CellId>;
}

function withScratchCell(notebookState: NotebookState): NotebookState {
  const config = { column: 0, hide_code: false, disabled: false };
  return {
    ...notebookState,
    cellData: {
      [SCRATCH_CELL_ID]: createCell({ id: SCRATCH_CELL_ID, config: config }),
      ...notebookState.cellData,
    },
    cellRuntime: {
      [SCRATCH_CELL_ID]: createCellRuntimeState(),
      ...notebookState.cellRuntime,
    },
    cellHandles: {
      [SCRATCH_CELL_ID]: createRef(),
      ...notebookState.cellHandles,
    },
  };
}

/**
 * Initial state of the notebook.
 */
export function initialNotebookState(): NotebookState {
  return withScratchCell({
    cellIds: MultiColumn.from([]),
    cellData: {},
    cellRuntime: {},
    cellHandles: {},
    history: [],
    scrollKey: null,
    cellLogs: [],
    untouchedNewCells: new Set<CellId>(),
  });
}

/** The target cell ID to create a new cell relative to. Can be:
 * - A CellId string for an existing cell
 * - "__end__" to append at the end of the first column
 * - {type: "__end__", columnId} to append at the end of a specific column
 */
export type CellPosition =
  | CellId
  | "__end__"
  | { type: "__end__"; columnId: CellColumnId };

export interface CreateNewCellAction {
  cellId: CellPosition;
  /** Whether to insert before (true) or after (false) the target cell */
  before: boolean;
  /** Initial code content for the new cell */
  code?: string;
  /** The last executed code for the new cell */
  lastCodeRun?: string;
  /** Timestamp of the last execution */
  lastExecutionTime?: number;
  /** Optional custom ID for the new cell. Auto-generated if not provided */
  newCellId?: CellId;
  /** Whether to focus the new cell after creation */
  autoFocus?: boolean;
  /** If true, skip creation if code already exists */
  skipIfCodeExists?: boolean;
  /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
  hideCode?: boolean;
}

/**
 * Actions and reducer for the notebook state.
 */
const {
  reducer,
  addMiddleware,
  createActions,
  useActions,
  valueAtom: notebookAtom,
} = createReducerAndAtoms(initialNotebookState, {
  createNewCell: (state, action: CreateNewCellAction) => {
    const {
      cellId,
      before,
      code,
      lastCodeRun = null,
      lastExecutionTime = null,
      autoFocus = true,
      skipIfCodeExists = false,
      hideCode = false,
    } = action;

    let columnId: CellColumnId;
    let cellIndex: number;

    // If skipIfCodeExists is true, check if the code already exists in the notebook
    if (skipIfCodeExists) {
      for (const cellId of state.cellIds.inOrderIds) {
        if (state.cellData[cellId]?.code === code) {
          return state;
        }
      }
    }

    if (cellId === "__end__") {
      const column = state.cellIds.atOrThrow(0);
      columnId = column.id;
      cellIndex = column.length;
    } else if (typeof cellId === "string") {
      const column = state.cellIds.findWithId(cellId);
      columnId = column.id;
      cellIndex = column.topLevelIds.indexOf(cellId);
    } else if (cellId.type === "__end__") {
      const column =
        state.cellIds.get(cellId.columnId) || state.cellIds.atOrThrow(0);
      columnId = column.id;
      cellIndex = column.length;
    } else {
      throw new Error("Invalid cellId");
    }

    const newCellId = action.newCellId || CellId.create();
    const insertionIndex = before ? cellIndex : cellIndex + 1;

    return {
      ...state,
      cellIds: state.cellIds.insertId(newCellId, columnId, insertionIndex),
      cellData: {
        ...state.cellData,
        [newCellId]: createCell({
          id: newCellId,
          code,
          lastCodeRun,
          config: createCellConfig({ hide_code: hideCode }),
          lastExecutionTime,
          edited: Boolean(code) && code !== lastCodeRun,
        }),
      },
      cellRuntime: {
        ...state.cellRuntime,
        [newCellId]: createCellRuntimeState(),
      },
      cellHandles: {
        ...state.cellHandles,
        [newCellId]: createRef(),
      },
      scrollKey: autoFocus ? newCellId : null,
      untouchedNewCells:
        hideCode && !code
          ? new Set([...state.untouchedNewCells, newCellId])
          : state.untouchedNewCells,
    };
  },
  moveCell: (
    state,
    action: {
      cellId: CellId;
      before?: boolean;
      direction?: "left" | "right";
    },
  ) => {
    const { cellId, before, direction } = action;

    if (before !== undefined && direction !== undefined) {
      Logger.warn(
        "Both before and direction specified for moveCell. Ignoring one.",
      );
    }

    // Handle left/right movement
    if (direction) {
      const fromColumn = state.cellIds.findWithId(cellId);
      const fromColumnIndex = state.cellIds.indexOf(fromColumn);
      const toColumnIndex =
        direction === "left" ? fromColumnIndex - 1 : fromColumnIndex + 1;
      const toColumn = state.cellIds.at(toColumnIndex);

      // If no column to move to, return unchanged state
      if (!toColumn) {
        return state;
      }

      return {
        ...state,
        cellIds: state.cellIds.moveAcrossColumns(
          fromColumn.id,
          cellId,
          toColumn.id,
          undefined,
        ),
        scrollKey: cellId,
      };
    }

    // Handle up/down movement
    const column = state.cellIds.findWithId(cellId);
    const cellIndex = column.indexOfOrThrow(cellId);

    if (before && cellIndex === 0) {
      return {
        ...state,
        cellIds: state.cellIds.moveWithinColumn(column.id, cellIndex, 0),
        scrollKey: cellId,
      };
    }
    if (!before && cellIndex === column.length - 1) {
      return {
        ...state,
        cellIds: state.cellIds.moveWithinColumn(
          column.id,
          cellIndex,
          column.length - 1,
        ),
        scrollKey: cellId,
      };
    }

    return before
      ? {
          ...state,
          cellIds: state.cellIds.moveWithinColumn(
            column.id,
            cellIndex,
            cellIndex - 1,
          ),
          scrollKey: cellId,
        }
      : {
          ...state,
          cellIds: state.cellIds.moveWithinColumn(
            column.id,
            cellIndex,
            cellIndex + 1,
          ),
          scrollKey: cellId,
        };
  },
  moveCellToIndex: (
    state,
    action: {
      cellId: CellId;
      columnId: CellColumnId;
      index: number;
    },
  ) => {
    const { cellId, columnId, index } = action;
    const fromColumn = state.cellIds.findWithId(cellId);
    const fromIndex = fromColumn.indexOfOrThrow(cellId);

    const destinationColumn = state.cellIds.get(columnId);
    if (!destinationColumn) {
      return state;
    }

    const clampedIndex = Math.max(0, Math.min(index, destinationColumn.length));
    const adjustedIndex =
      fromColumn.id === columnId && fromIndex < clampedIndex
        ? clampedIndex - 1
        : clampedIndex;

    if (fromColumn.id === columnId && fromIndex === adjustedIndex) {
      return state;
    }

    const withoutCell = state.cellIds.deleteById(cellId);
    const updatedColumn = withoutCell.get(columnId);
    if (!updatedColumn) {
      return state;
    }

    return {
      ...state,
      cellIds: withoutCell.insertId(
        cellId,
        columnId,
        Math.max(0, Math.min(adjustedIndex, updatedColumn.length)),
      ),
      scrollKey: null,
    };
  },
  dropCellOverCell: (state, action: { cellId: CellId; overCellId: CellId }) => {
    const { cellId, overCellId } = action;

    const fromColumn = state.cellIds.findWithId(cellId);
    const toColumn = state.cellIds.findWithId(overCellId);

    const fromIndex = fromColumn.indexOfOrThrow(cellId);
    const toIndex = toColumn.indexOfOrThrow(overCellId);

    if (fromColumn.id === toColumn.id) {
      if (fromIndex === toIndex) {
        return state;
      }
      return {
        ...state,
        cellIds: state.cellIds.moveWithinColumn(
          fromColumn.id,
          fromIndex,
          toIndex,
        ),
        scrollKey: null,
      };
    }

    return {
      ...state,
      cellIds: state.cellIds.moveAcrossColumns(
        fromColumn.id,
        cellId,
        toColumn.id,
        overCellId,
      ),
      scrollKey: null,
    };
  },
  dropCellOverColumn: (
    state,
    action: { cellId: CellId; columnId: CellColumnId },
  ) => {
    const { cellId, columnId } = action;
    const fromColumn = state.cellIds.findWithId(cellId);

    return {
      ...state,
      cellIds: state.cellIds.moveAcrossColumns(
        fromColumn.id,
        cellId,
        columnId,
        undefined,
      ),
    };
  },
  dropOverNewColumn: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    return {
      ...state,
      cellIds: state.cellIds.moveToNewColumn(cellId),
    };
  },
  moveColumn: (
    state,
    action: {
      column: CellColumnId;
      overColumn: CellColumnId | "_left_" | "_right_";
    },
  ) => {
    if (action.column === action.overColumn) {
      return state;
    }
    return {
      ...state,
      cellIds: state.cellIds.moveColumn(action.column, action.overColumn),
    };
  },
  focusCell: (
    state,
    action: { cellId: CellId; where: "before" | "after" | "exact" },
  ) => {
    const column = state.cellIds.findWithId(action.cellId);
    if (column.length === 0) {
      return state;
    }

    const { cellId, where } = action;
    const index = column.indexOfOrThrow(cellId);

    let focusIndex: number;
    if (where === "before") {
      focusIndex = clamp(index - 1, 0, column.length - 1);
    } else if (where === "after") {
      focusIndex = clamp(index + 1, 0, column.length - 1);
    } else {
      focusIndex = index;
    }
    const focusCellId = column.atOrThrow(focusIndex);
    // can scroll immediately, without setting scrollKey in state, because
    // CellArray won't need to re-render
    focusAndScrollCellIntoView({
      cellId: focusCellId,
      cell: state.cellHandles[focusCellId],
      isCodeHidden: isCellCodeHidden(state, focusCellId),
      codeFocus: where === "after" ? "top" : "bottom",
      variableName: undefined,
    });
    return state;
  },
  focusTopCell: (state) => {
    // TODO: focus the existing column, not the first column
    const column = state.cellIds.getColumns().at(0);
    if (column === undefined || column.length === 0) {
      return state;
    }

    const cellId = column.first();
    focusAndScrollCellIntoView({
      cellId: cellId,
      cell: state.cellHandles[cellId],
      isCodeHidden: isCellCodeHidden(state, cellId),
      codeFocus: undefined,
      variableName: undefined,
    });
    scrollToTop();
    return state;
  },
  focusBottomCell: (state) => {
    // TODO: focus the existing column, not the last column
    const column = state.cellIds.getColumns().at(-1);
    if (column === undefined || column.length === 0) {
      return state;
    }

    const cellId = column.last();
    focusAndScrollCellIntoView({
      cellId: cellId,
      cell: state.cellHandles[cellId],
      isCodeHidden: isCellCodeHidden(state, cellId),
      codeFocus: undefined,
      variableName: undefined,
    });
    scrollToBottom();
    return state;
  },
  sendToTop: (state, action: { cellId: CellId; scroll?: boolean }) => {
    const column = state.cellIds.findWithId(action.cellId);
    if (column.length === 0) {
      return state;
    }

    const { cellId, scroll = true } = action;
    const cellIndex = column.indexOfOrThrow(cellId);
    if (cellIndex === 0) {
      return state;
    }

    return {
      ...state,
      cellIds: state.cellIds.moveWithinColumn(column.id, cellIndex, 0),
      scrollKey: scroll ? cellId : null,
    };
  },
  sendToBottom: (state, action: { cellId: CellId; scroll?: boolean }) => {
    const column = state.cellIds.findWithId(action.cellId);
    if (column.length === 0) {
      return state;
    }

    const { cellId, scroll = true } = action;
    const cellIndex = column.indexOfOrThrow(cellId);
    const newIndex = column.length - 1;

    if (cellIndex === newIndex) {
      return state;
    }

    return {
      ...state,
      cellIds: state.cellIds.moveWithinColumn(column.id, cellIndex, newIndex),
      scrollKey: scroll ? cellId : null,
    };
  },
  addColumn: (state, action: { columnId: CellColumnId }) => {
    // Add column and new cell
    const newCellId = CellId.create();
    return {
      ...state,
      cellIds: state.cellIds.addColumn(action.columnId, [newCellId]),
      cellData: {
        ...state.cellData,
        [newCellId]: createCell({
          id: newCellId,
        }),
      },
      cellRuntime: {
        ...state.cellRuntime,
        [newCellId]: createCellRuntimeState(),
      },
      cellHandles: {
        ...state.cellHandles,
        [newCellId]: createRef(),
      },
      scrollKey: newCellId,
    };
  },
  addColumnBreakpoint: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    if (state.cellIds.getColumns()[0].inOrderIds[0] === cellId) {
      return state;
    }
    return {
      ...state,
      cellIds: state.cellIds.insertBreakpoint(cellId),
    };
  },
  deleteColumn: (state, action: { columnId: CellColumnId }) => {
    // Move all cells in the column to the previous column
    const { columnId } = action;
    return {
      ...state,
      cellIds: state.cellIds.delete(columnId),
    };
  },
  mergeAllColumns: (state) => {
    return {
      ...state,
      cellIds: state.cellIds.mergeAllColumns(),
    };
  },
  compactColumns: (state) => {
    return {
      ...state,
      cellIds: state.cellIds.compact(),
    };
  },
  deleteCell: (state, action: { cellId: CellId }) => {
    const cellId = action.cellId;

    // Can't delete the last cell, across all columns
    if (state.cellIds.hasOnlyOneId()) {
      return state;
    }

    let column: CollapsibleTree<CellId>;
    try {
      column = state.cellIds.findWithId(cellId);
    } catch (error) {
      // Expected for kernel-only cells or out-of-order transactions.
      Logger.warn("Skipping delete for missing cellId", {
        cellId,
        error,
      });
      return state;
    }

    const cellIndex = column.indexOfOrThrow(cellId);
    const focusIndex = cellIndex === 0 ? 1 : cellIndex - 1;
    let scrollKey: CellId | null = null;
    if (column.length > 1) {
      scrollKey = column.atOrThrow(focusIndex);
    }

    const editorView = state.cellHandles[cellId].current?.editorView;
    const serializedEditorState = editorView?.state.toJSON({
      history: historyField,
    });
    if (serializedEditorState) {
      serializedEditorState.doc = state.cellData[cellId].code;
    }

    // release the granular atom(s) created for this cell
    releaseCellAtoms(cellId);
    const prevData = state.cellData[cellId];
    return {
      ...state,
      cellIds: state.cellIds.deleteById(cellId),
      history: [
        ...state.history,
        {
          name: prevData.name,
          serializedEditorState: serializedEditorState,
          column: column.id,
          index: cellIndex,
          isSetupCell: cellId === SETUP_CELL_ID,
          config: prevData.config,
        },
      ],
      scrollKey: scrollKey,
    };
  },
  undoDeleteCell: (state) => {
    if (state.history.length === 0) {
      return state;
    }

    const mostRecentlyDeleted = state.history[state.history.length - 1];

    const {
      name,
      serializedEditorState = { doc: "" },
      column,
      index,
      isSetupCell,
      config,
    } = mostRecentlyDeleted;

    const cellId = isSetupCell ? SETUP_CELL_ID : CellId.create();
    const undoCell = createCell({
      id: cellId,
      name,
      code: serializedEditorState.doc,
      edited: serializedEditorState.doc.trim().length > 0,
      serializedEditorState,
      config,
    });

    return {
      ...state,
      cellIds: state.cellIds.insertId(cellId, column, index),
      cellData: {
        ...state.cellData,
        [cellId]: undoCell,
      },
      cellRuntime: {
        ...state.cellRuntime,
        [cellId]: createCellRuntimeState(),
      },
      cellHandles: {
        ...state.cellHandles,
        [cellId]: createRef(),
      },
      history: state.history.slice(0, -1),
      // Scroll to the newly created cell
      scrollKey: cellId,
    };
  },
  clearSerializedEditorState: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    return updateCellData({
      state,
      cellId,
      cellReducer: (cell) => {
        return {
          ...cell,
          serializedEditorState: null,
        };
      },
    });
  },
  updateCellCode: (
    state,
    action: {
      cellId: CellId;
      code: string;
      /**
       * Whether or not the update is a formatting change,
       * if so, the 'edited' state will be handled differently.
       */
      formattingChange: boolean;
    },
  ) => {
    const { cellId, code, formattingChange } = action;
    if (!state.cellData[cellId]) {
      return state;
    }

    return updateCellData({
      state,
      cellId,
      cellReducer: (cell) => {
        // Formatting-only change means we can re-use the last code run
        // if it was not previously edited. And we don't change the edited state.
        return formattingChange
          ? {
              ...cell,
              code: code,
              lastCodeRun: cell.edited ? cell.lastCodeRun : code,
            }
          : {
              ...cell,
              code: code,
              edited: code.trim() !== cell.lastCodeRun,
            };
      },
    });
  },
  updateCellName: (state, action: { cellId: CellId; name: string }) => {
    const { cellId, name } = action;
    return updateCellData({
      state,
      cellId,
      cellReducer: (cell) => {
        return {
          ...cell,
          name: name,
        };
      },
    });
  },
  updateCellConfig: (
    state,
    action: { cellId: CellId; config: Partial<CellConfig> },
  ) => {
    const { cellId, config } = action;
    return updateCellData({
      state,
      cellId,
      cellReducer: (cell) => {
        return {
          ...cell,
          config: { ...cell.config, ...config },
        };
      },
    });
  },
  prepareForRun: (state, action: { cellId: CellId }) => {
    const newState = updateCellRuntimeState({
      state,
      cellId: action.cellId,
      cellReducer: (cell) => {
        return prepareCellForExecution(cell);
      },
    });
    return updateCellData({
      state: newState,
      cellId: action.cellId,
      cellReducer: (cell) => {
        return {
          ...cell,
          edited: false,
          lastCodeRun: cell.code.trim(),
        };
      },
    });
  },
  handleCellMessage: (state, message: CellMessage) => {
    const cellId = message.cell_id;
    let nextState = updateCellRuntimeState({
      state,
      cellId,
      cellReducer: (cell) => {
        return transitionCell(cell, message);
      },
    });
    // When a cell is queued for execution, snapshot the current code
    // as lastCodeRun. This clears staleness for cells executed by the
    // kernel (e.g. via code_mode). If the user edits during execution,
    // code !== lastCodeRun keeps the cell stale.
    if (message.status === "queued") {
      nextState = updateCellData({
        state: nextState,
        cellId,
        cellReducer: (cell) => ({
          ...cell,
          lastCodeRun: cell.code.trim(),
          edited: false,
        }),
      });
    }
    return {
      ...nextState,
      cellLogs: [...nextState.cellLogs, ...getCellLogsForMessage(message)],
    };
  },
  setCellIds: (state, action: { cellIds: CellId[] }) => {
    const isTheSame = isEqual(state.cellIds.inOrderIds, action.cellIds);
    if (isTheSame) {
      return state;
    }

    // Create new cell data and runtime states for the new cell IDs
    const nextCellData = { ...state.cellData };
    const nextCellRuntime = { ...state.cellRuntime };
    const nextCellHandles = { ...state.cellHandles };

    for (const cellId of action.cellIds) {
      if (!(cellId in state.cellData)) {
        nextCellData[cellId] = createCell({ id: cellId });
      }
      if (!(cellId in state.cellRuntime)) {
        nextCellRuntime[cellId] = createCellRuntimeState();
      }
      if (!(cellId in state.cellHandles)) {
        nextCellHandles[cellId] = createRef();
      }
    }

    return {
      ...state,
      cellIds: MultiColumn.fromWithPreviousShape(action.cellIds, state.cellIds),
      cellData: nextCellData,
      cellRuntime: nextCellRuntime,
      cellHandles: nextCellHandles,
    };
  },
  /**
   * Rebuild the MultiColumn tree using each cell's `config.column` value.
   *
   * Used after a transaction whose `set-config` changes updated cells'
   * column metadata without physically moving them in the tree. Cells with
   * `config.column == null` inherit the column of the previous cell in the
   * given order (see `MultiColumn.fromIdsAndColumns`), which lets the server
   * send column changes only at column boundaries.
   */
  rebuildCellColumns: (state, action: { cellIds: CellId[] }) => {
    const newCellIds = MultiColumn.fromIdsAndColumns(
      action.cellIds.map((id) => [
        id,
        state.cellData[id]?.config.column ?? null,
      ]),
    );
    return { ...state, cellIds: newCellIds };
  },
  setCellCodes: (
    state,
    action: {
      codes: string[];
      ids: CellId[];
      codeIsStale: boolean;
      names?: string[];
      configs?: CellConfig[];
    },
  ) => {
    invariant(
      action.codes.length === action.ids.length,
      "Expected codes and ids to have the same length",
    );

    let nextState = { ...state };

    const cellReducer = ({
      cell,
      code,
      cellId,
      name,
      config,
    }: {
      cell: CellData | undefined;
      code: string;
      cellId: CellId;
      name?: string;
      config?: CellConfig;
    }) => {
      if (!cell) {
        return createCell({
          id: cellId,
          code,
          name: name,
          config: config,
          lastCodeRun: action.codeIsStale ? null : code,
          edited: action.codeIsStale && code.trim().length > 0,
        });
      }

      // If code is stale, we don't promote it to lastCodeRun
      const lastCodeRun = action.codeIsStale ? cell.lastCodeRun : code;

      // Mark as edited if the code has changed
      const edited = lastCodeRun
        ? lastCodeRun.trim() !== code.trim()
        : Boolean(code);

      // No change
      if (cell.code.trim() === code.trim()) {
        return {
          ...cell,
          code: code,
          edited,
          lastCodeRun,
          ...(name !== undefined && { name }),
          ...(config !== undefined && { config }),
        };
      }

      // Update codemirror if mounted
      // If RTC is enabled, the editor view will already be updated, so we don't need to do this
      if (!isRtcEnabled()) {
        const cellHandle = nextState.cellHandles[cellId].current;
        if (cellHandle?.editorViewOrNull) {
          updateEditorCodeFromPython(cellHandle.editorViewOrNull, code);
        }
      }

      return {
        ...cell,
        code: code,
        edited,
        lastCodeRun,
        ...(name !== undefined && { name }),
        ...(config !== undefined && { config }),
      };
    };

    for (let i = 0; i < action.ids.length; i++) {
      const cellId = action.ids[i];
      const code = action.codes[i];
      if (cellId === undefined || code === undefined) {
        continue;
      }
      const name = action.names?.[i];
      const config = action.configs?.[i];
      nextState = {
        ...nextState,
        cellData: {
          ...nextState.cellData,
          [cellId]: cellReducer({
            cell: nextState.cellData[cellId],
            code,
            cellId,
            name,
            config,
          }),
        },
      };
    }

    return nextState;
  },
  setStdinResponse: (
    state,
    action: { cellId: CellId; response: string; outputIndex: number },
  ) => {
    const { cellId, response, outputIndex } = action;
    return updateCellRuntimeState({
      state,
      cellId,
      cellReducer: (cell) => {
        const consoleOutputs = [...cell.consoleOutputs];
        const stdinOutput = consoleOutputs[outputIndex];
        if (stdinOutput == null || stdinOutput.channel !== "stdin") {
          Logger.warn("Expected stdin output");
          return cell;
        }

        consoleOutputs[outputIndex] = {
          channel: "stdin",
          mimetype: stdinOutput.mimetype,
          data: stdinOutput.data,
          timestamp: stdinOutput.timestamp,
          response,
        };

        return {
          ...cell,
          interrupted: false,
          consoleOutputs,
        };
      },
    });
  },
  setCells: (state, cells: CellData[]) => {
    const cellData = Object.fromEntries(cells.map((cell) => [cell.id, cell]));

    const cellRuntime = Object.fromEntries(
      cells.map((cell) => [cell.id, createCellRuntimeState()]),
    );

    return withScratchCell({
      ...state,
      cellIds: MultiColumn.fromIdsAndColumns(
        cells.map((cell) => [cell.id, cell.config.column]),
      ),
      cellData: cellData,
      cellRuntime: cellRuntime,
      cellHandles: Object.fromEntries(
        cells.map((cell) => [cell.id, createRef()]),
      ),
    });
  },
  /**
   * Move focus to next cell
   *
   * Creates a new cell if the current cell is the last one in the array.
   *
   * If needed, scrolls newly created or focused cell into view.
   *
   * Replicates Shift+Enter functionality of Jupyter
   */
  moveToNextCell: (
    state,
    action: { cellId: CellId; before: boolean; noCreate?: boolean },
  ) => {
    const { cellId, before, noCreate = false } = action;

    // Can't move focus of scratch cell
    if (cellId === SCRATCH_CELL_ID) {
      return state;
    }

    const column = state.cellIds.findWithId(cellId);
    const index = column.indexOfOrThrow(cellId);
    const nextCellIndex = before ? index - 1 : index + 1;

    const isPastLastCell = nextCellIndex === column.length;
    const isBeforeFirstCell = nextCellIndex === -1;

    // Create a new cell at the end and set scrollKey to focus it
    if (isPastLastCell && !noCreate) {
      const newCellId = CellId.create();
      return {
        ...state,
        cellIds: state.cellIds.insertId(newCellId, column.id, nextCellIndex),
        cellData: {
          ...state.cellData,
          [newCellId]: createCell({ id: newCellId }),
        },
        cellRuntime: {
          ...state.cellRuntime,
          [newCellId]: createCellRuntimeState(),
        },
        cellHandles: {
          ...state.cellHandles,
          [newCellId]: createRef(),
        },
        scrollKey: newCellId,
      };
    }

    if (isBeforeFirstCell && !noCreate) {
      const newCellId = CellId.create();
      return {
        ...state,
        cellIds: state.cellIds.insertId(newCellId, column.id, 0),
        cellData: {
          ...state.cellData,
          [newCellId]: createCell({ id: newCellId }),
        },
        cellRuntime: {
          ...state.cellRuntime,
          [newCellId]: createCellRuntimeState(),
        },
        cellHandles: {
          ...state.cellHandles,
          [newCellId]: createRef(),
        },
        scrollKey: newCellId,
      };
    }

    if ((isPastLastCell || isBeforeFirstCell) && noCreate) {
      // Trying to move past the last cell or before the first cell
      // with noCreate=true - do nothing
      return state;
    }

    if (nextCellIndex >= 0 && nextCellIndex < column.length) {
      // Move to valid cell within bounds
      const nextCellId = column.atOrThrow(nextCellIndex);
      // Just focus, no state change
      focusAndScrollCellIntoView({
        cellId: nextCellId,
        cell: state.cellHandles[nextCellId],
        isCodeHidden: isCellCodeHidden(state, nextCellId),
        codeFocus: before ? "bottom" : "top",
        variableName: undefined,
      });
    }

    return state;
  },
  markTouched: (state, action: { cellId: CellId }) => {
    const { cellId } = action;

    if (state.untouchedNewCells.has(cellId)) {
      const nextUntouchedNewCells = new Set(state.untouchedNewCells);
      nextUntouchedNewCells.delete(cellId);
      return {
        ...state,
        untouchedNewCells: nextUntouchedNewCells,
      };
    }

    return state;
  },
  markUntouched: (state, action: { cellId: CellId }) => {
    const { cellId } = action;

    if (!state.untouchedNewCells.has(cellId)) {
      const nextUntouchedNewCells = new Set(state.untouchedNewCells);
      nextUntouchedNewCells.add(cellId);
      return {
        ...state,
        untouchedNewCells: nextUntouchedNewCells,
      };
    }

    return state;
  },
  scrollToTarget: (state) => {
    // Scroll to the specified cell and clear the scroll key.
    const scrollKey = state.scrollKey;
    if (scrollKey === null) {
      return state;
    }

    const column = state.cellIds.findWithId(scrollKey);
    const index = column.indexOfOrThrow(scrollKey);

    const cellId =
      index === column.length - 1 ? column.last() : column.atOrThrow(index);

    focusAndScrollCellIntoView({
      cellId: cellId,
      cell: state.cellHandles[cellId],
      isCodeHidden: isCellCodeHidden(state, cellId),
      codeFocus: undefined,
      variableName: undefined,
    });

    return {
      ...state,
      scrollKey: null,
    };
  },
  foldAll: (state) => {
    const targets = Object.values(state.cellHandles).map(
      (handle) => handle.current?.editorView,
    );
    foldAllBulk(targets);
    return state;
  },
  unfoldAll: (state) => {
    const targets = Object.values(state.cellHandles).map(
      (handle) => handle.current?.editorView,
    );
    unfoldAllBulk(targets);
    return state;
  },
  clearLogs: (state) => {
    return {
      ...state,
      cellLogs: [],
    };
  },
  collapseCell: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    const column = state.cellIds.findWithId(cellId);

    // Get all the top-level outlines
    const outlines = column.topLevelIds.map((id) => {
      const cell = state.cellRuntime[id];
      return cell.outline;
    });

    // Find the start/end of the collapsed range
    const startIndex = column.indexOfOrThrow(cellId);
    const range = findCollapseRange(startIndex, outlines);
    if (!range) {
      return state;
    }
    const endCellId = column.atOrThrow(range[1]);

    return {
      ...state,
      // Collapse the range
      cellIds: state.cellIds.transformWithCellId(cellId, (column) =>
        column.collapse(cellId, endCellId),
      ),
      scrollKey: cellId,
    };
  },
  expandCell: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    return {
      ...state,
      cellIds: state.cellIds.transformWithCellId(cellId, (column) =>
        column.expand(cellId),
      ),
      scrollKey: cellId,
    };
  },
  collapseAllCells: (state) => {
    return {
      ...state,
      cellIds: state.cellIds.transformAll((column) => {
        // Get all the top-level outlines
        const outlines = column.topLevelIds.map((id) => {
          const cell = state.cellRuntime[id];
          return cell.outline;
        });

        // Find the start/end of the collapsed ranges
        const nodes = [...column.nodes];
        const rangeIndexes: {
          start: CellIndex;
          end: CellIndex;
        }[] = [];
        const reversedCollapseRanges = [];

        // Iterate in reverse order (bottom-up) to process children first
        let i = nodes.length - 1;
        while (i >= 0) {
          const range = findCollapseRange(i, outlines);
          if (range) {
            const startIndex = i;
            let endIndex = range[1];

            // Check if the parent's end point is inside any already-collapsed child range
            const parentEndInChild = rangeIndexes.find(
              (child) => child.start <= endIndex && child.end === endIndex,
            );

            if (parentEndInChild) {
              // Adjust the new endIndex to the child's start
              endIndex = parentEndInChild.start;
            }

            // Store this range for future child checks
            rangeIndexes.push({ start: startIndex, end: endIndex });

            // Add the range to the list of ranges
            const cellId = column.atOrThrow(startIndex);
            const until = column.atOrThrow(endIndex);
            reversedCollapseRanges.push({ id: cellId, until });
          } else {
            reversedCollapseRanges.push(null);
          }
          i--;
        }

        const collapseRanges = reversedCollapseRanges.toReversed();
        return column.collapseAll(collapseRanges);
      }),
    };
  },
  expandAllCells: (state) => {
    return {
      ...state,
      cellIds: state.cellIds.transformAll((column) => column.expandAll()),
    };
  },
  showCellIfHidden: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    const column = state.cellIds.findWithId(cellId);
    const prev = column;
    const result = column.findAndExpandDeep(cellId);

    if (result.equals(prev)) {
      return state;
    }

    return {
      ...state,
      cellIds: state.cellIds.transformWithCellId(cellId, () => result),
    };
  },
  splitCell: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    const column = state.cellIds.findWithId(cellId);
    const index = column.indexOfOrThrow(cellId);
    const cell = state.cellData[cellId];
    const cellHandle = state.cellHandles[cellId].current;

    if (cellHandle?.editorView == null) {
      return state;
    }

    const { beforeCursorCode, afterCursorCode } = splitEditor(
      cellHandle.editorView,
    );

    updateEditorCodeFromPython(cellHandle.editorView, beforeCursorCode);

    const newCellId = CellId.create();

    return {
      ...state,
      cellIds: state.cellIds.insertId(newCellId, column.id, index + 1),
      cellData: {
        ...state.cellData,
        [cellId]: {
          ...cell,
          code: beforeCursorCode,
          edited:
            Boolean(beforeCursorCode) &&
            beforeCursorCode.trim() !== cell.lastCodeRun?.trim(),
        },
        [newCellId]: createCell({
          id: newCellId,
          code: afterCursorCode,
          edited: Boolean(afterCursorCode),
        }),
      },
      cellRuntime: {
        ...state.cellRuntime,
        [cellId]: {
          ...state.cellRuntime[cellId],
          output: null,
          consoleOutputs: [],
        },
        [newCellId]: createCellRuntimeState(),
      },
      cellHandles: {
        ...state.cellHandles,
        [newCellId]: createRef(),
      },
      scrollKey: newCellId,
    };
  },
  undoSplitCell: (state, action: { cellId: CellId; snapshot: string }) => {
    const { cellId, snapshot } = action;

    const cell = state.cellData[cellId];
    const cellHandle = state.cellHandles[cellId].current;

    if (cellHandle?.editorView == null) {
      return state;
    }

    updateEditorCodeFromPython(cellHandle.editorView, snapshot);

    return {
      ...state,
      cellIds: state.cellIds.transformWithCellId(cellId, (column) => {
        const newCellIndex = column.indexOfOrThrow(cellId) + 1;
        return column.deleteAtIndex(newCellIndex);
      }),
      cellData: {
        ...state.cellData,
        [cellId]: {
          ...cell,
          code: snapshot,
          edited:
            Boolean(snapshot) && snapshot?.trim() !== cell.lastCodeRun?.trim(),
        },
      },
      cellRuntime: {
        ...state.cellRuntime,
        [cellId]: {
          ...state.cellRuntime[cellId],
          output: null,
          consoleOutputs: [],
        },
      },
      cellHandles: {
        ...state.cellHandles,
      },
    };
  },
  clearCellOutput: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    return updateCellRuntimeState({
      state,
      cellId,
      cellReducer: (cell) => ({
        ...cell,
        output: null,
        consoleOutputs: [],
      }),
    });
  },
  clearCellConsoleOutput: (state, action: { cellId: CellId }) => {
    const { cellId } = action;
    return updateCellRuntimeState({
      state,
      cellId,
      cellReducer: (cell) => ({
        ...cell,
        // Remove everything except unresponsed stdin
        consoleOutputs: cell.consoleOutputs.filter(
          (output) => output.channel === "stdin" && output.response == null,
        ),
      }),
    });
  },
  clearAllCellOutputs: (state) => {
    const newCellRuntime = { ...state.cellRuntime };
    for (const cellId of state.cellIds.inOrderIds) {
      newCellRuntime[cellId] = {
        ...newCellRuntime[cellId],
        output: null,
        consoleOutputs: [],
      };
    }
    return {
      ...state,
      cellRuntime: newCellRuntime,
    };
  },
  addSetupCellIfDoesntExist: (state, action: { code?: string }) => {
    const { code } = action;

    if (state.cellIds.setupCellExists()) {
      // Just focus on the existing setup cell
      return {
        ...state,
        scrollKey: SETUP_CELL_ID,
      };
    }

    return {
      ...state,
      cellIds: state.cellIds.insertId(
        SETUP_CELL_ID,
        state.cellIds.atOrThrow(0).id,
        0,
      ),
      cellData: {
        ...state.cellData,
        [SETUP_CELL_ID]: createCell({
          id: SETUP_CELL_ID,
          name: SETUP_CELL_ID,
          code,
          edited: Boolean(code),
        }),
      },
      cellRuntime: {
        ...state.cellRuntime,
        [SETUP_CELL_ID]: createCellRuntimeState(),
      },
      cellHandles: {
        ...state.cellHandles,
        [SETUP_CELL_ID]: createRef(),
      },
      scrollKey: SETUP_CELL_ID,
    };
  },
});

// We apply the middleware here (rather than inline in createReducerAndAtoms)
// so that the document transaction middleware can import CellActions and
// strictly type the dispatched actions without creating a circular dependency.
addMiddleware(documentTransactionMiddleware);

function isCellCodeHidden(state: NotebookState, cellId: CellId): boolean {
  return (
    Boolean(state.cellData[cellId].config.hide_code) &&
    !state.untouchedNewCells.has(cellId)
  );
}

// Helper function to update a cell in the array
function updateCellRuntimeState({
  state,
  cellId,
  cellReducer,
}: {
  state: NotebookState;
  cellId: CellId;
  cellReducer: ReducerWithoutAction<CellRuntimeState>;
}) {
  if (!(cellId in state.cellRuntime)) {
    Logger.warn(`Cell ${cellId} not found in state`);
    return state;
  }

  return {
    ...state,
    cellRuntime: {
      ...state.cellRuntime,
      [cellId]: cellReducer(state.cellRuntime[cellId]),
    },
  };
}

function updateCellData({
  state,
  cellId,
  cellReducer,
}: {
  state: NotebookState;
  cellId: CellId;
  cellReducer: ReducerWithoutAction<CellData>;
}) {
  if (!(cellId in state.cellData)) {
    Logger.warn(`Cell ${cellId} not found in state`);
    return state;
  }

  return {
    ...state,
    cellData: {
      ...state.cellData,
      [cellId]: cellReducer(state.cellData[cellId]),
    },
  };
}

export function getCellConfigs(state: NotebookState): CellConfig[] {
  const cells = state.cellData;

  // We set to null by default to prevent inconsistencies between undefined & null
  const defaultCellConfig: Partial<CellConfig> = { column: null };

  // Handle the case where there's only one column
  // We don't want to set the column config
  const hasMultipleColumns = state.cellIds.getColumns().length > 1;
  if (!hasMultipleColumns) {
    return state.cellIds.getColumns().flatMap((column) => {
      return column.inOrderIds.map((cellId) => {
        return {
          ...cells[cellId].config,
          ...defaultCellConfig,
        };
      });
    });
  }

  return state.cellIds.getColumns().flatMap((column, columnIndex) => {
    return column.inOrderIds.map((cellId, cellIndex) => {
      const config: Partial<CellConfig> = { ...defaultCellConfig };

      // Only set the column index for the first cell in the column
      if (cellIndex === 0) {
        config.column = columnIndex;
      }

      const newConfig = {
        ...cells[cellId].config,
        ...config,
      };

      return newConfig;
    });
  });
}

export {
  createActions as createNotebookActions,
  reducer as notebookReducer,
  notebookAtom,
};

/// ATOMS

export const cellIdsAtom = atom((get) => get(notebookAtom).cellIds);

export const hasOnlyOneCellAtom = atom((get) =>
  get(cellIdsAtom).hasOnlyOneId(),
);

export const hasDisabledCellsAtom = atom(
  (get) => disabledCellIds(get(notebookAtom)).length > 0,
);
export const hasEnabledCellsAtom = atom(
  (get) => enabledCellIds(get(notebookAtom)).length > 0,
);
export const canUndoDeletesAtom = atom((get) =>
  canUndoDeletes(get(notebookAtom)),
);

export const needsRunAtom = atom((get) => notebookNeedsRun(get(notebookAtom)));

export const cellErrorsAtom = atom((get) => {
  const { cellIds, cellRuntime, cellData } = get(notebookAtom);
  const errors = cellIds.inOrderIds
    .map((cellId) => {
      const cell = cellRuntime[cellId];
      const { name } = cellData[cellId];
      if (isErrorMime(cell.output?.mimetype)) {
        // Filter out ancestor-stopped errors
        // These are errors that are caused by a cell that was stopped,
        // but nothing the user can take action on.
        invariant(Array.isArray(cell.output.data), "Expected array data");
        const nonAncestorErrors = cell.output.data.filter(
          (error) => !error.type.includes("ancestor"),
        );

        if (nonAncestorErrors.length > 0) {
          return {
            output: { ...cell.output, data: nonAncestorErrors },
            cellId: cellId,
            cellName: name,
          };
        }
      }

      return null;
    })
    .filter(Boolean);
  return errors;
});

export const notebookOutline = atom((get) => {
  const { cellIds, cellRuntime } = get(notebookAtom);
  const outlines = cellIds.inOrderIds.map(
    (cellId) => cellRuntime[cellId].outline,
  );
  return mergeOutlines(outlines);
});

export const cellErrorCount = atom((get) => get(cellErrorsAtom).length);

export const cellIdToNamesMap = createDeepEqualAtom(
  atom((get) => {
    const { cellIds, cellData } = get(notebookAtom);
    const names: Record<CellId, string | undefined> = Objects.fromEntries(
      cellIds.inOrderIds.map((cellId) => [cellId, cellData[cellId]?.name]),
    );
    return names;
  }),
);

const scrollKeyAtom = atom((get) => get(notebookAtom).scrollKey);

/// HOOKS

/**
 * React-hook for the array of cells.
 */
export const useNotebook = () => useAtomValue(notebookAtom);

/**
 * React-hook for the array of cell IDs.
 */
export const useCellIds = () => useAtomValue(cellIdsAtom);

/**
 * React-hook for the dictionary of cell names
 */
export const useCellNames = () => useAtomValue(cellIdToNamesMap);

/**
 * React-hook for the array of cell errors.
 */
export const useCellErrors = () => useAtomValue(cellErrorsAtom);

/**
 * React-hook for the cell logs.
 */
export const useCellLogs = () => useAtomValue(notebookAtom).cellLogs;

/**
 * React-hook for the notebook scrollKey
 */
export const useScrollKey = () => useAtomValue(scrollKeyAtom);

/// IMPERATIVE GETTERS

/**
 * Get the array of cell IDs.
 */
export const getNotebook = () => store.get(notebookAtom);

/**
 * Get the array of cell IDs.
 */
export const getCells = () => store.get(notebookAtom).cellIds;

/**
 * Get the array of cell names
 */
export const getCellNames = () => {
  const { cellIds, cellData } = store.get(notebookAtom);
  return cellIds.inOrderIds.map((id) => cellData[id]?.name).filter(Boolean);
};

const cellDataAtoms = splitAtom(
  selectAtom(notebookAtom, (cells) =>
    cells.cellIds.inOrderIds.map((id) => cells.cellData[id]),
  ),
);
export const useCellDataAtoms = () => useAtom(cellDataAtoms);

export const cellsRuntimeAtom = atom((get) => get(notebookAtom).cellRuntime);

export const notebookIsRunningAtom = atom((get) =>
  notebookIsRunning(get(notebookAtom)),
);
export const onlyScratchpadIsRunningAtom = atom((get) => {
  const { cellRuntime } = get(notebookAtom);
  return Object.entries(cellRuntime).every(
    ([id, rt]) => rt.status !== "running" || id === SCRATCH_CELL_ID,
  );
});
export const notebookQueuedOrRunningCountAtom = atom((get) =>
  notebookQueueOrRunningCount(get(notebookAtom)),
);

export const numColumnsAtom = atom(
  (get) => get(notebookAtom).cellIds.colLength,
);
export const hasCellsAtom = atom(
  (get) => get(notebookAtom).cellIds.idLength > 0,
);
export const columnIdsAtom = atom((get) =>
  get(notebookAtom).cellIds.getColumnIds(),
);

export const cellDataAtom = atomFamily((cellId: CellId) =>
  atom((get) => get(notebookAtom).cellData[cellId]),
);
const cellRuntimeAtom = atomFamily((cellId: CellId) =>
  atom((get) => get(notebookAtom).cellRuntime[cellId]),
);
export const cellHandleAtom = atomFamily((cellId: CellId) =>
  atom((get) => get(notebookAtom).cellHandles[cellId]),
);
/**
 * Cleans up atomFamily cache entries for the given cell.
 *
 * Jotai's atomFamily retains a cache of created atoms, which can cause memory leaks
 * if not explicitly removed. This function removes the atoms associated with a specific
 * cellId to free up memory.
 *
 * @param cellId - The cell ID whose atoms should be removed.
 * @see https://jotai.org/docs/utilities/family#caveat-memory-leaks
 */
export function releaseCellAtoms(cellId: CellId) {
  cellDataAtom.remove(cellId);
  cellRuntimeAtom.remove(cellId);
  cellHandleAtom.remove(cellId);
}

/** Subscribes to reactive updates of the cell's data. */
export const useCellData = (cellId: CellId) =>
  useAtomValue(cellDataAtom(cellId));
/** Subscribes to reactive updates of the cell's runtime info. */
export const useCellRuntime = (cellId: CellId) =>
  useAtomValue(cellRuntimeAtom(cellId));
/** Subscribes to reactive updates of the cell's handle (e.g. refs or UI bindings). */
export const useCellHandle = (cellId: CellId) =>
  useAtomValue(cellHandleAtom(cellId));

/**
 * Get the editor views for all cells.
 */
export const getAllEditorViews = () => {
  const { cellIds, cellHandles } = store.get(notebookAtom);
  return cellIds.inOrderIds
    .map((cellId) => cellHandles[cellId]?.current?.editorViewOrNull)
    .filter(Boolean);
};

export const getCellEditorView = (cellId: CellId) => {
  const { cellHandles } = store.get(notebookAtom);
  return cellHandles[cellId].current?.editorView;
};

export function flattenTopLevelNotebookCells(
  state: NotebookState,
): (CellData & CellRuntimeState)[] {
  const { cellIds, cellData, cellRuntime } = state;
  return cellIds.getColumns().flatMap((column) =>
    column.topLevelIds.map((cellId) => ({
      ...cellData[cellId],
      ...cellRuntime[cellId],
    })),
  );
}

export function createUntouchedCellAtom(cellId: CellId): Atom<boolean> {
  return atom((get) => get(notebookAtom).untouchedNewCells.has(cellId));
}

export function createTracebackInfoAtom(
  cellId: CellId,
): Atom<TracebackInfo[] | undefined> {
  //use existing cellRuntimeAtom for intermediate computation
  const cellRuntime = cellRuntimeAtom(cellId);

  return atom((get) => {
    const data = get(cellRuntime);

    if (!data) {
      return undefined;
    }

    if (data.status === "queued" || data.status === "running") {
      return undefined;
    }

    const tracebackInfo: TracebackInfo[] = [];

    // Runtime errors (ZeroDivisionError, etc.)
    const outputs = data.consoleOutputs;
    if (outputs && outputs.length > 0) {
      const firstTraceback = outputs.find(
        (output) => output.mimetype === "application/vnd.marimo+traceback",
      );
      if (firstTraceback) {
        const traceback = firstTraceback.data as string;
        tracebackInfo.push(...extractAllTracebackInfo(traceback));
      }
    }

    // Syntax errors
    const output = data.output;
    if (output?.mimetype === "application/vnd.marimo+error") {
      const errors = output.data;
      if (Array.isArray(errors)) {
        for (const error of errors) {
          if (error.type === "syntax" && error.lineno != null) {
            tracebackInfo.push({
              kind: "cell",
              cellId: cellId,
              lineNumber: error.lineno,
            });
          }
          if (error.type === "import-star" && error.lineno != null) {
            tracebackInfo.push({
              kind: "cell",
              cellId: cellId,
              lineNumber: error.lineno,
            });
          }
        }
      }
    }

    return tracebackInfo.length > 0 ? tracebackInfo : undefined;
  });
}

/**
 * Use this hook to dispatch cell actions. This hook will not cause a re-render
 * when cells change.
 */
export function useCellActions(
  options: { skipMiddleware?: boolean } = {},
): CellActions {
  return useActions(options);
}

/**
 * Map of cell actions
 */
export type CellActions = ReturnType<typeof createActions>;

/**
 * This is exported for testing purposes only.
 */
export const exportedForTesting = {
  reducer,
  createActions,
  initialNotebookState,
  isCellCodeHidden,
  // Export atom families for testing cleanup
  cellDataAtom,
  cellRuntimeAtom,
  cellHandleAtom,
  createTracebackInfoAtom,
};