import type { components as Api } from "@marimo-team/openapi/src/api";
import type * as lsp from "vscode-languageclient/node";

import type {
  NotebookCellId,
  NotebookId,
} from "./schemas/MarimoNotebookDocument.ts";
import type * as Gen from "./schemas/Models.gen.ts";

export type { CellRuntimeState } from "@marimo-team/frontend/unstable_internal/core/cells/types.ts";

type Schemas = Api["schemas"];
export type CellId = Schemas["CellId"];
export type CellConfig = Schemas["CellConfig"];
export type VariableName = Schemas["VariableName"];

export type Notification = Schemas["KnownUnions"]["notification"];
export type NotificationOf<T extends Notification["op"]> = Extract<
  Notification,
  { op: T }
>;
export type CellOperationNotification = NotificationOf<"cell-op">;
export type VariableValuesNotification = NotificationOf<"variable-values">;
export type VariablesNotification = NotificationOf<"variables">;
export type DataColumnPreviewNotification =
  NotificationOf<"data-column-preview">;
export type DataSourceConnectionsNotification =
  NotificationOf<"data-source-connections">;
export type DatasetsNotification = NotificationOf<"datasets">;
export type SqlSchemaListPreviewNotification =
  NotificationOf<"sql-schema-list-preview">;
export type SqlTablePreviewNotification = NotificationOf<"sql-table-preview">;
export type SqlTableListPreviewNotification =
  NotificationOf<"sql-table-list-preview">;
export type DataSourceConnection = Schemas["DataSourceConnection"];
export type DataTable = Schemas["DataTable"];
export type DatabaseSchema = Schemas["Schema"];
export type DocumentTransactionNotification =
  NotificationOf<"notebook-document-transaction">;
/** A single change within a document transaction (create/edit/move/... a cell). */
export type DocumentChange =
  DocumentTransactionNotification["transaction"]["changes"][number];

export type MarimoConfig = Gen.MarimoConfig;
export type SessionInfo = Gen.SessionInfo;
export type SessionsSnapshot = Gen.ListSessionsResponse;

/**
 * Wire shape of one `marimo.api` call. The generated client
 * (`schemas/Models.gen.ts`, from the `API_METHODS` registry in
 * `src/marimo_lsp/api.py`) owns the per-method payload/response schemas;
 * this is what they erase to at the transport.
 */
export type MarimoApiCall = Gen.MarimoApiCall;

/** The kernel-bound `inner` payload of a generated `marimo.api` method. */
type InnerOf<Payload> = Payload extends { readonly inner: infer Inner }
  ? Inner
  : never;

/**
 * Subset of API methods allowed to be dispatched by the renderer.
 *
 * The renderer handles two types of messages:
 *
 * 1. **From marimo UI components** (`update-ui-element`, `invoke-function`):
 *    - Interactive widgets like sliders, buttons, etc. send these when users interact with them
 *    - These are part of marimo's core reactive execution model
 *
 * 2. **From HTML output generated by the extension** (`navigate-to-cell`):
 *    - Error messages and other HTML output can include interactive elements
 *    - When users click links in error messages, the HTML posts messages that the renderer forwards
 *    - Example: "Variable 'x' was defined in [cell-2]" - clicking "cell-2" triggers navigate_to_cell
 *
 * All messages flow through the renderer (which runs in the notebook output iframe) to the
 * extension (which has access to the VSCode API and notebook editor).
 */
type RendererCommandMap = {
  // Forward to extension marimo.api
  "update-ui-element": InnerOf<typeof Gen.UpdateUiElementPayload.Encoded>;
  "set-model-value": InnerOf<typeof Gen.SetModelValuePayload.Encoded>;
  "invoke-function": InnerOf<typeof Gen.InvokeFunctionPayload.Encoded>;
  // Custom
  "navigate-to-cell": { cellId: NotebookCellId };
  // Image toolbar: the sandboxed renderer can't read cross-origin image bytes,
  // so it asks the host to fetch them. `save-image` writes to disk; `copy-image`
  // returns the bytes (as a data URI, keyed by `requestId`) for the renderer to
  // place on the clipboard.
  "save-image": { src: string; suggestedName: string };
  "copy-image": { src: string; requestId: string };
};
type RendererCommandMessageOf<K extends keyof RendererCommandMap> = {
  [C in keyof RendererCommandMap]: {
    command: C;
    params: RendererCommandMap[C];
  };
}[K];

// renderer -> extension
export type RendererCommand = RendererCommandMessageOf<
  keyof RendererCommandMap
>;

/**
 * Reply to a renderer `copy-image` request. `dataUri` is null when the host
 * could not fetch/decode the image, so the renderer can fall back gracefully.
 */
export type ImageDataResult = {
  op: "image-data-result";
  requestId: string;
  dataUri: string | null;
};

// extension -> renderer
export type RendererReceiveMessage =
  | NotificationOf<"remove-ui-elements">
  | NotificationOf<"send-ui-element-message">
  | NotificationOf<"function-call-result">
  | NotificationOf<"model-lifecycle">
  | ImageDataResult;

// Language server -> client
type MarimoLspNotificationMap = {
  "marimo/operation": { notebookUri: NotebookId; operation: Notification };
  "marimo/sessionsChanged": SessionsSnapshot;
  "window/logMessage": lsp.LogMessageParams;
};
export type MarimoLspNotification = keyof MarimoLspNotificationMap;
export type MarimoLspNotificationOf<K extends MarimoLspNotification> = {
  [C in MarimoLspNotification]: MarimoLspNotificationMap[C];
}[K];

export type MarimoOperation = MarimoLspNotificationOf<"marimo/operation">;
export type MarimoSessionsChanged =
  MarimoLspNotificationOf<"marimo/sessionsChanged">;

/**
 * Commands that HTML output can send to the renderer.
 *
 * The extension generates HTML output (e.g., error messages) that may include interactive
 * elements. Since this HTML runs in the renderer's iframe, it can't directly call VSCode APIs.
 * Instead, the HTML uses `window.parent.postMessage()` to send messages to the renderer,
 * which forwards them to the extension.
 *
 * Flow: HTML onclick → window.parent.postMessage() → renderer.tsx → NotebookRuntime.ts
 *
 * Example: Error messages include `<a onclick="postMessage({command: 'navigate_to_cell', ...})">cell-2</a>`
 *
 * To add new interactions: Add to RendererCommandMap and include in MarimoHtmlPublish.
 */
export type MarimoHtmlPublish = keyof Pick<
  RendererCommandMap,
  "navigate-to-cell"
>;
export type MarimoHtmlPublishOf<K extends MarimoHtmlPublish> =
  RendererCommandMessageOf<K>;
export type MarimoHtmlPublishMessage = MarimoHtmlPublishOf<MarimoHtmlPublish>;

/**
 * Type-safe helper to create a cell navigation link in HTML output.
 *
 * @param cellUri The URI of the cell to navigate to (e.g., "vscode-notebook-cell://...")
 * @param label The display text for the link (e.g., "cell-2")
 * @returns HTML string with embedded onclick handler
 */
export function createCellNavigationLink(
  cellId: NotebookCellId,
  cellIndex: number,
): string {
  const msg: MarimoHtmlPublishOf<"navigate-to-cell"> = {
    command: "navigate-to-cell",
    params: { cellId },
  };
  const encoded = JSON.stringify(msg).replace(/"/g, "&quot;");
  return `<a href="#" data-message="${encoded}" onclick="event.preventDefault(); window.parent.postMessage(JSON.parse(this.getAttribute('data-message')), '*'); return false;">cell-${cellIndex}</a>`;
}
