import { DurableObject } from "cloudflare:workers";
import { containerLimits } from "./container-config";
import {
  abandonContainerSession,
  closeContainerSession,
  ContainerCapError,
  openContainerSession,
} from "./container-meter";
import { QuotaExceededError } from "../metering/quota";
import type { Reservation } from "../metering/quota";

const MINUTE_MS = 60 * 1000;
const CHECK_MS = MINUTE_MS;
const READY_ATTEMPTS = 60;
const READY_DELAY_MS = 500;
const ACTIVITY_WRITE_MS = 30 * 1000;
const MAX_SYNC_FILES = 300;
const DEFAULT_COLS = 120;
const DEFAULT_ROWS = 32;
const MAX_TERMINALS = 8;
const MAX_CONTROL_FRAME_BYTES = 1024;
const WORKSPACE = "/workspace";
const CONTAINER_HOME = "/home/dev";

/**
 * Writes (or deletes) one file under the checkout, refusing when any existing path component is a symlink:
 * the room is shared per branch but containers are per user, so a committed `link -> $HOME` must not let a
 * co-editor's room edit land in this user's home directory.
 */
const SYNC_SCRIPT = `set -e
root="$1"; rel="$2"; mode="$3"; dir="$root"
set -f; oldIFS=$IFS; IFS=/
for part in $rel; do
  if [ -L "$dir/$part" ]; then echo "symlink in path" >&2; exit 3; fi
  dir="$dir/$part"
done
IFS=$oldIFS; set +f
if [ "$mode" = "rm" ]; then rm -f "$root/$rel"; exit 0; fi
mkdir -p "$(dirname "$root/$rel")" && rm -f "$root/$rel" && cat > "$root/$rel"`;

type Meta = {
  accountId: string;
  repoId: string;
  branch: string;
  repoDir: string;
  /** The compute.editor hold this session runs against; its id is the editor_container_sessions row id. */
  reservation: Reservation;
  startedAt: number;
  lastActivity: number;
};

export type ContainerStart = {
  accountId: string;
  repoId: string;
  branch: string;
  repoName: string;
  cloneUrl: string;
  publicRemote: string;
  author: { name: string; email: string };
  /** Seconds to hold up front: the session cap, or less when the month's quota is nearly used. */
  budgetSeconds: number;
};

export type ContainerStartResult =
  | { ok: true; status: ContainerState }
  | { ok: false; code: "quota" | "cap" | "failed"; message: string };

export type ContainerState = {
  state: "stopped" | "running";
  startedAt: number | null;
  lastActivity: number | null;
  repoDir: string | null;
};

const encoder = new TextEncoder();

type Terminal = { process: ExecProcess; writer: WritableStreamDefaultWriter | null; pending: Promise<unknown> };

type TerminalControl = { type?: string; cols?: number; rows?: number };

function controlFrame(message: string | ArrayBuffer): TerminalControl | null {
  if (typeof message !== "string") return null;
  if (message.length > MAX_CONTROL_FRAME_BYTES) return {};
  try {
    return JSON.parse(message) as TerminalControl;
  } catch {
    return null;
  }
}

/** Binary frames are ArrayBuffers in workerd; dev proxies have been seen handing over Blobs. */
async function frameBytes(message: unknown): Promise<Uint8Array | null> {
  if (message instanceof ArrayBuffer) return new Uint8Array(message);
  if (message instanceof Blob) return new Uint8Array(await message.arrayBuffer());
  if (typeof message === "string") return encoder.encode(message);
  return null;
}

function safeRelative(path: string): string {
  const parts = path.split("/");
  if (path.startsWith("/") || parts.some((part) => part === ".." || part === "" || part === ".git")) {
    throw new Error(`Refusing to write ${path}`);
  }
  return path;
}

/**
 * One container per (account, repo, branch) for the editor's opt-in "Run in container" mode. It checks
 * out the branch read-only (commits still go through the forge, so protections apply), serves PTY
 * terminals over WebSockets, proxies previews of container ports, sleeps after the idle window and is
 * hard-stopped at the session cap. Each start/stop is metered in editor_container_sessions.
 */
export class EditorContainer extends DurableObject<Env> {
  private lastActivityWrite = 0;
  private starting: Promise<ContainerStartResult> | null = null;
  private readonly terminals = new Map<WebSocket, Terminal>();

  private async meta(): Promise<Meta | null> {
    return (await this.ctx.storage.get<Meta>("meta")) ?? null;
  }

  private get container(): Container {
    const container = this.ctx.container;
    if (!container) throw new Error("EditorContainer has no container binding");
    return container;
  }

  async status(): Promise<ContainerState> {
    const meta = await this.meta();
    if (!meta || !this.ctx.container?.running)
      return { state: "stopped", startedAt: null, lastActivity: null, repoDir: null };
    return { state: "running", startedAt: meta.startedAt, lastActivity: meta.lastActivity, repoDir: meta.repoDir };
  }

  /**
   * Boots (or reuses) the container and checks out the branch. The caller has authorized and checked
   * quota. Concurrent starts (a terminal and a task at once) share one boot, so only one session is metered.
   */
  async start(input: ContainerStart): Promise<ContainerStartResult> {
    // This object is keyed by (account, repo, branch), so concurrent callers always want the same container.
    this.starting ??= this.boot(input).finally(() => {
      this.starting = null;
    });
    return this.starting;
  }

  private async boot(input: ContainerStart): Promise<ContainerStartResult> {
    const existing = await this.meta();
    if (existing && this.container.running) {
      await this.touch(true);
      return { ok: true, status: await this.status() };
    }
    if (existing) await this.finish(existing, "restarted");
    let reservation: Reservation;
    try {
      reservation = await openContainerSession(input.accountId, input.repoId, input.branch, input.budgetSeconds);
    } catch (error) {
      if (error instanceof QuotaExceededError) return { ok: false, code: "quota", message: error.message };
      if (error instanceof ContainerCapError) return { ok: false, code: "cap", message: error.message };
      return { ok: false, code: "failed", message: error instanceof Error ? error.message : "Couldn't start" };
    }
    try {
      const repoDir = await this.bootContainer(input);
      const now = Date.now();
      await this.ctx.storage.put("meta", {
        accountId: input.accountId,
        repoId: input.repoId,
        branch: input.branch,
        repoDir,
        reservation,
        startedAt: now,
        lastActivity: now,
      } satisfies Meta);
      await this.ctx.storage.setAlarm(now + CHECK_MS);
    } catch (error) {
      console.error("Editor container boot failed", error);
      await this.ctx.container?.destroy("boot failed").catch(() => undefined);
      await abandonContainerSession(reservation).catch((inner: unknown) =>
        console.error("Couldn't release the container hold", inner),
      );
      return { ok: false, code: "failed", message: error instanceof Error ? error.message : "Couldn't start" };
    }
    return { ok: true, status: await this.status() };
  }

  /** Starts the container and checks the branch out; returns the checkout directory. */
  private async bootContainer(input: ContainerStart): Promise<string> {
    const container = this.container;
    if (!container.running) container.start({ enableInternet: true });
    const limits = containerLimits();
    await container.setInactivityTimeout(limits.capHours * 60 * MINUTE_MS);
    await this.waitReady();
    const repoDir = `${WORKSPACE}/${input.repoName.replace(/[^\w.-]/g, "_")}`;
    const clone = await container.exec(
      [
        "sh",
        "-c",
        'rm -rf "$1" && git clone --quiet --depth 50 --branch "$2" "$3" "$1" && cd "$1" && git remote set-url origin "$4" && git config user.name "$5" && git config user.email "$6"',
        "sh",
        repoDir,
        input.branch,
        input.cloneUrl,
        input.publicRemote,
        input.author.name,
        input.author.email,
      ],
      { stdout: "ignore", stderr: "pipe", env: { GIT_TERMINAL_PROMPT: "0" } },
    );
    const [code, stderr] = await Promise.all([clone.exitCode, clone.stderr ? new Response(clone.stderr).text() : ""]);
    if (code !== 0) {
      throw new Error(`Checkout failed: ${stderr.replaceAll(input.cloneUrl, "<remote>").slice(0, 300)}`);
    }
    return repoDir;
  }

  /** Stops the container and settles its usage; returns the seconds billed. */
  async stop(reason = "stopped"): Promise<number> {
    const meta = await this.meta();
    return meta ? this.finish(meta, reason) : 0;
  }

  /** Writes the room's pending edits into the checkout so terminals and tasks see what the editor shows. */
  async sync(changes: Array<{ path: string; content: string | null }>): Promise<number> {
    const meta = await this.meta();
    if (!meta || !this.container.running) throw new Error("The container isn't running");
    let written = 0;
    for (const change of changes.slice(0, MAX_SYNC_FILES)) {
      const relative = safeRelative(change.path);
      const command = ["sh", "-c", SYNC_SCRIPT, "sh", meta.repoDir, relative, change.content === null ? "rm" : "write"];
      const bytes = change.content === null ? null : encoder.encode(change.content);
      const process = await this.container.exec(command, {
        stdin: bytes
          ? new ReadableStream<Uint8Array>({
              start(controller) {
                controller.enqueue(bytes);
                controller.close();
              },
            })
          : undefined,
        stdout: "ignore",
        stderr: "ignore",
      });
      if ((await process.exitCode) === 0) written++;
    }
    await this.touch(true);
    return written;
  }

  /** Forwards an HTTP request to a port inside the container (port previews). */
  async proxy(port: number, request: Request): Promise<Response> {
    if (!(await this.meta()) || !this.container.running)
      return new Response("Container not running\n", { status: 503 });
    await this.touch(false);
    return this.container.getTcpPort(port).fetch(request);
  }

  /**
   * PTY terminal over a WebSocket: binary frames carry I/O, JSON text frames carry resize, exit and errors.
   * Sockets use the hibernation API, as the room does, so message types are what workerd delivers
   * (string or ArrayBuffer) rather than whatever the dev proxy wraps them in.
   */
  async fetch(request: Request): Promise<Response> {
    if (request.headers.get("upgrade") !== "websocket") return new Response("Expected websocket", { status: 426 });
    const meta = await this.meta();
    const url = new URL(request.url);
    const command = url.searchParams.get("command");
    const cols = Number(url.searchParams.get("cols")) || DEFAULT_COLS;
    const rows = Number(url.searchParams.get("rows")) || DEFAULT_ROWS;
    const pair = new WebSocketPair();
    const socket = pair[1];
    this.ctx.acceptWebSocket(socket);
    // Refusals are delivered over the accepted socket so the terminal can show them instead of a bare close code.
    if (!meta || !this.ctx.container?.running) {
      this.refuse(socket, "The container isn't running. Start it again from the status bar.");
    } else if (this.terminals.size >= MAX_TERMINALS) {
      this.refuse(socket, `At most ${MAX_TERMINALS} terminals can be open in one container.`);
    } else {
      void this.runTerminal(socket, meta, command, cols, rows);
    }
    const protocol = request.headers.get("x-editor-protocol");
    return new Response(null, {
      status: 101,
      webSocket: pair[0],
      headers: protocol ? { "sec-websocket-protocol": protocol } : {},
    });
  }

  private refuse(socket: WebSocket, message: string): void {
    socket.send(JSON.stringify({ type: "error", message }));
    socket.close(1011, "refused");
  }

  async webSocketMessage(socket: WebSocket, message: string | ArrayBuffer): Promise<void> {
    const terminal = this.terminals.get(socket);
    if (!terminal) return socket.close(1011, "terminal is gone");
    void this.touch(false);
    const control = controlFrame(message);
    if (control) {
      if (control.type === "resize" && control.cols && control.rows)
        terminal.process.resize(control.cols, control.rows);
      return;
    }
    const bytes = await frameBytes(message);
    if (!bytes || !terminal.writer) return;
    // Writes are chained so a slow stdin never reorders keystrokes.
    terminal.pending = terminal.pending.then(() => terminal.writer?.write(bytes)).catch(() => undefined);
    await terminal.pending;
  }

  webSocketClose(socket: WebSocket): void {
    this.terminals.get(socket)?.process.kill();
    this.terminals.delete(socket);
  }

  webSocketError(socket: WebSocket): void {
    this.webSocketClose(socket);
  }

  private async runTerminal(socket: WebSocket, meta: Meta, command: string | null, cols: number, rows: number) {
    try {
      const process = await this.container.exec(command ? ["bash", "-lc", command] : ["bash", "-l"], {
        cwd: meta.repoDir,
        env: { TERM: "xterm-256color", HOME: CONTAINER_HOME, GITFOR_BRANCH: meta.branch },
        pty: { cols, rows },
        stdin: "pipe",
        stdout: "pipe",
        stderr: "combined",
      });
      this.terminals.set(socket, { process, writer: process.stdin?.getWriter() ?? null, pending: Promise.resolve() });
      const reader = process.stdout?.getReader();
      if (reader) {
        for (;;) {
          const { done, value } = await reader.read();
          if (done) break;
          void this.touch(false);
          socket.send(value);
        }
      }
      const code = await process.exitCode;
      socket.send(JSON.stringify({ type: "exit", code }));
      socket.close(1000, "exited");
    } catch (error) {
      console.error("Editor container terminal failed", error);
      try {
        socket.send(
          JSON.stringify({ type: "error", message: error instanceof Error ? error.message : "Terminal failed" }),
        );
        socket.close(1011, "terminal failed");
      } catch {
        // socket already closed
      }
    } finally {
      this.terminals.delete(socket);
    }
  }

  async alarm(): Promise<void> {
    const meta = await this.meta();
    if (!meta) return;
    const limits = containerLimits();
    const now = Date.now();
    const elapsed = now - meta.startedAt;
    if (!this.container.running) {
      await this.finish(meta, "exited");
    } else if (elapsed >= limits.capHours * 60 * MINUTE_MS) {
      await this.finish(meta, "cap");
    } else if (elapsed >= meta.reservation.quantity * 1000) {
      await this.finish(meta, "quota");
    } else if (now - meta.lastActivity >= limits.idleMinutes * MINUTE_MS) {
      await this.finish(meta, "idle");
    } else {
      await this.ctx.storage.setAlarm(now + CHECK_MS);
    }
  }

  private async touch(force: boolean): Promise<void> {
    const now = Date.now();
    if (!force && now - this.lastActivityWrite < ACTIVITY_WRITE_MS) return;
    this.lastActivityWrite = now;
    const meta = await this.meta();
    if (meta) await this.ctx.storage.put("meta", { ...meta, lastActivity: now } satisfies Meta);
  }

  private async finish(meta: Meta, reason: string): Promise<number> {
    await this.ctx.storage.delete("meta");
    await this.ctx.storage.deleteAlarm();
    if (this.ctx.container?.running) await this.ctx.container.destroy(reason).catch(() => undefined);
    return closeContainerSession(meta.reservation, reason);
  }

  private async waitReady(): Promise<void> {
    for (let attempt = 0; attempt < READY_ATTEMPTS; attempt++) {
      try {
        const probe = await this.container.exec(["true"], { stdout: "ignore", stderr: "ignore" });
        if ((await probe.exitCode) === 0) return;
      } catch {
        // still booting
      }
      await new Promise((resolve) => setTimeout(resolve, READY_DELAY_MS));
    }
    throw new Error("The container didn't start in time");
  }
}