/// <reference lib="dom" />

import "./styles.css";
import { flushSync } from "react-dom";
import * as ReactDOM from "react-dom/client";
import styleText from "virtual:stylesheet";
import type * as vscode from "vscode-notebook-renderer";

import { assert, unreachable } from "../assert.ts";
import type { MarimoHtmlPublishMessage } from "../types.ts";
import { CellOutput } from "./CellOutput.tsx";
import {
  initImageTransport,
  resolveImageDataResult,
} from "./imageTransport.ts";
import {
  type CellId,
  type CellRuntimeState,
  handleFunctionCallResult,
  handleModelLifecycle,
  handleRemoveUIElements,
  handleSendUiElementMessage,
  initialize,
} from "./marimo-frontend.ts";
import { createRequestClient, isTypedRequestContext } from "./utils.ts";

const TABLE_EXPORT_LIMIT_MB = 50;
const TABLE_EXPORT_LIMIT_BYTES = TABLE_EXPORT_LIMIT_MB * 1024 * 1024;

export const activate: vscode.ActivationFunction = (context) => {
  assert(
    isTypedRequestContext(context),
    `Expected {"requiresMessaging": "always"} for marimo outputs.`,
  );

  initialize(createRequestClient(context), {
    limitBytes: TABLE_EXPORT_LIMIT_BYTES,
    unavailableMessage: `This table is too large to download from a VS Code notebook output (limit: ${TABLE_EXPORT_LIMIT_MB}MB). Open the notebook in a browser to export it, or filter the data first.`,
  });
  initImageTransport((message) => context.postMessage(message));

  /**
   * Bridge for HTML output interactions.
   *
   * Notebook cell output HTML runs in an isolated context and can't directly execute VSCode
   * commands or access the notebook editor. To enable interactive elements (like clickable
   * cell references in error messages), we:
   *
   * 1. Embed onclick handlers in the HTML that post messages to window.parent
   * 2. Catch those messages here in the renderer
   * 3. Forward them to the extension via context.postMessage()
   * 4. The extension handles them in NotebookRuntime.ts
   *
   * This enables features like "This variable was defined in [cell-2]" where clicking
   * "cell-2" navigates to that cell in the notebook.
   */
  window.addEventListener(
    "message",
    (event: MessageEvent<MarimoHtmlPublishMessage>) => {
      switch (event.data.command) {
        case "navigate-to-cell":
          context.postMessage(event.data);
          break;
        default:
          unreachable(
            event.data.command,
            `Unhandled HTML publish message: ${JSON.stringify(event.data)}`,
          );
      }
    },
  );

  context.onDidReceiveMessage((msg) => {
    switch (msg.op) {
      case "send-ui-element-message": {
        handleSendUiElementMessage(msg);
        return;
      }
      case "remove-ui-elements": {
        handleRemoveUIElements(msg);
        return;
      }
      case "function-call-result": {
        handleFunctionCallResult(msg);
        return;
      }
      case "model-lifecycle": {
        handleModelLifecycle(msg);
        return;
      }
      case "image-data-result": {
        resolveImageDataResult(msg.requestId, msg.dataUri);
        return;
      }
      default: {
        unreachable(
          msg,
          `Unhandled message from VS Code extension: ${JSON.stringify(msg)}`,
        );
      }
    }
  });

  // Inject the final compiled CSS from our Vite plugin
  // The title="marimo" tags these styles to be copied
  // into the ShadowDOM (for our UI elements).
  {
    const style = document.createElement("style");
    style.title = "marimo";
    style.textContent = styleText;
    document.head.appendChild(style);
  }

  const registry = new Map<HTMLElement, ReactDOM.Root>();

  return {
    renderOutputItem(data, element, signal) {
      const root = registry.get(element) ?? ReactDOM.createRoot(element);
      const { cellId, state }: { cellId: CellId; state: CellRuntimeState } =
        data.json();
      // Render synchronously so the output's real height is in the DOM before
      // VS Code measures `element` to reserve layout space. A plain async
      // `root.render` paints after `renderOutputItem` returns, so VS Code can
      // reserve too little height for tall outputs (e.g. tracebacks) and the
      // next cell overlaps them until a re-layout (#579, #598).
      flushSync(() => {
        root.render(<CellOutput cellId={cellId} state={state} />);
      });
      registry.set(element, root);
      signal.addEventListener("abort", () => {
        root.unmount();
        registry.delete(element);
      });
    },
  };
};
