/** Shared utilities for spawning isolated prompt-native Pi subprocesses. */

import { spawn, type ChildProcess } from "node:child_process";
import * as fs from "node:fs";
import * as path from "node:path";
import type { Message } from "@earendil-works/pi-ai";
import type { BackendName, UsageStats } from "./agent-runner-types.js";
import { CHILD_ORCHESTRATION_TOOL_NAMES } from "./leaf-policy.js";

export function emptyUsage(): UsageStats {
  return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 };
}

function getPiInvocation(args: string[]): { command: string; args: string[] } {
  const currentScript = process.argv[1];
  const isBunVirtualPath = currentScript?.startsWith("/$bunfs/");

  if (currentScript && !isBunVirtualPath && fs.existsSync(currentScript)) {
    return { command: process.execPath, args: [currentScript, ...args] };
  }

  const execName = path.basename(process.execPath).toLowerCase();
  if (!/^(node|bun)(\.exe)?$/.test(execName)) return { command: process.execPath, args };
  return { command: "pi", args };
}

export interface ToolExecutionStartEvent {
  toolCallId: string;
  toolName: string;
  args: Record<string, unknown>;
}

export interface SpawnPiAgentOptions {
  cwd: string;
  prompt: string;
  label?: string;
  model?: string;
  thinking?: string;
  tools?: string[];
  /** Agent runtime to spawn. Defaults to `pi`. */
  backend?: BackendName;
  env?: NodeJS.ProcessEnv;
  signal?: AbortSignal;
  onMessage?: (msg: Message) => void;
  onToolResult?: (msg: Message) => void;
  onToolExecutionStart?: (event: ToolExecutionStartEvent) => void;
  /** Injectable only for lifecycle tests. */
  terminationGraceMs?: number;
  /** Injectable only for lifecycle tests. */
  terminationForceWaitMs?: number;
}

export interface SpawnPiAgentResult {
  exitCode: number;
  messages: Message[];
  stderr: string;
  wasAborted: boolean;
  usage: UsageStats;
  model?: string;
  stopReason?: string;
  errorMessage?: string;
}

export const PROCESS_TREE_GRACE_MS = 5_000;
export const PROCESS_TREE_FORCE_WAIT_MS = 5_000;
const PROCESS_TREE_POLL_MS = 25;

export interface TerminateProcessTreeOptions {
  /** Tests can inject a short grace period; production keeps the five-second default. */
  graceMs?: number;
  /** Tests can inject a short forced-exit deadline; production remains bounded at five seconds. */
  forceWaitMs?: number;
  platform?: NodeJS.Platform;
}

export function processTreeKillCommand(
  pid: number,
  force: boolean,
  platform = process.platform,
): { command: string; args: string[] } | undefined {
  if (!Number.isInteger(pid) || pid <= 0) return undefined;
  if (platform === "win32")
    return { command: "taskkill", args: ["/pid", String(pid), "/t", ...(force ? ["/f"] : [])] };
  return undefined;
}

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

function isProcessAlive(pid: number): boolean {
  try {
    process.kill(pid, 0);
    return true;
  } catch (error) {
    if ((error as NodeJS.ErrnoException).code === "ESRCH") return false;
    if ((error as NodeJS.ErrnoException).code === "EPERM") return true;
    throw error;
  }
}

function isPosixProcessGroupAlive(pid: number): boolean {
  try {
    process.kill(-pid, 0);
    return true;
  } catch (error) {
    if ((error as NodeJS.ErrnoException).code === "ESRCH") return false;
    if ((error as NodeJS.ErrnoException).code === "EPERM") return true;
    throw error;
  }
}

function signalPosixProcessGroup(pid: number, signal: NodeJS.Signals): void {
  try {
    process.kill(-pid, signal);
  } catch (error) {
    if ((error as NodeJS.ErrnoException).code !== "ESRCH") throw error;
  }
}

async function waitForExit(isAlive: () => boolean, timeoutMs?: number): Promise<boolean> {
  const deadline = timeoutMs === undefined ? undefined : Date.now() + timeoutMs;
  while (isAlive()) {
    if (deadline !== undefined && Date.now() >= deadline) return false;
    await sleep(
      deadline === undefined
        ? PROCESS_TREE_POLL_MS
        : Math.max(1, Math.min(PROCESS_TREE_POLL_MS, deadline - Date.now())),
    );
  }
  return true;
}

async function runTaskkill(pid: number, force: boolean, timeoutMs: number): Promise<void> {
  const command = processTreeKillCommand(pid, force, "win32");
  if (!command) return;
  await new Promise<void>((resolve, reject) => {
    let stderr = "";
    let settled = false;
    let timer: ReturnType<typeof setTimeout> | undefined;
    const settle = (error?: Error) => {
      if (settled) return;
      settled = true;
      if (timer) clearTimeout(timer);
      if (error) reject(error);
      else resolve();
    };
    let killer: ChildProcess;
    try {
      killer = spawn(command.command, command.args, {
        stdio: ["ignore", "ignore", "pipe"],
        windowsHide: true,
      });
    } catch (error) {
      settle(error instanceof Error ? error : new Error(String(error)));
      return;
    }
    timer = setTimeout(
      () => {
        try {
          killer.kill();
        } catch {
          // The timeout failure below is the actionable termination error.
        }
        settle(
          new Error(
            `taskkill ${force ? "/T /F" : "/T"} timed out for process ${pid} after ${timeoutMs}ms.`,
          ),
        );
      },
      Math.max(1, timeoutMs),
    );
    killer.stderr?.on("data", (data: Buffer) => {
      stderr += data.toString();
    });
    killer.once("error", (error) => settle(error));
    killer.once("close", (code) => {
      if (code === 0) settle();
      else
        settle(
          new Error(
            `taskkill ${force ? "/T /F" : "/T"} failed for process ${pid}${stderr ? `: ${stderr.trim()}` : ` (exit ${code ?? "unknown"})`}`,
          ),
        );
    });
  });
}

/**
 * Terminate a child tree and wait for the tree's lifecycle to finish. POSIX targets the detached
 * process group and checks group liveness; Windows waits for taskkill /T and escalates to /F.
 */
export async function terminateProcessTree(
  proc: Pick<ChildProcess, "pid">,
  {
    graceMs = PROCESS_TREE_GRACE_MS,
    forceWaitMs = PROCESS_TREE_FORCE_WAIT_MS,
    platform = process.platform,
  }: TerminateProcessTreeOptions = {},
): Promise<void> {
  const pid = proc.pid;
  if (!pid) return;
  if (platform !== "win32") {
    signalPosixProcessGroup(pid, "SIGTERM");
    if (await waitForExit(() => isPosixProcessGroupAlive(pid), graceMs)) return;
    signalPosixProcessGroup(pid, "SIGKILL");
    if (!(await waitForExit(() => isPosixProcessGroupAlive(pid), forceWaitMs)))
      throw new Error(`Process group ${pid} survived forced termination.`);
    return;
  }

  let gracefulFailure: Error | undefined;
  try {
    await runTaskkill(pid, false, graceMs);
  } catch (error) {
    gracefulFailure = error instanceof Error ? error : new Error(String(error));
  }
  if (!isProcessAlive(pid)) return;

  try {
    await runTaskkill(pid, true, forceWaitMs);
  } catch (error) {
    const forceFailure = error instanceof Error ? error : new Error(String(error));
    if (isProcessAlive(pid))
      throw new Error(
        `Failed to terminate process tree ${pid}: ${gracefulFailure?.message ?? "graceful taskkill did not stop it"}; ${forceFailure.message}`,
      );
  }
  if (!(await waitForExit(() => isProcessAlive(pid), forceWaitMs)))
    throw new Error(`Process tree ${pid} survived forced termination.`);
}

export function buildPiAgentArgs(options: SpawnPiAgentOptions): string[] {
  const args: string[] = ["--mode", "json", "-p", "--no-session", "--no-prompt-templates"];
  if (options.model) args.push("--model", options.model);
  if (options.thinking) args.push("--thinking", options.thinking);
  if (options.tools) {
    if (options.tools.length > 0) args.push("--tools", options.tools.join(","));
    else args.push("--no-tools");
  }
  args.push("--exclude-tools", CHILD_ORCHESTRATION_TOOL_NAMES.join(","));
  return args;
}

/** End stdin after writing the caller's complete prompt without transformation. */
export function sendPromptToStdin(
  stdin: { end(chunk: string, encoding?: BufferEncoding): unknown },
  prompt: string,
): void {
  stdin.end(prompt, "utf8");
}

/** Shared fields every backend's spawn result carries; backends extend it with their own payload. */
export interface BaseProcessResult {
  exitCode: number;
  stderr: string;
  wasAborted: boolean;
  errorMessage?: string;
}

export interface RunLineDelimitedProcessOptions {
  command: string;
  args: string[];
  cwd: string;
  prompt: string;
  env?: NodeJS.ProcessEnv;
  signal?: AbortSignal;
  /** Windows launchers resolved to a shell alias (e.g. `pi`, `claude.cmd`) need a shell. */
  needsShell?: boolean;
  /** Called once per complete stdout line, plus once for any trailing partial line at close. */
  onLine: (line: string) => void;
  /** Injectable only for lifecycle tests. */
  terminationGraceMs?: number;
  /** Injectable only for lifecycle tests. */
  terminationForceWaitMs?: number;
}

/**
 * Spawn a detached child, stream the caller's complete prompt over stdin unchanged, and parse
 * newline-delimited stdout via `onLine`. The shared lifecycle owns process-group termination on
 * abort, trailing-line flushing, and exit-code reconciliation so each backend only supplies its
 * own line parser. The caller's `result` accumulates backend payload in `onLine`; this runner
 * mutates the shared `exitCode`, `stderr`, `wasAborted`, and `errorMessage` fields.
 */
export function runLineDelimitedProcess<R extends BaseProcessResult>(
  result: R,
  options: RunLineDelimitedProcessOptions,
): Promise<R> {
  return new Promise<R>((resolve) => {
    const proc = spawn(options.command, options.args, {
      cwd: options.cwd,
      env: options.env ?? process.env,
      shell: options.needsShell ?? false,
      detached: process.platform !== "win32",
      stdio: ["pipe", "pipe", "pipe"],
    });
    let buffer = "";
    let settled = false;
    let stdinErrored = false;
    let abortRequested = false;
    let termination: Promise<void> | undefined;
    let terminationFailed = false;
    let finalizing = false;

    const removeAbortListener = () => options.signal?.removeEventListener("abort", killProc);
    const recordFailure = (error: unknown) => {
      const message = error instanceof Error ? error.message : String(error);
      result.errorMessage ??= message;
      result.stderr += result.stderr ? `\n${message}` : message;
    };
    const finish = (exitCode: number) => {
      if (settled || finalizing) return;
      finalizing = true;
      void (async () => {
        try {
          await termination;
          if (buffer.trim()) options.onLine(buffer);
        } catch (error) {
          terminationFailed = true;
          recordFailure(error);
          exitCode = 1;
        } finally {
          settled = true;
          removeAbortListener();
          result.exitCode = terminationFailed ? 1 : exitCode;
          resolve(result);
        }
      })();
    };

    proc.stdout.on("data", (data: Buffer) => {
      try {
        buffer += data.toString();
        const lines = buffer.split("\n");
        buffer = lines.pop() || "";
        for (const line of lines) options.onLine(line);
      } catch (error) {
        recordFailure(error);
      }
    });
    proc.stderr.on("data", (data: Buffer) => {
      result.stderr += data.toString();
    });
    proc.stdin.once("error", (error) => {
      stdinErrored = true;
      result.errorMessage ??= error.message;
      result.stderr += error.message;
    });
    proc.once("close", (code) => finish(stdinErrored ? 1 : (code ?? 1)));
    proc.once("error", (error) => {
      recordFailure(error);
      finish(1);
    });

    function killProc(): void {
      if (abortRequested) return;
      abortRequested = true;
      result.wasAborted = true;
      termination = terminateProcessTree(proc, {
        graceMs: options.terminationGraceMs,
        forceWaitMs: options.terminationForceWaitMs,
      });
      void termination.catch(() => finish(1));
    }
    if (options.signal) {
      if (options.signal.aborted) killProc();
      else options.signal.addEventListener("abort", killProc, { once: true });
    }

    // end() is intentional: the child reads the whole prompt from stdin before executing it.
    sendPromptToStdin(proc.stdin, options.prompt);
  });
}

/** Spawn Pi and send the caller-provided complete prompt, unchanged, over stdin. */
export function spawnPiAgent(options: SpawnPiAgentOptions): Promise<SpawnPiAgentResult> {
  const result: SpawnPiAgentResult = {
    exitCode: 0,
    messages: [],
    stderr: "",
    wasAborted: false,
    usage: emptyUsage(),
  };

  const invocation = getPiInvocation(buildPiAgentArgs(options));
  const needsShell = process.platform === "win32" && invocation.command === "pi";

  const processLine = (line: string) => {
    if (!line.trim()) return;
    let event: any;
    try {
      event = JSON.parse(line);
    } catch {
      return;
    }
    if (event.type === "message_end" && event.message) {
      const msg = event.message as Message;
      result.messages.push(msg);
      if (msg.role === "assistant") {
        result.usage.turns++;
        const usage = msg.usage;
        if (usage) {
          result.usage.input += usage.input || 0;
          result.usage.output += usage.output || 0;
          result.usage.cacheRead += usage.cacheRead || 0;
          result.usage.cacheWrite += usage.cacheWrite || 0;
          result.usage.cost += usage.cost?.total || 0;
          result.usage.contextTokens = usage.totalTokens || 0;
        }
        if (!result.model && msg.model) result.model = msg.model;
        if (msg.stopReason) result.stopReason = msg.stopReason;
        if (msg.errorMessage) result.errorMessage = msg.errorMessage;
      }
      options.onMessage?.(msg);
    }
    if (event.type === "tool_result_end" && event.message) {
      result.messages.push(event.message as Message);
      options.onToolResult?.(event.message as Message);
    }
    if (event.type === "tool_execution_start" && event.toolName) {
      options.onToolExecutionStart?.({
        toolCallId: event.toolCallId ?? "",
        toolName: event.toolName,
        args: event.args ?? {},
      });
    }
  };

  return runLineDelimitedProcess(result, {
    command: invocation.command,
    args: invocation.args,
    cwd: options.cwd,
    prompt: options.prompt,
    env: options.env,
    signal: options.signal,
    needsShell,
    onLine: processLine,
    terminationGraceMs: options.terminationGraceMs,
    terminationForceWaitMs: options.terminationForceWaitMs,
  });
}