import { mkdirSync, readFileSync, rmdirSync, writeFileSync } from "node:fs";
import { isPidAlive, killPid, processTreePids, waitForPidsExit } from "agent-relay-sdk/process-utils";

export interface ExecResult {
  ok: boolean;
  exitCode: number | null;
  stdout: string;
  stderr: string;
  timedOut?: boolean;
  outputTruncated?: boolean;
  outputLimitExceeded?: boolean;
  /** Raw, undecoded stdout bytes — populated only when `stdoutBytes: true` was requested.
   * `stdout` is still populated too (best-effort UTF-8 decode, for logging), but callers that
   * need byte-exact paths (e.g. filenames containing invalid-UTF-8 bytes) must read this instead
   * of `stdout`, which is lossy by construction (`TextDecoder` replaces invalid sequences). */
  stdoutBytes?: Uint8Array;
}

interface ExecOptions {
  cwd?: string;
  env?: Record<string, string | undefined>;
  stdout?: "pipe" | "ignore";
  stderr?: "pipe" | "ignore";
  trimStdout?: boolean;
  trimStderr?: boolean;
  timeoutMs?: number;
  timeoutLabel?: string;
  streamDrainGraceMs?: number;
  maxOutputBytes?: number;
  signal?: AbortSignal;
  /** Write these raw bytes to the child's stdin, then close it. Opt-in; when omitted stdin is
   * "ignore" (the existing, unchanged default for every current caller). */
  stdin?: Uint8Array;
  /** Capture stdout as raw bytes (`stdoutBytes`) in addition to the lossily-decoded `stdout`
   * string. Opt-in and purely additive — every existing caller is unaffected. For callers that
   * need to round-trip exact filename bytes through git plumbing (e.g. `git diff -z` records
   * containing a non-UTF-8 path) the strings this module otherwise hands back have already lost
   * that information via `TextDecoder`; this is the only way to get it back. */
  stdoutBytes?: boolean;
  /** Reap ALL descendants on EVERY exit path — normal exit included — not just on timeout/abort
   * (#968 hole 2). Set for scheduler commands, which may background/daemonize children that would
   * otherwise outlive the parent shell. On Linux this confines the command to a dedicated cgroup v2
   * and kills the whole cgroup (bulletproof: a `setsid`-detached child cannot escape cgroup
   * membership); elsewhere it falls back to a detached process-group + descendant-tree kill. */
  reapProcessGroup?: boolean;
}

const DEFAULT_STREAM_DRAIN_GRACE_MS = 1_000;
const DEFAULT_KILL_GRACE_MS = 1_000;

let commandCgroupSeq = 0;

/** Create a dedicated child cgroup (v2) for a command so its entire process tree — including any
 * `setsid`-detached grandchildren that reparent away and escape the process group — can be reaped
 * in one shot via `cgroup.kill`. Returns the new cgroup's absolute path, or undefined when cgroup v2
 * isn't available/writable (non-Linux, cgroup v1, or missing delegation) so callers fall back. */
function createCommandCgroup(): string | undefined {
  if (process.platform !== "linux") return undefined;
  try {
    const self = readFileSync("/proc/self/cgroup", "utf8").trim();
    const idx = self.indexOf("::"); // cgroup v2 line: "0::/rel/path"
    if (idx < 0) return undefined;
    const rel = self.slice(idx + 2).trim();
    const path = `/sys/fs/cgroup${rel}/arcmd-${process.pid}-${Date.now()}-${commandCgroupSeq++}`;
    mkdirSync(path);
    return path;
  } catch {
    return undefined;
  }
}

function cgroupPids(cgPath: string): number[] {
  try {
    return readFileSync(`${cgPath}/cgroup.procs`, "utf8")
      .split("\n")
      .map((line) => Number(line.trim()))
      .filter((pid) => Number.isFinite(pid) && pid > 0);
  } catch {
    return [];
  }
}

/** Kill every process in the command's cgroup and remove the cgroup. `cgroup.kill` SIGKILLs the
 * whole subtree atomically; we then wait out any lingering PID-table entries and rmdir (retrying
 * a few times, never hanging) so we don't leak an empty cgroup dir. */
async function reapCommandCgroup(cgPath: string): Promise<void> {
  try { writeFileSync(`${cgPath}/cgroup.kill`, "1"); } catch {}
  await waitForPidsExit(cgroupPids(cgPath), DEFAULT_KILL_GRACE_MS);
  for (const pid of cgroupPids(cgPath)) killPid(pid, "SIGKILL");
  for (let attempt = 0; attempt < 5; attempt++) {
    try { rmdirSync(cgPath); return; } catch {}
    await sleep(50);
  }
}

/** Fallback reaper for hosts without cgroup v2: the command was spawned detached (its own process
 * group, pgid === proc.pid), so signal the whole group plus any still-living descendants. Reuses the
 * shared process helpers. NOTE: a `setsid`-detached child escapes the process group entirely — only
 * the cgroup path above reaps those; this is the best achievable without cgroups. */
async function reapDetachedProcessGroup(proc: Bun.Subprocess): Promise<void> {
  try { process.kill(-proc.pid, "SIGTERM"); } catch {}
  const pids = await processTreePids([proc.pid]).catch(() => [proc.pid]);
  for (const pid of pids) killPid(pid, "SIGTERM");
  const all = [proc.pid, ...pids];
  if (await waitForPidsExit(all, DEFAULT_KILL_GRACE_MS)) return;
  try { process.kill(-proc.pid, "SIGKILL"); } catch {}
  for (const pid of all.filter(isPidAlive)) killPid(pid, "SIGKILL");
}

interface StreamCapture {
  done: Promise<void>;
  text(): string;
  truncated(): boolean;
  limitExceeded(): boolean;
  cancel(): void;
}

function captureStream(
  stream: ReadableStream<Uint8Array> | undefined,
  options: { maxOutputBytes?: number; onLimitExceeded?: () => void } = {},
): StreamCapture {
  if (!stream) return { done: Promise.resolve(), text: () => "", truncated: () => false, limitExceeded: () => false, cancel: () => {} };
  const reader = stream.getReader();
  const decoder = new TextDecoder();
  let output = "";
  let canceled = false;
  let truncated = false;
  let limitExceeded = false;
  const done = (async () => {
    try {
      while (!canceled) {
        const chunk = await reader.read();
        if (chunk.done) break;
        output += decoder.decode(chunk.value, { stream: true });
        if (options.maxOutputBytes && output.length > options.maxOutputBytes) {
          truncated = true;
          output = output.slice(output.length - options.maxOutputBytes);
          if (!limitExceeded) {
            limitExceeded = true;
            options.onLimitExceeded?.();
          }
        }
      }
      output += decoder.decode();
      if (options.maxOutputBytes && output.length > options.maxOutputBytes) {
        truncated = true;
        output = output.slice(output.length - options.maxOutputBytes);
      }
    } catch {
      // Intentional cancellation on process timeout or a stuck post-exit pipe.
    } finally {
      try { reader.releaseLock(); } catch {}
    }
  })();
  return {
    done,
    text: () => output,
    truncated: () => truncated,
    limitExceeded: () => limitExceeded,
    cancel: () => {
      canceled = true;
      void reader.cancel().catch(() => {});
    },
  };
}

interface BytesCapture {
  done: Promise<void>;
  bytes(): Uint8Array;
  cancel(): void;
}

/** Byte-preserving sibling of {@link captureStream} — accumulates raw chunks with no
 * `TextDecoder` pass, so filename bytes that aren't valid UTF-8 survive intact. No truncation
 * support: only used for small, bounded plumbing output (git diff/check-ignore records for a
 * single land), never for arbitrary command stdout. */
function captureStreamBytes(stream: ReadableStream<Uint8Array> | undefined): BytesCapture {
  if (!stream) return { done: Promise.resolve(), bytes: () => new Uint8Array(0), cancel: () => {} };
  const reader = stream.getReader();
  const chunks: Uint8Array[] = [];
  let canceled = false;
  const done = (async () => {
    try {
      while (!canceled) {
        const chunk = await reader.read();
        if (chunk.done) break;
        chunks.push(chunk.value);
      }
    } catch {
      // Intentional cancellation on process timeout or a stuck post-exit pipe.
    } finally {
      try { reader.releaseLock(); } catch {}
    }
  })();
  return {
    done,
    bytes: () => {
      const total = chunks.reduce((n, c) => n + c.length, 0);
      const out = new Uint8Array(total);
      let offset = 0;
      for (const c of chunks) { out.set(c, offset); offset += c.length; }
      return out;
    },
    cancel: () => {
      canceled = true;
      void reader.cancel().catch(() => {});
    },
  };
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function timeoutMessage(cmd: string[], options: ExecOptions): string {
  const label = options.timeoutLabel ?? cmd.join(" ");
  return `${label} timed out after ${options.timeoutMs}ms`;
}

function abortMessage(cmd: string[], options: ExecOptions): string {
  const reason = options.signal?.reason;
  if (reason instanceof Error && reason.message) return reason.message;
  if (typeof reason === "string" && reason) return reason;
  const label = options.timeoutLabel ?? cmd.join(" ");
  return `${label} aborted`;
}

async function terminateProcessTree(proc: Bun.Subprocess): Promise<void> {
  const pids = await processTreePids([proc.pid]).catch(() => [proc.pid]);
  for (const target of pids) killPid(target, "SIGTERM");
  try { proc.kill("SIGTERM"); } catch {}
  const exited = await waitForPidsExit(pids, DEFAULT_KILL_GRACE_MS);
  const alive = exited ? [] : pids.filter(isPidAlive);
  for (const target of alive) killPid(target, "SIGKILL");
  if (alive.includes(proc.pid)) {
    try { proc.kill("SIGKILL"); } catch {}
  }
  if (alive.length > 0) await waitForPidsExit(alive, DEFAULT_KILL_GRACE_MS);
}

export async function execProcess(cmd: string[], options: ExecOptions = {}): Promise<ExecResult> {
  if (options.signal?.aborted) {
    return { ok: false, exitCode: null, stdout: "", stderr: abortMessage(cmd, options) };
  }
  // #968 hole 2 — for reap-on-every-exit commands, confine the whole tree to a dedicated cgroup so a
  // `setsid`-detached child can't survive the parent. The wrapper joins the cgroup (writes its own
  // pid to cgroup.procs) BEFORE exec-ing the real command, so every descendant is a member from birth
  // (race-free); `exec "$0" "$@"` passes the original argv through untouched.
  const cgroup = options.reapProcessGroup ? createCommandCgroup() : undefined;
  const spawnCmd = cgroup
    ? ["sh", "-c", `echo $$ > ${cgroup}/cgroup.procs 2>/dev/null || true; exec "$0" "$@"`, ...cmd]
    : cmd;
  const proc = Bun.spawn(spawnCmd, {
    cwd: options.cwd,
    env: options.env,
    stdin: options.stdin !== undefined ? "pipe" : "ignore",
    stdout: options.stdout ?? "pipe",
    stderr: options.stderr ?? "pipe",
    // Own process group so the fallback reaper can group-signal without touching the orchestrator.
    ...(options.reapProcessGroup ? { detached: true } : {}),
  });
  if (options.stdin !== undefined) {
    // Bun's Subprocess<stdin> type is only narrowed for a statically-literal `stdin: "pipe"`
    // spawn option; ours is chosen at runtime above, so TS can't see stdin is definitely a
    // FileSink here even though it always is when options.stdin was provided.
    const stdin = proc.stdin as NonNullable<typeof proc.stdin>;
    stdin.write(options.stdin);
    await stdin.end();
  }
  let timedOut = false;
  let aborted = false;
  let outputLimitExceeded = false;
  let abortStarted = false;
  let timeout: ReturnType<typeof setTimeout> | undefined;
  let timeoutResolve: ((value: null) => void) | undefined;
  const timeoutPromise = new Promise<null>((resolve) => { timeoutResolve = resolve; });
  let abortDone: Promise<void> = Promise.resolve();
  const abortProcess = (killParentImmediately = false) => {
    if (abortStarted) return;
    abortStarted = true;
    aborted = true;
    if (killParentImmediately) {
      try { proc.kill("SIGTERM"); } catch {}
    }
    abortDone = terminateProcessTree(proc).finally(() => timeoutResolve?.(null));
  };
  const abortForOutputLimit = () => {
    outputLimitExceeded = true;
    abortProcess();
  };
  // `stdoutBytes` reads proc.stdout itself (a ReadableStream only one reader can consume), so it's
  // mutually exclusive with the string capture below — the string `stdout` field is then derived
  // from the captured bytes after the fact instead of a second independent read.
  const stdoutBytesCapture = options.stdoutBytes && options.stdout !== "ignore" ? captureStreamBytes(proc.stdout) : undefined;
  const stdoutCapture = stdoutBytesCapture
    ? undefined
    : options.stdout === "ignore"
      ? captureStream(undefined)
      : captureStream(proc.stdout, { maxOutputBytes: options.maxOutputBytes, onLimitExceeded: abortForOutputLimit });
  const stderrCapture = options.stderr === "ignore"
    ? captureStream(undefined)
    : captureStream(proc.stderr, { maxOutputBytes: options.maxOutputBytes, onLimitExceeded: abortForOutputLimit });
  const abortFromSignal = () => abortProcess(true);
  options.signal?.addEventListener("abort", abortFromSignal, { once: true });
  if (options.timeoutMs && options.timeoutMs > 0) {
    timeout = setTimeout(() => {
      timedOut = true;
      abortProcess();
    }, options.timeoutMs);
    timeout.unref?.();
  }

  const exitCode = await (timeout || options.signal ? Promise.race([proc.exited, timeoutPromise]) : proc.exited);
  options.signal?.removeEventListener("abort", abortFromSignal);
  if (timeout) clearTimeout(timeout);
  if (timedOut || aborted) {
    await abortDone;
    (stdoutBytesCapture ?? stdoutCapture)!.cancel();
    stderrCapture.cancel();
  }

  const drainGraceMs = options.streamDrainGraceMs ?? DEFAULT_STREAM_DRAIN_GRACE_MS;
  await Promise.race([Promise.allSettled([(stdoutBytesCapture ?? stdoutCapture)!.done, stderrCapture.done]), sleep(drainGraceMs)]);
  (stdoutBytesCapture ?? stdoutCapture)!.cancel();
  stderrCapture.cancel();

  // #968 hole 2 — reap the command's whole process tree on EVERY exit path (normal exit included),
  // so a backgrounded/daemonized child can't outlive the command. Runs regardless of exit reason.
  if (options.reapProcessGroup) {
    if (cgroup) await reapCommandCgroup(cgroup);
    else await reapDetachedProcessGroup(proc);
  }

  const stdoutBytes = stdoutBytesCapture?.bytes();
  const stdout = stdoutBytes !== undefined ? new TextDecoder().decode(stdoutBytes) : stdoutCapture!.text();
  let stderr = stderrCapture.text();
  if (timedOut) {
    const msg = timeoutMessage(cmd, options);
    stderr = stderr ? `${stderr}\n${msg}` : msg;
  } else if (outputLimitExceeded) {
    const msg = `output exceeded ${options.maxOutputBytes} bytes`;
    stderr = stderr ? `${stderr}\n${msg}` : msg;
  } else if (aborted) {
    const msg = abortMessage(cmd, options);
    stderr = stderr ? `${stderr}\n${msg}` : msg;
  }
  const reportedExitCode = timedOut || aborted || outputLimitExceeded ? null : exitCode;
  return {
    ok: !timedOut && !aborted && !outputLimitExceeded && reportedExitCode === 0,
    exitCode: reportedExitCode,
    stdout: options.trimStdout === false ? stdout : stdout.trim(),
    stderr: options.trimStderr === false ? stderr : stderr.trim(),
    ...(timedOut ? { timedOut } : {}),
    ...(stdoutBytes !== undefined ? { stdoutBytes } : {}),
    ...((stdoutCapture?.truncated() ?? false) || stderrCapture.truncated() ? { outputTruncated: true } : {}),
    ...(outputLimitExceeded ? { outputLimitExceeded: true } : {}),
  };
}