type CommandResult = { exitCode: number; output: string };

const launchScript = `
import os
import sys

pid_file, job_id, *command = sys.argv[1:]
if os.getpgrp() != os.getpid():
    os.setsid()
os.environ["SPOON_CODEX_JOB_ID"] = job_id
with open(pid_file, "w", encoding="utf-8") as handle:
    handle.write(str(os.getpid()))
    handle.flush()
    os.fsync(handle.fileno())
os.execvpe(command[0], command, os.environ)
`;

const abortScript = `
import os
import signal
import sys
import time

pid_file, job_id = sys.argv[1:]
pid = None
for _ in range(20):
    try:
        with open(pid_file, encoding="utf-8") as handle:
            pid = int(handle.read().strip())
        break
    except FileNotFoundError:
        time.sleep(0.1)
if pid is None:
    sys.exit(0)

try:
    with open(f"/proc/{pid}/environ", "rb") as handle:
        environment = handle.read().split(b"\\0")
    marker = f"SPOON_CODEX_JOB_ID={job_id}".encode()
    if marker not in environment or os.getpgid(pid) != pid:
        raise RuntimeError("Codex process identity did not match the requested job")
    os.killpg(pid, signal.SIGTERM)
    for _ in range(20):
        time.sleep(0.1)
        try:
            os.killpg(pid, 0)
        except ProcessLookupError:
            break
    else:
        os.killpg(pid, signal.SIGKILL)
except ProcessLookupError:
    pass
finally:
    try:
        os.unlink(pid_file)
    except FileNotFoundError:
        pass
`;

export const managedCodexCommand = (args: {
  jobId: string;
  pidFile: string;
  command: string[];
}) => [
  'python3',
  '-c',
  launchScript,
  args.pidFile,
  args.jobId,
  ...args.command,
];

export const abortManagedCodexProcess = async (args: {
  jobId: string;
  pidFile: string;
  execute: (command: string[]) => Promise<CommandResult>;
}) => {
  const result = await args.execute([
    'python3',
    '-c',
    abortScript,
    args.pidFile,
    args.jobId,
  ]);
  if (result.exitCode !== 0) {
    throw new Error(`Failed to abort Codex process: ${result.output}`);
  }
};
