import { randomBytes } from "node:crypto";
import { spawn, type ChildProcess } from "node:child_process";
import * as path from "node:path";
import { fileURLToPath } from "node:url";
import { safeStringify, toSerializable } from "../_workflows/serialization.ts";

const MAX_SOURCE_BYTES = 512 * 1024;
const MAX_ARGS_BYTES = 256 * 1024;
const MAX_RESULT_BYTES = 1024 * 1024;
const MAX_ASSIGNMENT_MESSAGE_BYTES = 512 * 1024;
const MAX_ASSIGNMENT_REQUESTS = 32;
const MAX_STAGE_UPDATES = 64;
const DEFAULT_SCRIPT_RUNTIME_MS = 10 * 60 * 1000;

export interface ScriptAssignmentResult {
  ok: boolean;
  output: string;
  structured?: unknown;
  error?: string;
}

export interface RunScriptSandboxOptions {
  source: string;
  args: unknown;
  cwd: string;
  signal: AbortSignal;
  runtimeTimeoutMs?: number;
  onAssign: (
    role: string,
    prompt: string,
    signal: AbortSignal,
    requestId: number,
  ) => Promise<ScriptAssignmentResult>;
  onStage: (title: string) => void;
}

function byteLength(value: string) {
  return Buffer.byteLength(value, "utf8");
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

function errorText(error: unknown) {
  return error instanceof Error ? error.message : String(error);
}

function terminateChild(child: ChildProcess) {
  if (child.exitCode !== null || child.signalCode !== null) return;
  child.kill("SIGTERM");
  const force = setTimeout(() => {
    if (child.exitCode === null && child.signalCode === null)
      child.kill("SIGKILL");
  }, 1_000);
  force.unref?.();
}

export function validateScriptIpcEnvelope(raw: unknown, token: string) {
  if (!isRecord(raw) || raw.token !== token || typeof raw.kind !== "string") {
    throw new Error("Script sandbox sent an invalid IPC message");
  }
  return raw;
}

export function parseScriptAssignmentPayload(payloadJson: unknown) {
  if (
    typeof payloadJson !== "string" ||
    byteLength(payloadJson) > MAX_ASSIGNMENT_MESSAGE_BYTES
  ) {
    throw new Error("Script sandbox sent an oversized assignment request");
  }
  let payload: unknown;
  try {
    payload = JSON.parse(payloadJson);
  } catch {
    throw new Error("Script sandbox sent malformed assignment JSON");
  }
  if (
    !isRecord(payload) ||
    !Number.isSafeInteger(payload.id) ||
    typeof payload.id !== "number" ||
    payload.id < 1 ||
    typeof payload.role !== "string" ||
    payload.role.trim().length === 0 ||
    payload.role.length > 80 ||
    typeof payload.prompt !== "string" ||
    payload.prompt.length > 100_000
  ) {
    throw new Error("Script sandbox sent an invalid assignment request");
  }
  return {
    id: payload.id,
    role: payload.role.trim(),
    prompt: payload.prompt,
  };
}

export class ScriptRequestBudget {
  private readonly ids = new Set<number>();
  private count = 0;

  accept(id: number) {
    if (this.ids.has(id)) {
      throw new Error("Script sandbox sent a duplicate assignment request");
    }
    if (++this.count > MAX_ASSIGNMENT_REQUESTS) {
      throw new Error("Script sandbox exceeded its assignment request budget");
    }
    this.ids.add(id);
  }
}

/**
 * Execute orchestration code in a separate, permission-restricted Node process.
 * The child can only invoke the narrow agent/phase IPC protocol and is always
 * terminated on completion, cancellation, or protocol failure. The workflow
 * itself and its assignment requests have no wall-clock deadline. Active requests
 * are aborted only when the workflow is cancelled or the sandbox is cleaned up.
 */
export function runScriptSandbox(options: RunScriptSandboxOptions) {
  if (!process.allowedNodeEnvironmentFlags.has("--permission")) {
    return Promise.reject(
      new Error("This Node runtime cannot enforce script child permissions"),
    );
  }
  if (byteLength(options.source) > MAX_SOURCE_BYTES) {
    return Promise.reject(
      new Error(`Script source exceeds the ${MAX_SOURCE_BYTES} byte limit`),
    );
  }

  try {
    const exactArgsJson = JSON.stringify({
      defined: options.args !== undefined,
      value: options.args,
    });
    if (exactArgsJson && byteLength(exactArgsJson) > MAX_ARGS_BYTES) {
      return Promise.reject(new Error("Script args exceed the IPC limit"));
    }
  } catch {
    // The bounded serializer below handles cycles and hostile values.
  }

  const argsJson = safeStringify(
    { defined: options.args !== undefined, value: options.args },
    { maxBytes: MAX_ARGS_BYTES + 1, maxDepth: 16, maxNodes: 10_000 },
  );
  if (byteLength(argsJson) > MAX_ARGS_BYTES) {
    return Promise.reject(new Error("Script args exceed the IPC limit"));
  }

  return new Promise<unknown>((resolve, reject) => {
    const workerPath = fileURLToPath(
      new URL("./sandbox-child.cjs", import.meta.url),
    );
    const child = spawn(
      process.execPath,
      [
        "--permission",
        `--allow-fs-read=${path.dirname(workerPath)}`,
        "--max-old-space-size=128",
        "--stack-size=2048",
        workerPath,
      ],
      {
        cwd: options.cwd,
        env: {
          PATH: process.env.PATH ?? "",
          NODE_NO_WARNINGS: "1",
        },
        stdio: ["ignore", "ignore", "ignore", "ipc"],
      },
    );
    const token = randomBytes(24).toString("hex");
    const requestBudget = new ScriptRequestBudget();
    const activeAssignmentRequests = new Map<number, AbortController>();
    let stageCount = 0;
    let finished = false;
    const runtimeTimer = setTimeout(
      () => finish(new Error("Script exceeded its wall-clock time limit")),
      options.runtimeTimeoutMs ?? DEFAULT_SCRIPT_RUNTIME_MS,
    );
    runtimeTimer.unref?.();

    const cleanup = () => {
      for (const abortController of activeAssignmentRequests.values()) {
        abortController.abort(new Error("Script stopped"));
      }
      activeAssignmentRequests.clear();
      clearTimeout(runtimeTimer);
      options.signal.removeEventListener("abort", onAbort);
      child.removeAllListeners("message");
      child.removeAllListeners("error");
      child.removeAllListeners("exit");
      terminateChild(child);
    };
    const finish = (error?: Error, value?: unknown) => {
      if (finished) return;
      finished = true;
      cleanup();
      if (error) reject(error);
      else resolve(value);
    };
    const onAbort = () => finish(new Error("Script was aborted"));

    options.signal.addEventListener("abort", onAbort, { once: true });
    if (options.signal.aborted) {
      onAbort();
      return;
    }

    child.on("error", (error) => finish(error));
    child.on("exit", (code, exitSignal) => {
      if (!finished) {
        finish(
          new Error(
            `Script sandbox exited before completion (${exitSignal ?? code ?? "unknown"})`,
          ),
        );
      }
    });
    child.on("message", (message: unknown) => {
      let raw: Record<string, unknown>;
      try {
        raw = validateScriptIpcEnvelope(message, token);
      } catch (error) {
        finish(error instanceof Error ? error : new Error(String(error)));
        return;
      }
      if (raw.kind === "stage") {
        if (++stageCount > MAX_STAGE_UPDATES) {
          finish(new Error("Script sandbox exceeded its stage update budget"));
          return;
        }
        if (
          typeof raw.payloadJson !== "string" ||
          byteLength(raw.payloadJson) > 4096
        ) {
          finish(new Error("Script sandbox sent an invalid stage update"));
          return;
        }
        try {
          const payload: unknown = JSON.parse(raw.payloadJson);
          if (!isRecord(payload) || typeof payload.title !== "string") {
            throw new Error("invalid title");
          }
          options.onStage(payload.title.slice(0, 160));
        } catch {
          finish(new Error("Script sandbox sent an invalid stage update"));
        }
        return;
      }
      if (raw.kind === "assign") {
        let payload: ReturnType<typeof parseScriptAssignmentPayload>;
        try {
          payload = parseScriptAssignmentPayload(raw.payloadJson);
        } catch (error) {
          finish(error instanceof Error ? error : new Error(String(error)));
          return;
        }
        try {
          requestBudget.accept(payload.id);
        } catch (error) {
          finish(error instanceof Error ? error : new Error(String(error)));
          return;
        }
        const id = payload.id;
        const abortController = new AbortController();
        const sendResult = (result: ScriptAssignmentResult) => {
          if (!activeAssignmentRequests.delete(id)) return;
          if (finished || !child.connected) return;
          const normalized = toSerializable(result, {
            maxDepth: 16,
            maxNodes: 10_000,
            maxStringBytes: 128 * 1024,
          });
          let resultJson = JSON.stringify(normalized);
          if (byteLength(resultJson) > MAX_ASSIGNMENT_MESSAGE_BYTES) {
            resultJson = JSON.stringify({
              ok: false,
              output: "",
              error: "Assignment result exceeded the script IPC output limit",
            });
          }
          child.send({ token, kind: "assignmentResult", id, resultJson });
        };
        activeAssignmentRequests.set(id, abortController);
        void options
          .onAssign(
            payload.role.trim(),
            payload.prompt,
            abortController.signal,
            payload.id,
          )
          .then(sendResult)
          .catch((error) =>
            sendResult({ ok: false, output: "", error: errorText(error) }),
          );
        return;
      }
      if (raw.kind === "result") {
        if (
          typeof raw.resultJson !== "string" ||
          byteLength(raw.resultJson) > MAX_RESULT_BYTES
        ) {
          finish(new Error("Script result exceeded the IPC limit"));
          return;
        }
        try {
          const normalized = toSerializable(JSON.parse(raw.resultJson));
          finish(undefined, JSON.parse(JSON.stringify(normalized)));
        } catch (error) {
          finish(
            new Error(`Script returned invalid JSON: ${errorText(error)}`),
          );
        }
        return;
      }
      if (raw.kind === "error" && typeof raw.error === "string") {
        finish(new Error(raw.error.slice(0, 16 * 1024)));
        return;
      }
      finish(new Error("Script sandbox sent an unknown IPC message"));
    });

    child.send(
      {
        kind: "init",
        token,
        source: options.source,
        argsJson,
      },
      (error) => {
        if (error) finish(error);
      },
    );
  });
}