/* Copyright 2026 Marimo. All rights reserved. */
import { act, fireEvent, render, screen, within } from "@testing-library/react";
import { createStore, Provider } from "jotai";
import { expect, it, vi } from "vitest";
import { MockRequestClient } from "@/__mocks__/requests";
import { ModalProvider } from "@/components/modal/ImperativeModal";
import { TooltipProvider } from "@/components/ui/tooltip";
import { alertAtom } from "@/core/alerts/state";
import type { EnvironmentOperation } from "@/core/alerts/environment";
import { kernelStartupErrorAtom } from "@/core/errors/state";
import { connectionAtom, startupProgressAtom } from "@/core/network/connection";
import { requestClientAtom } from "@/core/network/requests";
import { sandboxAtom } from "@/core/packages/sandbox-state";
import { WebSocketClosedReason, ConnectionState } from "@/core/websocket/types";
import PackagesPanel from "../packages-panel";
import { PanelSectionProvider } from "../panel-context";
import { RuntimeStatusToggle } from "../runtime-status-toggle";

vi.mock("@/utils/reload-safe", () => ({ reloadSafe: vi.fn() }));

const preparation: EnvironmentOperation = {
  operation_id: "prepare-1",
  action: "prepare",
  source: "kernel",
  status: { kind: "succeeded" },
  packages: {},
  logs: {
    environment:
      "Using Python\nResolved dependencies\nDownloaded numpy\nInstalled numpy\n",
  },
};

function mount({
  preparing = false,
  backend = "uv",
}: { preparing?: boolean; backend?: "uv" | "pixi" | null } = {}) {
  const store = createStore();
  store.set(sandboxAtom, {
    backend,
    manifest: "dependencies = []",
    filename: "notebook.py",
  });
  store.set(
    connectionAtom,
    preparing
      ? { state: ConnectionState.CONNECTING, phase: "preparing-environment" }
      : { state: ConnectionState.OPEN },
  );
  store.set(alertAtom, (state) => ({
    ...state,
    environments: {
      ...state.environments,
      kernel: {
        restart_required: false,
        operations: backend ? [preparation] : [],
      },
    },
  }));
  if (backend) {
    store.set(startupProgressAtom, {
      phase: preparing ? "preparing-environment" : "starting-kernel",
      logs: preparing
        ? ""
        : "Loading runtime\nLaunching kernel\nKernel connected\n",
      log_mode: "replace",
    });
  }
  const client = MockRequestClient.create({
    sendRestart: vi.fn(async () => null),
    getDependencyTree: vi.fn(async () => ({
      context: backend
        ? { kind: "sandbox" as const, backend }
        : { kind: "package-manager" as const, name: "pip" },
      tree: {
        name: "<root>",
        version: null,
        tags: [],
        dependencies: [
          { name: "numpy", version: "2.0.0", tags: [], dependencies: [] },
        ],
      },
    })),
    getPackageList: vi.fn(async () => ({
      packages: [{ name: "numpy", version: "2.0.0" }],
    })),
  });
  store.set(requestClientAtom, client);
  const { unmount } = render(
    <Provider store={store}>
      <TooltipProvider>
        <ModalProvider>
          <RuntimeStatusToggle section="sidebar" />
          <PanelSectionProvider value="sidebar">
            <PackagesPanel />
          </PanelSectionProvider>
        </ModalProvider>
      </TooltipProvider>
    </Provider>,
  );
  return { store, client, unmount };
}

it("opens restored preparation and kernel logs without reopening completed startup on refresh", async () => {
  mount();
  await screen.findByText("numpy");
  const toggle = screen.getByRole("button", { name: /^uv sandbox, / });
  expect(toggle).toHaveAttribute("aria-expanded", "false");
  expect(
    screen.queryByRole("region", { name: "Sandbox details" }),
  ).not.toBeInTheDocument();
  fireEvent.click(toggle);
  const details = within(
    screen.getByRole("region", { name: "Sandbox details" }),
  );
  expect(details.getByText("Environment prepared")).toBeVisible();
  expect(details.getByText("Kernel started")).toBeVisible();
  fireEvent.click(
    details.getByRole("button", {
      name: "Expand environment preparation output",
    }),
  );
  fireEvent.click(
    details.getByRole("button", { name: "Expand kernel startup output" }),
  );
  expect(
    details.getByLabelText("environment preparation output").textContent,
  ).toBe(preparation.logs.environment);
  expect(details.getByLabelText("kernel startup output")).toHaveTextContent(
    "Loading runtime",
  );
  fireEvent.click(toggle);
  fireEvent.click(toggle);
  expect(
    details.getByLabelText("environment preparation output"),
  ).toBeVisible();
  expect(details.getByLabelText("kernel startup output")).toBeVisible();
});

it("restores a restart-required sync with the kernel restart action", async () => {
  const { store, client } = mount();
  act(() =>
    store.set(alertAtom, (state) => ({
      ...state,
      environments: {
        ...state.environments,
        kernel: {
          restart_required: true,
          operations: [
            preparation,
            {
              operation_id: "sync-python",
              action: "sync",
              source: "kernel",
              status: {
                kind: "restart-required",
                reason: "Python version changed",
              },
              packages: {},
              logs: { environment: "Resolved Python 3.14\n" },
            },
          ],
        },
      },
    })),
  );
  await screen.findByText("numpy");
  expect(
    screen.getByRole("heading", { name: "Sandbox restart required" }),
  ).toBeVisible();
  expect(
    screen.queryByRole("button", { name: "Retry sync" }),
  ).not.toBeInTheDocument();
  expect(screen.getByLabelText("Error details")).toHaveTextContent(
    "Python version changed",
  );
  // A later package operation can displace the sync, but not its restart.
  act(() =>
    store.set(alertAtom, (state) => ({
      ...state,
      environments: {
        ...state.environments,
        kernel: {
          restart_required: true,
          operations: [
            preparation,
            {
              operation_id: "install-1",
              action: "install",
              source: "kernel",
              status: { kind: "succeeded" },
              packages: { numpy: "succeeded" },
              logs: {},
            },
          ],
        },
      },
    })),
  );
  expect(
    screen.getByRole("heading", { name: "Sandbox restart required" }),
  ).toBeVisible();
  fireEvent.click(screen.getByRole("button", { name: "Restart Kernel" }));
  expect(screen.getByRole("alertdialog")).toHaveTextContent("Restart Kernel");
  expect(client.sendRestart).not.toHaveBeenCalled();
  fireEvent.click(screen.getByRole("button", { name: "Confirm Restart" }));
  expect(client.sendRestart).toHaveBeenCalledTimes(1);
});

it("keeps sync output open through completion and restores it quietly after refresh", () => {
  const { store, unmount } = mount();
  const restore = (status: EnvironmentOperation["status"], logs: string) =>
    store.set(alertAtom, (state) => ({
      ...state,
      environments: {
        ...state.environments,
        kernel: {
          restart_required: false,
          operations: [
            preparation,
            {
              operation_id: "sync-1",
              action: "sync",
              source: "kernel",
              status,
              packages: {},
              logs: { environment: logs },
            },
          ],
        },
      },
    }));
  act(() => restore({ kind: "running" }, "Resolving dependencies\n"));
  const preview = screen.getByRole("button", {
    name: "Expand sandbox sync output",
  });
  fireEvent.pointerDown(preview);
  fireEvent.click(preview);
  const output = screen.getByLabelText("sandbox sync output");
  act(() =>
    restore({ kind: "running" }, "Resolving dependencies\nDownloading numpy\n"),
  );
  expect(output).toHaveTextContent("Downloading numpy");
  act(() =>
    restore(
      { kind: "succeeded" },
      "Resolving dependencies\nDownloading numpy\nDone\n",
    ),
  );
  expect(output).toBeVisible();
  expect(output).toHaveTextContent("Done");
  expect(screen.getByText("Environment prepared")).toBeVisible();
  expect(screen.getByText("Kernel started")).toBeVisible();

  const snapshot = store.get(alertAtom);
  unmount();
  const refreshed = mount();
  act(() => refreshed.store.set(alertAtom, snapshot));
  const toggle = screen.getByRole("button", { name: /^uv sandbox, / });
  expect(toggle).toHaveAttribute("aria-expanded", "false");
  fireEvent.click(toggle);
  fireEvent.click(
    screen.getByRole("button", { name: "Expand sandbox sync output" }),
  );
  expect(screen.getByLabelText("sandbox sync output").textContent).toBe(
    "Resolving dependencies\nDownloading numpy\nDone\n",
  );
});

it.each([false, true])(
  "retains the progression after completion with inspecting=%s",
  async (inspecting) => {
    const { store, client } = mount({ preparing: true });
    const toggle = screen.getByRole("button", { name: /^uv sandbox, / });
    expect(toggle).toHaveAttribute("aria-expanded", "true");
    expect(client.getDependencyTree).not.toHaveBeenCalled();
    if (inspecting) {
      const preview = screen.getByRole("button", {
        name: "Expand environment preparation output",
      });
      fireEvent.pointerDown(preview);
      fireEvent.click(preview);
    }
    act(() => {
      store.set(startupProgressAtom, {
        phase: "starting-kernel",
        logs: "Launching kernel\n",
        log_mode: "replace",
      });
      store.set(connectionAtom, {
        state: ConnectionState.CONNECTING,
        phase: "starting-kernel",
      });
    });
    act(() => store.set(connectionAtom, { state: ConnectionState.OPEN }));
    await screen.findByText("numpy");
    expect(toggle).toHaveAttribute("aria-expanded", String(inspecting));
    if (!inspecting) {
      fireEvent.click(toggle);
    }
    expect(screen.getByText("Environment prepared")).toBeVisible();
    expect(screen.getByText("Kernel started")).toBeVisible();
    expect(
      screen.getByRole("button", {
        name: `${inspecting ? "Collapse" : "Expand"} environment preparation output`,
      }),
    ).toBeVisible();
    expect(
      screen.getByRole("button", { name: "Expand kernel startup output" }),
    ).toHaveTextContent("Launching kernel");
  },
);

it("opens a startup error even after collapsing preparation and keeps both step outputs available", () => {
  const { store } = mount({ preparing: true, backend: "pixi" });
  const toggle = screen.getByRole("button", { name: /^pixi sandbox, / });
  fireEvent.click(toggle);
  act(() => {
    store.set(startupProgressAtom, {
      phase: "starting-kernel",
      logs: "Loading runtime\nImport failed\n",
      log_mode: "replace",
    });
    store.set(kernelStartupErrorAtom, "Could not import the runtime");
    store.set(connectionAtom, {
      state: ConnectionState.CLOSED,
      phase: "starting-kernel",
      code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
      reason: "Kernel startup failed",
    });
  });
  expect(toggle).toHaveAttribute("aria-expanded", "true");
  expect(toggle).toHaveTextContent("pixi sandbox");
  expect(within(toggle).getByRole("status")).toHaveAccessibleName(
    "Kernel failed to start",
  );
  expect(screen.getByLabelText("Error details")).toHaveTextContent(
    "Could not import the runtime",
  );
  expect(screen.getByText("Environment prepared")).toBeVisible();
  expect(
    screen.getByRole("button", {
      name: "Expand environment preparation output",
    }),
  ).toHaveTextContent("Installed numpy");
  expect(
    screen.getByRole("button", { name: "Expand kernel startup output" }),
  ).toHaveTextContent("Import failed");
  fireEvent.click(toggle);
  expect(toggle).toHaveAttribute("aria-expanded", "false");
  fireEvent.click(toggle);
  expect(screen.getByLabelText("Error details")).toBeVisible();
});

it.each(["uv", "pixi"] as const)(
  "shows connection status and errors before the %s startup phase arrives",
  (backend) => {
    const { store } = mount({ backend, preparing: true });
    act(() => store.set(connectionAtom, { state: ConnectionState.CONNECTING }));
    const details = within(
      screen.getByRole("region", { name: "Sandbox details" }),
    );
    expect(details.getByRole("status")).toHaveTextContent("Connecting…");
    expect(details.queryByRole("list")).not.toBeInTheDocument();

    act(() => {
      store.set(kernelStartupErrorAtom, "Kernel process exited unexpectedly");
      store.set(connectionAtom, {
        state: ConnectionState.CLOSED,
        code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
        reason: "Kernel startup failed",
      });
    });
    expect(details.getByRole("status")).toHaveTextContent(
      "Kernel failed to start",
    );
    expect(details.getByLabelText("Error details")).toHaveTextContent(
      "Kernel process exited unexpectedly",
    );
    expect(details.queryByRole("list")).not.toBeInTheDocument();
  },
);

it("uses the same disclosure for an existing environment and updates it through connection changes", async () => {
  const { store } = mount({ backend: null });
  await screen.findByText("numpy");
  const toggle = screen.getByRole("button", { name: /^Runtime details, / });
  expect(toggle).toHaveTextContent(/^Runtime$/);
  expect(toggle).toHaveAccessibleName("Runtime details, Kernel ready");
  expect(within(toggle).getByRole("status")).toHaveAccessibleName(
    "Kernel ready",
  );
  expect(toggle).toHaveAttribute("aria-expanded", "false");
  fireEvent.click(toggle);
  const details = screen.getByRole("region", { name: "Runtime details" });
  expect(toggle).toHaveAttribute("aria-controls", details.id);
  expect(within(details).getByText("Kernel started")).toBeVisible();
  expect(within(details).getAllByRole("listitem")).toHaveLength(1);
  expect(within(details).queryByRole("button", { name: "Sync" })).toBeNull();

  act(() => store.set(connectionAtom, { state: ConnectionState.CONNECTING }));
  expect(toggle).toHaveTextContent(/^Runtime$/);
  expect(within(toggle).getByRole("status")).toHaveAccessibleName(
    "Connecting…",
  );
  expect(within(details).getByText("Starting kernel")).toBeVisible();
  fireEvent.click(toggle);
  expect(details).not.toBeVisible();
  act(() =>
    store.set(connectionAtom, {
      state: ConnectionState.CLOSED,
      code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
      reason: "Kernel startup failed",
    }),
  );
  expect(toggle).toHaveAttribute("aria-expanded", "true");
  expect(toggle).toHaveTextContent(/^Runtime$/);
  expect(within(toggle).getByRole("status")).toHaveAccessibleName(
    "Kernel failed to start",
  );
  expect(details).toBeVisible();
  expect(toggle).toHaveAccessibleName(
    "Runtime details, Kernel failed to start",
  );

  act(() => store.set(connectionAtom, { state: ConnectionState.NOT_STARTED }));
  expect(toggle).toHaveTextContent(/^Runtime$/);
  expect(within(toggle).getByRole("status")).toHaveAccessibleName(
    "Kernel not connected",
  );
});