import * as vscode from "vscode";
import { EDITOR_WS_PROTOCOL } from "../lib/editor-protocol";
import type { ForgeApi } from "./forge";
import type { RoomClient } from "./room-client";

type ContainerStatus = {
  state: "stopped" | "starting" | "running";
  allowed: boolean;
  startedAt: number | null;
  idleMinutes: number;
  capHours: number;
  minutesUsed: number;
  quotaMinutes: number;
};

const NORMAL_CLOSE = 1000;
const MAX_PORT = 65535;
const encoder = new TextEncoder();
const decoder = new TextDecoder();

/** A Pseudoterminal backed by a process in the session's container, streamed over a WebSocket. */
function containerPty(api: ForgeApi, command: string | null): vscode.Pseudoterminal {
  const write = new vscode.EventEmitter<string>();
  const close = new vscode.EventEmitter<number>();
  let socket: WebSocket | null = null;
  return {
    onDidWrite: write.event,
    onDidClose: close.event,
    open(dimensions) {
      const url = new URL(api.roomUrl().replace(/\/room$/, "/container/terminal"));
      if (command) url.searchParams.set("command", command);
      if (dimensions) {
        url.searchParams.set("cols", String(dimensions.columns));
        url.searchParams.set("rows", String(dimensions.rows));
      }
      socket = new WebSocket(url, [EDITOR_WS_PROTOCOL, api.token]);
      socket.binaryType = "arraybuffer";
      socket.onmessage = (event) => {
        if (typeof event.data === "string") {
          const message = JSON.parse(event.data) as { type: string; code?: number; message?: string };
          if (message.type === "exit") close.fire(message.code ?? 0);
          if (message.type === "error") write.fire(`\r\n\x1b[31m${message.message ?? "error"}\x1b[0m\r\n`);
          return;
        }
        write.fire(
          decoder.decode(new Uint8Array(event.data as ArrayBuffer), { stream: true }).replace(/(?<!\r)\n/g, "\r\n"),
        );
      };
      socket.onclose = (event) => {
        if (event.code !== NORMAL_CLOSE)
          write.fire(`\r\n\x1b[33m[container connection closed: ${event.reason || event.code}]\x1b[0m\r\n`);
        close.fire(0);
      };
    },
    handleInput(data) {
      if (socket?.readyState === WebSocket.OPEN) socket.send(encoder.encode(data));
    },
    setDimensions(dimensions) {
      if (socket?.readyState === WebSocket.OPEN) {
        socket.send(JSON.stringify({ type: "resize", cols: dimensions.columns, rows: dimensions.rows }));
      }
    },
    close() {
      socket?.close(NORMAL_CLOSE);
    },
  };
}

/**
 * "Run in container": opt-in compute for the session. Terminals, tasks (npm scripts and `gitfor` tasks)
 * and port previews run in a per-(account, repo, branch) container that sleeps when idle and is capped.
 * Live room edits are synced into the container before each terminal or task starts.
 */
export function registerContainer(
  context: vscode.ExtensionContext,
  api: ForgeApi,
  room: RoomClient,
  startNow: boolean,
): void {
  const item = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 90);
  item.command = "gitfor.container.start";
  item.text = "$(vm) Run in container";
  item.tooltip = "Start a container with a terminal, tasks and ports for this branch (metered)";

  // The plan gate is re-read on every refresh, so an upgrade or a lapsed plan is reflected without a reload.
  const refresh = async () => {
    const status = await api.request<ContainerStatus>("container/status").catch(() => null);
    if (!status) return;
    if (!status.allowed) return item.hide();
    item.show();
    item.text = status.state === "running" ? "$(vm-active) Container running" : "$(vm) Run in container";
    item.command = status.state === "running" ? "gitfor.container.menu" : "gitfor.container.start";
    item.tooltip = `${status.minutesUsed}/${status.quotaMinutes} container minutes used this month; sleeps after ${status.idleMinutes} idle minutes, stops after ${status.capHours}h`;
  };

  const start = async (): Promise<boolean> => {
    try {
      await vscode.window.withProgress(
        { location: vscode.ProgressLocation.Notification, title: "Starting container" },
        () => api.request<ContainerStatus>("container/start", { method: "POST", body: {} }),
      );
      await refresh();
      return true;
    } catch (error) {
      void vscode.window.showErrorMessage(error instanceof Error ? error.message : "Couldn't start the container");
      return false;
    }
  };

  const sync = () => api.request("container/sync", { method: "POST", body: {} }).catch(() => undefined);

  const openTerminal = async (command: string | null, name: string) => {
    if (!room.writable) return void vscode.window.showWarningMessage("Containers need write access to this branch.");
    if (!(await start())) return;
    await sync();
    const terminal = vscode.window.createTerminal({ name, pty: containerPty(api, command) });
    terminal.show();
  };

  const taskProvider: vscode.TaskProvider = {
    async provideTasks() {
      const tasks: vscode.Task[] = [];
      const status = await api.request<ContainerStatus>("container/status").catch(() => null);
      if (!status?.allowed) return tasks;
      try {
        const manifest = JSON.parse(
          room.text("package.json").toString() || (await api.blob("package.json")).content || "{}",
        );
        for (const script of Object.keys(manifest.scripts ?? {})) {
          tasks.push(containerTask(api, script, `npm run ${script}`));
        }
      } catch {
        return tasks;
      }
      return tasks;
    },
    resolveTask(task) {
      const command = typeof task.definition.command === "string" ? task.definition.command : null;
      return command ? containerTask(api, task.name, command) : undefined;
    },
  };

  context.subscriptions.push(
    item,
    vscode.commands.registerCommand("gitfor.container.start", () => openTerminal(null, "container")),
    vscode.commands.registerCommand("gitfor.container.terminal", () => openTerminal(null, "container")),
    vscode.commands.registerCommand("gitfor.container.stop", async () => {
      await api.request("container/stop", { method: "POST", body: {} }).catch(() => undefined);
      await refresh();
    }),
    vscode.commands.registerCommand("gitfor.container.sync", sync),
    vscode.commands.registerCommand("gitfor.container.openPort", async () => {
      const value = await vscode.window.showInputBox({ prompt: "Port to preview", placeHolder: "3000" });
      const port = Number(value);
      if (!Number.isInteger(port) || port <= 0 || port > MAX_PORT) return;
      const preview = await api.request<{ url: string }>("container/preview", { method: "POST", body: { port } });
      await vscode.env.openExternal(vscode.Uri.parse(preview.url));
    }),
    vscode.commands.registerCommand("gitfor.container.menu", async () => {
      const pick = await vscode.window.showQuickPick([
        "New terminal",
        "Preview a port",
        "Sync edits to container",
        "Stop container",
      ]);
      const commands: Record<string, string> = {
        "New terminal": "gitfor.container.terminal",
        "Preview a port": "gitfor.container.openPort",
        "Sync edits to container": "gitfor.container.sync",
        "Stop container": "gitfor.container.stop",
      };
      if (pick) await vscode.commands.executeCommand(commands[pick]);
    }),
    vscode.tasks.registerTaskProvider("gitfor", taskProvider),
  );
  void refresh().then(() => {
    if (startNow) return vscode.commands.executeCommand("gitfor.container.start");
  });
}

function containerTask(api: ForgeApi, name: string, command: string): vscode.Task {
  return new vscode.Task(
    { type: "gitfor", command },
    vscode.TaskScope.Workspace,
    name,
    "container",
    new vscode.CustomExecution(async () => {
      await api.request("container/start", { method: "POST", body: {} });
      await api.request("container/sync", { method: "POST", body: {} }).catch(() => undefined);
      return containerPty(api, command);
    }),
  );
}