import { defineCommand } from "citty";
import { ctx } from "../lib/context";
import {
  openExecSession,
  type ExecApiSession,
  type OpenExecOptions,
} from "../lib/exec-api";

export type ExecSessionOpener = typeof openExecSession;

function terminalSize(): { columns?: number; rows?: number } {
  return {
    columns: process.stdout.columns,
    rows: process.stdout.rows,
  };
}

export async function runExec(
  project: string,
  deployment: string,
  command: string[],
  signal: AbortSignal,
  opener: ExecSessionOpener = openExecSession,
  options?: OpenExecOptions,
): Promise<number> {
  if (!deployment) throw new Error("Deployment name is required");
  if (command.length === 0) throw new Error("Command is required");
  const tty = Boolean(process.stdin.isTTY && process.stdout.isTTY);
  const session: ExecApiSession = await opener(
    project,
    { deployment, command, tty, ...(tty ? terminalSize() : {}) },
    signal,
    options,
  );
  const onData = (chunk: Buffer | string) =>
    session.sendStdin(
      typeof chunk === "string" ? new TextEncoder().encode(chunk) : chunk,
    );
  const onEnd = () => session.sendStdin(new Uint8Array(), true);
  const onResize = () => {
    const { columns, rows } = terminalSize();
    if (columns && rows) session.resize(columns, rows);
  };
  process.stdin.on("data", onData);
  process.stdin.once("end", onEnd);
  if (tty) process.stdout.on("resize", onResize);
  try {
    for await (const frame of session) {
      if (frame.type === "stdout") process.stdout.write(frame.data);
      else if (frame.type === "stderr") process.stderr.write(frame.data);
      else if (frame.type === "error") throw new Error(frame.message);
      else return frame.exitCode;
    }
    if (signal.aborted) return 0;
    throw new Error("Exec connection closed without an exit status");
  } finally {
    process.stdin.off("data", onData);
    process.stdin.off("end", onEnd);
    process.stdout.off("resize", onResize);
    session.close();
  }
}

export const exec = defineCommand({
  meta: {
    name: "exec",
    description: "Execute a command inside a deployment pod",
  },
  async run({ args }) {
    const [deployment, ...command] = args._;
    if (!deployment) throw new Error("Deployment name is required");
    if (command.length === 0) throw new Error("Command is required");
    const controller = new AbortController();
    const abort = () => controller.abort();
    process.on("SIGINT", abort);
    process.on("SIGTERM", abort);
    const raw = Boolean(process.stdin.isTTY);
    if (raw) {
      process.stdin.setRawMode?.(true);
      process.stdin.resume();
    }
    try {
      const exitCode = await runExec(
        ctx().project,
        deployment,
        command,
        controller.signal,
      );
      if (exitCode !== 0) process.exitCode = exitCode;
    } catch (error) {
      if (!controller.signal.aborted) throw error;
    } finally {
      process.off("SIGINT", abort);
      process.off("SIGTERM", abort);
      if (raw) {
        process.stdin.setRawMode?.(false);
        process.stdin.pause();
      }
    }
  },
});