// Every run gets a sandbox: the only place its tools may read, write, or run
// commands. The local sandbox gives each run its own working directory and,
// where bubblewrap is available, runs commands isolated from the machine:
// no network, a read-only system, the user's home and /tmp hidden, and an
// empty environment. Cloud sandboxes such as Modal implement the same
// interface.

import { execFile, execFileSync } from "node:child_process";
import { existsSync } from "node:fs";
import { lstat, mkdir, readFile, realpath, rm, writeFile } from "node:fs/promises";
import path from "node:path";

export type ExecResult = {
  exitCode: number;
  stdout: string;
  stderr: string;
  timedOut: boolean;
};

export interface Sandbox {
  readonly kind: string;
  /**
   * True when commands cannot reach the network, the user's files, or the
   * host's environment. Tools that run model-written code should refuse to
   * run when this is false.
   */
  readonly isolated: boolean;
  /** Absolute working directory inside the sandbox. */
  readonly workdir: string;
  exec(command: string, args?: string[], options?: ExecOptions): Promise<ExecResult>;
  /** Write a file at a path relative to the working directory; it may not escape it. */
  writeFile(relativePath: string, content: string): Promise<void>;
  readFile(relativePath: string): Promise<string>;
  dispose(): Promise<void>;
}

export type ExecOptions = {
  timeoutMs?: number;
  signal?: AbortSignal;
  /**
   * Host paths the command may read, such as a language runtime installed
   * under the hidden home directory. Mounted read-only; nothing else of the
   * hidden directories is exposed.
   */
  readOnlyPaths?: string[];
  /** Extra environment variables. The host environment is never passed through. */
  env?: Record<string, string>;
};

export interface SandboxProvider {
  readonly kind: string;
  create(runId: string): Promise<Sandbox>;
}

/** For model calls that must not touch anything, such as a judge grading a run. */
export const noSandbox: Sandbox = {
  kind: "none",
  isolated: true,
  workdir: "",
  async exec() {
    throw new Error("This step has no sandbox; it cannot run commands.");
  },
  async writeFile() {
    throw new Error("This step has no sandbox; it cannot write files.");
  },
  async readFile() {
    throw new Error("This step has no sandbox; it cannot read files.");
  },
  async dispose() {},
};

/**
 * "bwrap" requires bubblewrap and fails if it cannot isolate; "none" runs
 * commands directly in the working directory (scoped, not isolated); "auto"
 * uses bubblewrap when it works here and falls back to "none".
 */
export type Isolation = "auto" | "bwrap" | "none";

const MAX_OUTPUT = 4 * 1024 * 1024;
const MAX_FILE = 1024 * 1024;

/** A working directory per run under `root`. Commands run there under a timeout, isolated when possible. */
export function localSandboxes(root: string, options: { keep?: boolean; isolation?: Isolation } = {}): SandboxProvider {
  const requested = options.isolation ?? "auto";
  const useBwrap = requested === "none" ? false : bwrapWorks();
  if (requested === "bwrap" && !useBwrap) {
    throw new Error("Sandbox isolation needs bubblewrap (bwrap) with unprivileged user namespaces, which is not working here.");
  }
  return {
    kind: useBwrap ? "local-bwrap" : "local",
    async create(runId) {
      const workdir = path.resolve(root, runId);
      await mkdir(workdir, { recursive: true, mode: 0o700 });
      const real = await realpath(workdir);
      return makeSandbox(real, useBwrap, async () => {
        if (!options.keep) await rm(real, { recursive: true, force: true });
      });
    },
  };
}

/**
 * A sandbox over an existing directory, such as the project an interactive
 * session works in. Commands are isolated as in localSandboxes (with this
 * directory writable and the rest of home hidden); disposing it leaves the
 * directory alone.
 */
export async function directorySandbox(dir: string, options: { isolation?: Isolation } = {}): Promise<Sandbox> {
  const requested = options.isolation ?? "auto";
  const useBwrap = requested === "none" ? false : bwrapWorks();
  if (requested === "bwrap" && !useBwrap) {
    throw new Error("Sandbox isolation needs bubblewrap (bwrap) with unprivileged user namespaces, which is not working here.");
  }
  return makeSandbox(await realpath(dir), useBwrap, async () => {});
}

function makeSandbox(real: string, useBwrap: boolean, dispose: () => Promise<void>): Sandbox {
  return {
    kind: useBwrap ? "local-bwrap" : "local",
    isolated: useBwrap,
    workdir: real,
    exec(command, args = [], { timeoutMs = 30_000, signal, readOnlyPaths = [], env = {} } = {}) {
      for (const p of readOnlyPaths) if (!path.isAbsolute(p)) throw new Error(`Read-only path "${p}" must be absolute.`);
      const [file, argv] = useBwrap ? ["bwrap", [...bwrapArgs(real, { readOnlyPaths, env }), "--", command, ...args]] : [command, args];
      return new Promise((resolve) => {
        execFile(
          file,
          argv,
          { cwd: real, env: { ...sandboxEnv(real), ...env }, timeout: timeoutMs, signal, maxBuffer: MAX_OUTPUT, killSignal: "SIGKILL" },
          (error, stdout, stderr) => {
            const err = error as (NodeJS.ErrnoException & { killed?: boolean; code?: number | string }) | null;
            resolve({
              exitCode: err ? (typeof err.code === "number" ? err.code : 1) : 0,
              stdout: String(stdout),
              stderr: err && !stderr ? err.message : String(stderr),
              timedOut: Boolean(err?.killed),
            });
          },
        );
      });
    },
    async writeFile(relativePath, content) {
      if (Buffer.byteLength(content) > MAX_FILE) throw new Error(`File is larger than ${MAX_FILE} bytes.`);
      const target = await containedPath(real, relativePath);
      await mkdir(path.dirname(target), { recursive: true });
      // Re-check after creating parents: a symlinked directory could point outside.
      await containedPath(real, relativePath);
      await writeFile(target, content);
    },
    async readFile(relativePath) {
      return readFile(await containedPath(real, relativePath), "utf8");
    },
    dispose,
  };
}

/** Commands see only a minimal environment: never the host's keys or tokens. */
function sandboxEnv(workdir: string): NodeJS.ProcessEnv {
  return { PATH: "/usr/local/bin:/usr/bin:/bin", HOME: workdir, TMPDIR: "/tmp", LANG: "C.UTF-8" };
}

/**
 * Read-only view of the system with a fresh /dev, /proc, and /tmp; the user's
 * home and other user-data mounts hidden; only the working directory
 * writable; no network or other shared namespaces; the process tree killed
 * with its parent.
 */
export function bwrapArgs(workdir: string, options: { readOnlyPaths?: string[]; env?: Record<string, string> } = {}): string[] {
  const hidden = ["/home", "/root", "/run/user", "/mnt", "/media", "/srv"].filter((p) => existsSync(p));
  return [
    "--ro-bind", "/", "/",
    "--dev", "/dev",
    "--proc", "/proc",
    "--tmpfs", "/tmp",
    ...hidden.flatMap((p) => ["--tmpfs", p]),
    ...(options.readOnlyPaths ?? []).flatMap((p) => ["--ro-bind", p, p]),
    "--bind", workdir, workdir,
    "--unshare-all",
    "--die-with-parent",
    "--new-session",
    "--clearenv",
    ...Object.entries({ ...sandboxEnv(workdir), ...options.env }).flatMap(([k, v]) => ["--setenv", k, v!]),
    "--chdir", workdir,
  ];
}

let bwrapProbe: boolean | undefined;
function bwrapWorks(): boolean {
  if (bwrapProbe !== undefined) return bwrapProbe;
  try {
    execFileSync("bwrap", ["--ro-bind", "/", "/", "--unshare-all", "--die-with-parent", "--", "true"], { stdio: "ignore", timeout: 5_000 });
    bwrapProbe = true;
  } catch {
    bwrapProbe = false;
  }
  return bwrapProbe;
}

async function containedPath(root: string, relativePath: string): Promise<string> {
  if (!relativePath || path.isAbsolute(relativePath)) throw new Error(`Use a path relative to the working directory, not "${relativePath}".`);
  const target = path.resolve(root, relativePath);
  if (target !== root && !target.startsWith(root + path.sep)) throw new Error(`Path "${relativePath}" leaves the working directory.`);
  // Resolve the deepest existing ancestor so a symlink cannot lead outside.
  let probe = target;
  while (!existsSync(probe)) probe = path.dirname(probe);
  const resolved = await realpath(probe);
  if (resolved !== root && !resolved.startsWith(root + path.sep)) throw new Error(`Path "${relativePath}" leaves the working directory.`);
  if (existsSync(target) && (await lstat(target)).isSymbolicLink()) throw new Error(`Path "${relativePath}" is a symlink.`);
  return target;
}