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

/**
 * Tests for toDocumentChanges (action → change mapping) as a pure function.
 *
 * Each test calls the reducer to produce prevState/newState, then calls
 * toDocumentChanges directly — no middleware, no side effects. Basic round-trip
 * correctness is covered by document-roundtrip.test.ts. This file focuses on:
 *
 * - anchorOf edge cases (before vs after, first cell, __end__)
 * - Field name mapping (hide_code → hideCode)
 * - Column structure actions emitting reorder-cells
 * - Actions that should NOT emit changes
 */

import { python } from "@codemirror/lang-python";
import { EditorState } from "@codemirror/state";
import { EditorView } from "@codemirror/view";
import {
  afterAll,
  afterEach,
  beforeAll,
  beforeEach,
  describe,
  expect,
  it,
  vi,
} from "vitest";
import { cellId } from "@/__tests__/branded";
import type { CellHandle } from "@/components/editor/notebook-cell";
import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
import { requestClientAtom } from "@/core/network/requests";
import type { EditRequests, RunRequests } from "@/core/network/types";
import { store } from "@/core/state/jotai";
import { Deferred } from "@/utils/Deferred";
import { MultiColumn } from "@/utils/id-tree";
import { exportedForTesting, type NotebookState } from "../cells";
import {
  type CellAction,
  abortDocumentResync,
  beginDocumentResync,
  coalesceChanges,
  completeDocumentResync,
  type DocumentChange,
  flushDocumentChanges,
  exportedForTesting as middlewareExports,
  toDocumentChanges,
  withDocumentSave,
} from "../document-changes";
import { CellId } from "../ids";

const { initialNotebookState, reducer } = exportedForTesting;

function createEditor(code: string) {
  const state = EditorState.create({
    doc: code,
    extensions: [
      python(),
      adaptiveLanguageConfiguration({
        cellId: cellId("cell1"),
        completionConfig: {
          activate_on_typing: true,
          signature_hint_on_typing: false,
          copilot: false,
          codeium_api_key: null,
        },
        hotkeys: new OverridingHotkeyProvider({}),
        placeholderType: "marimo-import",
        lspConfig: {},
      }),
    ],
  });
  return new EditorView({ state, parent: document.body });
}

/** Dispatch an action through the reducer and auto-create editor handles. */
function dispatch(state: NotebookState, action: CellAction): NotebookState {
  const next = reducer(state, action);
  for (const [cellIdString, handle] of Object.entries(next.cellHandles)) {
    const cid = cellIdString as CellId;
    if (!handle.current) {
      const view = createEditor(next.cellData[cid].code);
      const h: CellHandle = { editorView: view, editorViewOrNull: view };
      next.cellHandles[cid] = { current: h };
    }
  }
  return next;
}

/** Dispatch an action and return the changes it produces. */
function resolve(state: NotebookState, action: CellAction) {
  const next = dispatch(state, action);
  const changes = toDocumentChanges(state, next, action);
  return { next, changes };
}

let state: NotebookState;

let i = 0;
const originalCreate = CellId.create.bind(CellId);

beforeAll(() => {
  CellId.create = () => cellId(`${i++}`);
});

beforeEach(() => {
  i = 0;
  state = initialNotebookState();
  state.cellIds = MultiColumn.from([]);
});

afterAll(() => {
  CellId.create = originalCreate;
});

function setup(...codes: string[]) {
  for (const code of codes) {
    state = dispatch(state, {
      type: "createNewCell",
      payload: {
        cellId: "__end__",
        before: false,
        code,
        newCellId: CellId.create(),
      },
    });
  }
}

describe("toDocumentChanges", () => {
  describe("anchorOf edge cases", () => {
    it("uses before anchor when new cell is first", () => {
      setup("a");
      const [a] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "createNewCell",
        payload: {
          cellId: a,
          before: true,
          code: "first",
          newCellId: CellId.create(),
        },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "before": "0",
            "cellId": "1",
            "code": "first",
            "config": {
              "column": null,
              "disabled": false,
              "expand_output": false,
              "hide_code": false,
            },
            "name": "_",
            "type": "create-cell",
          },
        ]
      `);
    });

    it("uses after anchor for __end__ insertion", () => {
      setup("a");
      const { changes } = resolve(state, {
        type: "createNewCell",
        payload: {
          cellId: "__end__",
          before: false,
          code: "last",
          newCellId: CellId.create(),
        },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "after": "0",
            "cellId": "1",
            "code": "last",
            "config": {
              "column": null,
              "disabled": false,
              "expand_output": false,
              "hide_code": false,
            },
            "name": "_",
            "type": "create-cell",
          },
        ]
      `);
    });

    it("uses before when sendToTop moves cell to first position", () => {
      setup("a", "b", "c");
      const [, b] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "sendToTop",
        payload: { cellId: b },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "before": "0",
            "cellId": "1",
            "type": "move-cell",
          },
        ]
      `);
    });
  });

  describe("field name mapping", () => {
    it("maps hide_code to hideCode in set-config", () => {
      setup("a");
      const [a] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "updateCellConfig",
        payload: { cellId: a, config: { hide_code: true } },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "cellId": "0",
            "column": null,
            "disabled": false,
            "expandOutput": false,
            "hideCode": true,
            "type": "set-config",
          },
        ]
      `);
    });

    it("maps expand_output to expandOutput in set-config", () => {
      setup("a");
      const [a] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "updateCellConfig",
        payload: { cellId: a, config: { expand_output: true } },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "cellId": "0",
            "column": null,
            "disabled": false,
            "expandOutput": true,
            "hideCode": false,
            "type": "set-config",
          },
        ]
      `);
    });

    it("includes full CellConfig in create-cell", () => {
      setup("a");

      const { changes } = resolve(state, {
        type: "createNewCell",
        payload: {
          cellId: "__end__",
          before: false,
          code: "hidden",
          newCellId: CellId.create(),
          config: { expand_output: true },
          hideCode: true,
        },
      });

      expect(changes[0]).toMatchObject({
        type: "create-cell",
        config: { hide_code: true, expand_output: true },
      });
    });
  });

  describe("column structure actions", () => {
    it("dropOverNewColumn emits set-config + reorder-cells", () => {
      setup("a", "b");
      const [, b] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "dropOverNewColumn",
        payload: { cellId: b },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "cellId": "1",
            "column": 1,
            "disabled": false,
            "expandOutput": false,
            "hideCode": false,
            "type": "set-config",
          },
          {
            "cellIds": [
              "0",
              "1",
            ],
            "type": "reorder-cells",
          },
        ]
      `);
    });

    it("addColumnBreakpoint emits set-config + reorder-cells", () => {
      setup("a", "b", "c");
      const [, b] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "addColumnBreakpoint",
        payload: { cellId: b },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "cellId": "1",
            "column": 1,
            "disabled": false,
            "expandOutput": false,
            "hideCode": false,
            "type": "set-config",
          },
          {
            "cellId": "2",
            "column": 1,
            "disabled": false,
            "expandOutput": false,
            "hideCode": false,
            "type": "set-config",
          },
          {
            "cellIds": [
              "0",
              "1",
              "2",
            ],
            "type": "reorder-cells",
          },
        ]
      `);
    });

    it("mergeAllColumns emits set-config + reorder-cells", () => {
      setup("a", "b", "c");
      const [, b] = state.cellIds.inOrderIds;
      state = dispatch(state, {
        type: "addColumnBreakpoint",
        payload: { cellId: b },
      });

      const { changes } = resolve(state, {
        type: "mergeAllColumns",
        payload: {},
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "cellId": "1",
            "column": 0,
            "disabled": false,
            "expandOutput": false,
            "hideCode": false,
            "type": "set-config",
          },
          {
            "cellId": "2",
            "column": 0,
            "disabled": false,
            "expandOutput": false,
            "hideCode": false,
            "type": "set-config",
          },
          {
            "cellIds": [
              "0",
              "1",
              "2",
            ],
            "type": "reorder-cells",
          },
        ]
      `);
    });
  });

  describe("cell lifecycle actions", () => {
    it("addColumn emits create-cell + column layout changes", () => {
      setup("a");
      const [a] = state.cellIds.inOrderIds;
      const columnId = state.cellIds.findWithId(a).id;

      const { changes } = resolve(state, {
        type: "addColumn",
        payload: { columnId },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "after": "0",
            "cellId": "1",
            "code": "",
            "config": {
              "column": null,
              "disabled": false,
              "expand_output": false,
              "hide_code": false,
            },
            "name": "_",
            "type": "create-cell",
          },
          {
            "cellId": "1",
            "column": 1,
            "disabled": false,
            "expandOutput": false,
            "hideCode": false,
            "type": "set-config",
          },
          {
            "cellIds": [
              "0",
              "1",
            ],
            "type": "reorder-cells",
          },
        ]
      `);
    });

    it("undoDeleteCell emits create-cell for restored cell", () => {
      setup("a", "b");
      const [, b] = state.cellIds.inOrderIds;

      // Delete cell b, then undo
      state = dispatch(state, {
        type: "deleteCell",
        payload: { cellId: b },
      });

      const { changes } = resolve(state, {
        type: "undoDeleteCell",
        payload: {},
      });

      // Restored cell should have original code
      expect(changes[0]).toMatchObject({
        type: "create-cell",
        code: "b",
      });
    });

    it("splitCell emits set-code + create-cell", () => {
      setup("line1\nline2");
      const [a] = state.cellIds.inOrderIds;

      // Position cursor at end of "line1" (position 5)
      const view = state.cellHandles[a].current!.editorView!;
      view.dispatch({ selection: { anchor: 5 } });

      const { changes } = resolve(state, {
        type: "splitCell",
        payload: { cellId: a },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "cellId": "0",
            "code": "line1",
            "type": "set-code",
          },
          {
            "after": "0",
            "cellId": "1",
            "code": "line2",
            "config": {
              "column": null,
              "disabled": false,
              "expand_output": false,
              "hide_code": false,
            },
            "name": "_",
            "type": "create-cell",
          },
        ]
      `);
    });

    it("undoSplitCell emits set-code + delete-cell", () => {
      setup("line1\nline2");
      const [a] = state.cellIds.inOrderIds;

      // Split the cell first
      const view = state.cellHandles[a].current!.editorView!;
      view.dispatch({ selection: { anchor: 5 } });
      const snapshot = view.state.doc.toString();
      state = dispatch(state, {
        type: "splitCell",
        payload: { cellId: a },
      });
      const [, newCell] = state.cellIds.inOrderIds;

      // Now undo the split
      const { changes } = resolve(state, {
        type: "undoSplitCell",
        payload: { cellId: a, snapshot },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "cellId": "0",
            "code": "line1
        line2",
            "type": "set-code",
          },
          {
            "cellId": "${newCell}",
            "type": "delete-cell",
          },
        ]
      `);
    });

    it("moveToNextCell emits create-cell when past last cell", () => {
      setup("a");
      const [a] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "moveToNextCell",
        payload: { cellId: a, before: false },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "after": "0",
            "cellId": "1",
            "code": "",
            "config": {
              "column": null,
              "disabled": false,
              "expand_output": false,
              "hide_code": false,
            },
            "name": "_",
            "type": "create-cell",
          },
        ]
      `);
    });

    it("moveToNextCell emits nothing when moving within bounds", () => {
      setup("a", "b");
      const [a] = state.cellIds.inOrderIds;

      const { changes } = resolve(state, {
        type: "moveToNextCell",
        payload: { cellId: a, before: false },
      });

      expect(changes).toHaveLength(0);
    });

    it("addSetupCellIfDoesntExist emits create-cell when new", () => {
      setup("a");

      const { changes } = resolve(state, {
        type: "addSetupCellIfDoesntExist",
        payload: { code: "import marimo as mo" },
      });

      expect(changes).toMatchInlineSnapshot(`
        [
          {
            "before": "0",
            "cellId": "setup",
            "code": "import marimo as mo",
            "config": {
              "column": null,
              "disabled": false,
              "expand_output": false,
              "hide_code": false,
            },
            "name": "setup",
            "type": "create-cell",
          },
        ]
      `);
    });

    it("addSetupCellIfDoesntExist emits nothing when already exists", () => {
      setup("a");
      // Add setup cell first
      state = dispatch(state, {
        type: "addSetupCellIfDoesntExist",
        payload: { code: "import marimo as mo" },
      });

      // Try to add again — should just focus
      const { changes } = resolve(state, {
        type: "addSetupCellIfDoesntExist",
        payload: {},
      });

      expect(changes).toHaveLength(0);
    });
  });

  describe("actions that should NOT emit changes", () => {
    it("focusCell returns empty", () => {
      setup("a", "b");
      const [a] = state.cellIds.inOrderIds;
      const { changes } = resolve(state, {
        type: "focusCell",
        payload: { cellId: a, where: "after" },
      });
      expect(changes).toHaveLength(0);
    });

    it("prepareForRun returns empty", () => {
      setup("a");
      const [a] = state.cellIds.inOrderIds;
      const { changes } = resolve(state, {
        type: "prepareForRun",
        payload: { cellId: a },
      });
      expect(changes).toHaveLength(0);
    });
  });
});

describe("coalesceChanges", () => {
  const X = cellId("X");
  const Y = cellId("Y");
  const A = cellId("A");
  const config = { column: null, disabled: false, hide_code: false };

  it("drops set-code that precedes a delete of the same cell", () => {
    const changes: DocumentChange[] = [
      { type: "set-code", cellId: X, code: "x = 2" },
      { type: "delete-cell", cellId: X },
    ];
    expect(coalesceChanges(changes)).toEqual([
      { type: "delete-cell", cellId: X },
    ]);
  });

  it("keeps create+delete for a created-then-deleted cell, dropping the edit", () => {
    const changes: DocumentChange[] = [
      { type: "create-cell", cellId: X, code: "", name: "_", config },
      { type: "set-code", cellId: X, code: "x = 1" },
      { type: "delete-cell", cellId: X },
    ];
    expect(coalesceChanges(changes)).toEqual([
      { type: "create-cell", cellId: X, code: "", name: "_", config },
      { type: "delete-cell", cellId: X },
    ]);
  });

  it("drops set-code, set-name, set-config, and move-cell for a deleted cell", () => {
    const changes: DocumentChange[] = [
      { type: "create-cell", cellId: X, code: "", name: "_", config },
      { type: "set-code", cellId: X, code: "x = 1" },
      { type: "set-name", cellId: X, name: "foo" },
      {
        type: "set-config",
        cellId: X,
        column: null,
        disabled: true,
        hideCode: false,
        expandOutput: false,
      },
      { type: "move-cell", cellId: X, after: A },
      { type: "delete-cell", cellId: X },
    ];
    expect(coalesceChanges(changes)).toEqual([
      { type: "create-cell", cellId: X, code: "", name: "_", config },
      { type: "delete-cell", cellId: X },
    ]);
  });

  it("leaves a plain delete of an unedited cell unchanged", () => {
    const changes: DocumentChange[] = [{ type: "delete-cell", cellId: X }];
    expect(coalesceChanges(changes)).toEqual([
      { type: "delete-cell", cellId: X },
    ]);
  });

  it("does not touch set-* for cells that are not deleted, preserving order", () => {
    const changes: DocumentChange[] = [
      { type: "set-code", cellId: A, code: "a = 1" },
      { type: "set-code", cellId: Y, code: "y = 1" },
      { type: "delete-cell", cellId: X },
    ];
    expect(coalesceChanges(changes)).toEqual([
      { type: "set-code", cellId: A, code: "a = 1" },
      { type: "set-code", cellId: Y, code: "y = 1" },
      { type: "delete-cell", cellId: X },
    ]);
  });

  it("collapses repeated set-code for a surviving cell to the last value", () => {
    const changes: DocumentChange[] = [
      { type: "set-code", cellId: X, code: "x = 1" },
      { type: "set-name", cellId: X, name: "foo" },
      { type: "set-code", cellId: X, code: "x = 2" },
    ];
    expect(coalesceChanges(changes)).toEqual([
      { type: "set-name", cellId: X, name: "foo" },
      { type: "set-code", cellId: X, code: "x = 2" },
    ]);
  });

  it("keeps both creates so a downstream anchor stays resolvable", () => {
    const changes: DocumentChange[] = [
      { type: "create-cell", cellId: X, code: "", name: "_", config, after: A },
      { type: "create-cell", cellId: Y, code: "", name: "_", config, after: X },
      { type: "delete-cell", cellId: X },
    ];
    expect(coalesceChanges(changes)).toEqual([
      { type: "create-cell", cellId: X, code: "", name: "_", config, after: A },
      { type: "create-cell", cellId: Y, code: "", name: "_", config, after: X },
      { type: "delete-cell", cellId: X },
    ]);
  });
});

describe("document transaction middleware", () => {
  let sent: DocumentChange[][];

  const updateCode = (cellId: CellId, code: string) => {
    state = dispatch(state, {
      type: "updateCellCode",
      payload: { cellId, code, formattingChange: false },
    });
  };

  beforeEach(() => {
    vi.useFakeTimers();
    sent = [];
    middlewareExports.cancelPendingChanges();
    const captureClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        sent.push(changes);
        return null;
      },
    };
    store.set(
      requestClientAtom,
      captureClient as unknown as EditRequests & RunRequests,
    );
  });

  afterEach(() => {
    middlewareExports.cancelPendingChanges();
    store.set(requestClientAtom, null);
    vi.useRealTimers();
  });

  it("coalesces edit + delete of the same cell within one debounce window", async () => {
    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    // Discard the create-cell from setup; it would flush as its own batch.
    middlewareExports.cancelPendingChanges();

    updateCode(x, "x = 2");
    state = dispatch(state, {
      type: "deleteCell",
      payload: { cellId: x },
    });

    expect(sent).toEqual([]);

    await vi.advanceTimersByTimeAsync(400);

    expect(sent).toHaveLength(1);
    expect(sent[0]).toEqual([{ type: "delete-cell", cellId: x }]);
  });

  it("sends edit and delete as separate clean transactions across windows", async () => {
    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();

    updateCode(x, "x = 2");
    await vi.advanceTimersByTimeAsync(400);

    state = dispatch(state, {
      type: "deleteCell",
      payload: { cellId: x },
    });
    await vi.advanceTimersByTimeAsync(400);

    expect(sent).toHaveLength(2);
    expect(sent[0]).toEqual([{ type: "set-code", cellId: x, code: "x = 2" }]);
    expect(sent[1]).toEqual([{ type: "delete-cell", cellId: x }]);
  });

  it("flushes pending changes before a dependent operation", async () => {
    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");

    await flushDocumentChanges();

    expect(sent).toEqual([[{ type: "set-code", cellId: x, code: "x = 2" }]]);
  });

  it("keeps a failed transaction sticky before a dependent operation", async () => {
    const failingClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        sent.push(changes);
        throw new Error("connection lost");
      },
    };
    store.set(
      requestClientAtom,
      failingClient as unknown as EditRequests & RunRequests,
    );

    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");
    await vi.advanceTimersByTimeAsync(400);

    updateCode(x, "x = 3");
    await expect(flushDocumentChanges()).rejects.toThrow("connection lost");

    expect(sent).toEqual([[{ type: "set-code", cellId: x, code: "x = 2" }]]);
  });

  it("resumes after a full document save establishes a fresh baseline", async () => {
    let attempt = 0;
    const recoveringClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        sent.push(changes);
        attempt += 1;
        if (attempt === 1) {
          throw new Error("connection lost");
        }
        return null;
      },
    };
    store.set(
      requestClientAtom,
      recoveringClient as unknown as EditRequests & RunRequests,
    );

    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");
    await vi.advanceTimersByTimeAsync(400);

    updateCode(x, "x = 3");
    await withDocumentSave(async () => {
      updateCode(x, "x = 4");
    });

    await vi.advanceTimersByTimeAsync(400);
    expect(sent).toEqual([
      [{ type: "set-code", cellId: x, code: "x = 2" }],
      [{ type: "set-code", cellId: x, code: "x = 4" }],
    ]);
  });

  it("restores unsent edits when a full document save fails", async () => {
    const firstStarted = new Deferred<void>();
    const firstRelease = new Deferred<void>();
    const failingClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        sent.push(changes);
        if (sent.length === 1) {
          firstStarted.resolve();
          await firstRelease.promise;
          throw new Error("connection lost");
        }
        return null;
      },
    };
    store.set(
      requestClientAtom,
      failingClient as unknown as EditRequests & RunRequests,
    );

    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");
    vi.advanceTimersByTime(400);
    await firstStarted.promise;

    updateCode(x, "x = 3");
    await vi.advanceTimersByTimeAsync(400);
    firstRelease.resolve();
    await expect(flushDocumentChanges()).rejects.toThrow("connection lost");

    updateCode(x, "x = 4");
    const firstResync = beginDocumentResync();
    abortDocumentResync(firstResync);
    const secondResync = beginDocumentResync();
    expect(secondResync?.includedTransactions).toEqual([
      [{ type: "set-code", cellId: x, code: "x = 3" }],
    ]);
    expect(secondResync?.includedChanges).toEqual([
      { type: "set-code", cellId: x, code: "x = 4" },
    ]);
    await completeDocumentResync(secondResync);

    updateCode(x, "x = 5");
    await vi.advanceTimersByTimeAsync(400);
    expect(sent.at(-1)).toEqual([
      { type: "set-code", cellId: x, code: "x = 5" },
    ]);
  });

  it("keeps edits debounced while an earlier transaction is in flight", async () => {
    const firstRelease = new Deferred<void>();
    const firstStarted = new Deferred<void>();
    const delayedClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        sent.push(changes);
        if (sent.length === 1) {
          firstStarted.resolve();
          await firstRelease.promise;
        }
        return null;
      },
    };
    store.set(
      requestClientAtom,
      delayedClient as unknown as EditRequests & RunRequests,
    );

    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");
    vi.advanceTimersByTime(400);
    await firstStarted.promise;

    updateCode(x, "x = 3");
    firstRelease.resolve();
    await vi.advanceTimersByTimeAsync(399);
    expect(sent).toHaveLength(1);

    await vi.advanceTimersByTimeAsync(1);
    expect(sent).toEqual([
      [{ type: "set-code", cellId: x, code: "x = 2" }],
      [{ type: "set-code", cellId: x, code: "x = 3" }],
    ]);
  });

  it("waits for every transaction queued while flushing", async () => {
    const releases = Array.from({ length: 3 }, () => new Deferred<void>());
    const starts = Array.from({ length: 3 }, () => new Deferred<void>());
    const orderedClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        const index = sent.length;
        sent.push(changes);
        starts[index].resolve();
        await releases[index].promise;
        return null;
      },
    };
    store.set(
      requestClientAtom,
      orderedClient as unknown as EditRequests & RunRequests,
    );

    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");
    vi.advanceTimersByTime(400);
    await starts[0].promise;

    const flushed = flushDocumentChanges();

    updateCode(x, "x = 3");
    await vi.advanceTimersByTimeAsync(400);
    expect(sent).toHaveLength(1);

    releases[0].resolve();
    await starts[1].promise;

    updateCode(x, "x = 4");
    await vi.advanceTimersByTimeAsync(400);
    expect(sent).toHaveLength(2);

    let flushResolved = false;
    void flushed.then(() => {
      flushResolved = true;
    });
    releases[1].resolve();
    await starts[2].promise;
    expect(flushResolved).toBe(false);

    releases[2].resolve();
    await flushed;
    expect(sent).toEqual([
      [{ type: "set-code", cellId: x, code: "x = 2" }],
      [{ type: "set-code", cellId: x, code: "x = 3" }],
      [{ type: "set-code", cellId: x, code: "x = 4" }],
    ]);
  });

  it("orders a full save between earlier and later transactions", async () => {
    const firstStarted = new Deferred<void>();
    const firstRelease = new Deferred<void>();
    const saveStarted = new Deferred<void>();
    const saveRelease = new Deferred<void>();
    const events: string[] = [];
    const delayedClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        const change = changes[0];
        events.push("code" in change ? change.code : change.type);
        if (events.length === 1) {
          firstStarted.resolve();
          await firstRelease.promise;
        }
        return null;
      },
    };
    store.set(
      requestClientAtom,
      delayedClient as unknown as EditRequests & RunRequests,
    );

    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");
    vi.advanceTimersByTime(400);
    await firstStarted.promise;

    const saving = withDocumentSave(async () => {
      events.push("save");
      saveStarted.resolve();
      await saveRelease.promise;
    });
    expect(events).toEqual(["x = 2"]);

    firstRelease.resolve();
    await saveStarted.promise;
    updateCode(x, "x = 3");
    await vi.advanceTimersByTimeAsync(400);
    expect(events).toEqual(["x = 2", "save"]);

    saveRelease.resolve();
    await saving;
    await flushDocumentChanges();
    expect(events).toEqual(["x = 2", "save", "x = 3"]);
  });

  it("ignores a failure from a transaction reset while in flight", async () => {
    const firstStarted = new Deferred<void>();
    const firstRelease = new Deferred<void>();
    const firstFinished = new Deferred<void>();
    const delayedClient: Pick<
      EditRequests & RunRequests,
      "sendDocumentTransaction"
    > = {
      sendDocumentTransaction: async ({ changes }) => {
        sent.push(changes);
        if (sent.length === 1) {
          firstStarted.resolve();
          try {
            await firstRelease.promise;
          } finally {
            firstFinished.resolve();
          }
        }
        return null;
      },
    };
    store.set(
      requestClientAtom,
      delayedClient as unknown as EditRequests & RunRequests,
    );

    setup("x = 1");
    const [x] = state.cellIds.inOrderIds;
    middlewareExports.cancelPendingChanges();
    updateCode(x, "x = 2");
    vi.advanceTimersByTime(400);
    await firstStarted.promise;

    middlewareExports.cancelPendingChanges();
    firstRelease.reject(new Error("stale connection failure"));
    await firstFinished.promise;
    await Promise.resolve();

    updateCode(x, "x = 3");
    await flushDocumentChanges();
    expect(sent.at(-1)).toEqual([
      { type: "set-code", cellId: x, code: "x = 3" },
    ]);
  });
});