import {
  type DownloadSizeLimit,
  downloadSizeLimitAtom,
} from "@marimo-team/frontend/unstable_internal/components/data-table/download-policy/atoms.ts";
/**
 * @module
 *
 * Type-safe imports from @marimo-team/frontend
 *
 * See marimo-frontend-untyped.js for details on why this split exists.
 */
import type { CellId } from "@marimo-team/frontend/unstable_internal/core/cells/ids.ts";
import { FUNCTIONS_REGISTRY } from "@marimo-team/frontend/unstable_internal/core/functions/FunctionRegistry.ts";
import { initialModeAtom } from "@marimo-team/frontend/unstable_internal/core/mode.ts";
import { requestClientAtom } from "@marimo-team/frontend/unstable_internal/core/network/requests.ts";
/**
 * Type imports from @marimo-team/frontend
 *
 * These network types are imported WITHOUT ?nocheck because they're the most
 * likely to change as the API evolves. By keeping these type-checked, we ensure
 * our RequestClient interface stays in sync with marimo's actual API contracts.
 * Type errors here indicate breaking changes we need to handle.
 */
import type {
  EditRequests,
  RunRequests,
} from "@marimo-team/frontend/unstable_internal/core/network/types.ts";
import { store } from "@marimo-team/frontend/unstable_internal/core/state/jotai.ts";
import { safeExtractSetUIElementMessageBuffers } from "@marimo-team/frontend/unstable_internal/utils/json/base64.ts";

import type { NotificationOf, CellRuntimeState } from "../types.ts";

import "@marimo-team/frontend/unstable_internal/css/common.css";
import "@marimo-team/frontend/unstable_internal/css/globals.css";
import "@marimo-team/frontend/unstable_internal/css/codehilite.css";
import "@marimo-team/frontend/unstable_internal/css/katex.min.css";
import "@marimo-team/frontend/unstable_internal/css/md.css";
import "@marimo-team/frontend/unstable_internal/css/admonition.css";
import "@marimo-team/frontend/unstable_internal/css/md-tooltip.css";
import "@marimo-team/frontend/unstable_internal/css/table.css";
// @ts-expect-error - untyped JS bridge to avoid type-checking marimo internals (see marimo-frontend-untyped.js)
import * as untyped from "./marimo-frontend-untyped.js";

export { useTheme } from "@marimo-team/frontend/unstable_internal/theme/useTheme.ts";

export type RequestClient = EditRequests & RunRequests;
export type { CellRuntimeState, CellId, DownloadSizeLimit };

/**
 * Initialize marimo UI components in the VS Code renderer environment.
 * This provides a minimal setup to hydrate web components without the full kernel.
 *
 * `downloadSizeLimit`, when provided, gates table exports/clipboard copies that
 * exceed the host's limit. Applied here so all VS Code–specific store overrides
 * stay in one place and the raw jotai `store` never escapes this module.
 */
export function initialize(
  client: RequestClient,
  downloadSizeLimit?: DownloadSizeLimit,
) {
  store.set(initialModeAtom, "edit");
  store.set(requestClientAtom, client);
  // Trust data: URIs from plugins without waiting for a run through marimo's
  // React UI — VS Code gates notebook open on workspace trust, and cell runs
  // flow through the extension's NotebookController (bypassing useRunCells),
  // so the atom would otherwise stay false forever.
  //
  // TODO: route through the untyped bridge because marimo 0.23.2 defines this
  // atom in `components/editor/cell/useRunCells.ts`, whose transitive closure
  // (CodeMirror, cells state, data-table UI, etc.) blows up tsc. Once the atom
  // moves to a leaf file upstream (marimo-team/marimo#TODO), import directly.
  store.set(untyped.hasRunAnyCellAtom, true);

  if (downloadSizeLimit) {
    store.set(downloadSizeLimitAtom, downloadSizeLimit);
  }

  untyped.initializePlugins();
  // Start the RuntimeState to listen for UI element value changes
  // This connects the UI element events to the request client
  untyped.RuntimeState.INSTANCE.start(client.sendComponentValues);
}

// vendored from https://github.com/marimo-team/marimo/blob/111b24f/frontend/src/core/websocket/useMarimoWebSocket.tsx#L110-L134
export function handleSendUiElementMessage(
  msg: NotificationOf<"send-ui-element-message">,
) {
  const uiElement = msg.ui_element;
  const message = msg.message;
  const buffers = safeExtractSetUIElementMessageBuffers(msg);

  if (uiElement) {
    untyped.UI_ELEMENT_REGISTRY.broadcastMessage(uiElement, message, buffers);
  }
}

export function handleModelLifecycle(msg: NotificationOf<"model-lifecycle">) {
  // TODO: route through the untyped bridge because
  // `plugins/impl/anywidget/registry.ts` transitively imports `trusted-url.ts`,
  // which in marimo 0.23.2 pulls in `useRunCells.ts` and its heavy closure.
  // Drop once the upstream atom move lands (marimo-team/marimo#TODO).
  void untyped.handleWidgetMessage(untyped.WIDGET_REGISTRY, msg);
}

export function handleFunctionCallResult(
  msg: NotificationOf<"function-call-result">,
) {
  FUNCTIONS_REGISTRY.resolve(msg.function_call_id, msg);
  return;
}

export function handleRemoveUIElements(
  msg: NotificationOf<"remove-ui-elements">,
) {
  // This removes the element from the registry to (1) clean-up
  // memory and (2) make sure that the old value doesn't get re-used
  // if the same cell-id is later reused for another element.
  const cellId = msg.cell_id;
  untyped.UI_ELEMENT_REGISTRY.removeElementsByCell(cellId);
}

type OutputMessage = NonNullable<CellRuntimeState["output"]>;

export const OutputRenderer: React.FC<{
  message: OutputMessage;
  cellId?: CellId;
}> = untyped.OutputRenderer;

export const ConsoleOutput: React.FC<{
  cellId: CellId;
  cellName: string;
  consoleOutputs: Array<OutputMessage>;
  stale: boolean;
  debuggerActive: boolean;
  onSubmitDebugger: (text: string, index: number) => void;
}> = untyped.ConsoleOutput;

// SAFETY: `untyped` is an intentional escape hatch for marimo's frontend
// exports (see module docstring + marimo-frontend-untyped.js). Each re-export
// asserts the expected shape at this single boundary.
export const TooltipProvider: React.FC<
  React.PropsWithChildren<{ container: HTMLElement | null }>
  // oxlint-disable-next-line typescript/no-unsafe-type-assertion
> = untyped.TooltipProvider as React.FC<
  React.PropsWithChildren<{ container: HTMLElement | null }>
>;
