import { spawn } from "node:child_process";
import { createServer, request as httpRequest } from "node:http";
import { mkdtemp, readFile, writeFile, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join, dirname } from "node:path";

const abi = "quixos-transaction-ports-v1";
const fixedError = (code) => Object.assign(Error(code), { code });
const object = (v) => !!v && typeof v === "object" && !Array.isArray(v);
const absolute = (value) =>
  typeof value === "string" &&
  value.startsWith("/") &&
  !/[\s\u0000-\u001f]/.test(value) &&
  !value.split("/").includes("..");

/** Trusted stdio adapter inside one workspace VM. Engine processes never own
 * coordinator credentials. The supervising service must impose VM/OS limits.
 * Per-frame process replacement prevents retained globals/handles crossing calls. */
export async function serveRestrictedRegistry(config, io = process) {
  if (
    !object(config) ||
    !absolute(config.workerd) ||
    !absolute(config.worker) ||
    !absolute(config.module) ||
    !/^[$A-Z_a-z][$\w]*$/.test(config.registryExport) ||
    !Array.isArray(config.exports) ||
    !config.exports.length ||
    config.exports.length > 10000 ||
    config.exports.some((x) => typeof x !== "string" || !x || x.length > 4096) ||
    new Set(config.exports).size !== config.exports.length ||
    !Array.isArray(config.sdk) ||
    config.sdk.length > 8 ||
    config.sdk.some(
      (m) =>
        !object(m) ||
        !absolute(m.path) ||
        typeof m.name !== "string" ||
        !m.name ||
        /[\u0000-\u001f]/.test(m.name) ||
        ["entry.mjs", "package.mjs"].includes(m.name),
    ) ||
    new Set(config.sdk.map((m) => m.name)).size !== config.sdk.length
  )
    throw fixedError("INVALID_RESTRICTED_CONFIGURATION");
  const limits = {
    maxBytes: config.maxBytes ?? 1024 * 1024,
    maxPortCalls: config.maxPortCalls ?? 1000,
    deadlineMs: config.deadlineMs ?? 4000,
    maxQueued: config.maxQueued ?? 8,
  };
  if (
    !Number.isSafeInteger(limits.maxBytes) ||
    limits.maxBytes < 1024 ||
    limits.maxBytes > 4 * 1024 * 1024 ||
    !Number.isSafeInteger(limits.maxPortCalls) ||
    limits.maxPortCalls < 1 ||
    limits.maxPortCalls > 10000 ||
    !Number.isSafeInteger(limits.deadlineMs) ||
    limits.deadlineMs < 1 ||
    limits.deadlineMs > 60000 ||
    !Number.isSafeInteger(limits.maxQueued) ||
    limits.maxQueued < 0 ||
    limits.maxQueued > 64
  )
    throw fixedError("INVALID_RESTRICTED_LIMITS");
  const packageSource = await readFile(config.module, "utf8");
  if (Buffer.byteLength(packageSource) > 1024 * 1024) throw fixedError("RESTRICTED_SOURCE_LIMIT");
  const workerSource = await readFile(config.worker, "utf8");
  if (Buffer.byteLength(workerSource) > 65536) throw fixedError("RESTRICTED_ENTRY_LIMIT");
  const sdk = await Promise.all(config.sdk.map(async (m) => ({ ...m, source: await readFile(m.path, "utf8") })));
  if (sdk.some((m) => Buffer.byteLength(m.source) > 2 * 1024 * 1024)) throw fixedError("RESTRICTED_SDK_LIMIT");
  let active,
    buffered = Buffer.alloc(0),
    stopped = false;
  const queued = [];
  const send = (value) => {
    if (stopped) return;
    const line = JSON.stringify({ abi, ...value }) + "\n";
    if (
      Buffer.byteLength(line) > limits.maxBytes ||
      io.stdout.writableLength + Buffer.byteLength(line) > limits.maxBytes * 2
    )
      throw fixedError("RESTRICTED_BACKPRESSURE");
    io.stdout.write(line);
  };
  const retire = async (call, reason) => {
    if (call.retiring) return call.retiring;
    call.open = false;
    call.retiring = (async () => {
      clearTimeout(call.timer);
      const index = queued.indexOf(call);
      if (index >= 0) queued.splice(index, 1);
      for (const pending of call.pending.values()) pending.reject(reason);
      call.pending.clear();
      // Cancellation can arrive during async directory/socket preparation.
      // Wait for it to settle before claiming every process/handle is gone.
      await call.preparing?.catch(() => {});
      if (call.child?.pid) {
        try {
          process.kill(-call.child.pid, "SIGKILL");
        } catch (error) {
          if (error.code !== "ESRCH") throw error;
        }
        await call.exited;
      }
      call.server?.closeAllConnections();
      if (call.server) await new Promise((resolve) => call.server.close(resolve));
      if (call.directory) await rm(call.directory, { recursive: true, force: true });
      if (active === call) active = undefined;
    })();
    return call.retiring;
  };
  const fail = async (error) => {
    if (stopped) return;
    stopped = true;
    io.stdin.destroy();
    for (const call of [...queued]) await retire(call, error);
    if (active) await retire(active, error);
    io.stdout.end();
    io.exitCode = 1;
  };
  const body = async (stream) => {
    const chunks = [];
    let bytes = 0;
    for await (const chunk of stream) {
      bytes += chunk.length;
      if (bytes > limits.maxBytes) throw fixedError("RESTRICTED_MESSAGE_LIMIT");
      chunks.push(chunk);
    }
    return JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(Buffer.concat(chunks)));
  };
  const pump = () => {
    if (!stopped && !active && queued.length) void invoke(queued.shift()).catch(fail);
  };
  const finish = async (call, reason, reply) => {
    await retire(call, reason);
    send(reply);
    pump();
  };
  const invoke = async (call) => {
    const { message } = call;
    active = call;
    const terminate = (code) => {
      if (!call.open) return;
      void finish(call, fixedError(code), {
        invocation: message.invocation,
        status: "failed",
        error: { code, message: code },
      }).catch(fail);
    };
    call.terminate = terminate;
    try {
      let engine;
      call.preparing = (async () => {
        call.directory = await mkdtemp(join(tmpdir(), "qx-frame-"));
        if (!call.open) return;
        const broker = join(call.directory, "broker.sock");
        engine = join(call.directory, "engine.sock");
        call.server = createServer(async (request, response) => {
          try {
            if (!call.open || request.method !== "POST" || request.url !== "/port") throw fixedError("EXPIRED_CONTEXT");
            const port = await body(request);
            if (
              !call.open ||
              !object(port) ||
              Object.keys(port).some((key) => !["portId", "action", "input", "key"].includes(key)) ||
              typeof port.portId !== "string" ||
              !object(port.action) ||
              !Object.hasOwn(port, "input")
            )
              throw fixedError("INVALID_RESTRICTED_PORT");
            if (++call.sequence > limits.maxPortCalls || call.pending.size >= limits.maxPortCalls)
              throw fixedError("RESTRICTED_PORT_LIMIT");
            const sequence = call.sequence;
            const result = new Promise((resolve, reject) => call.pending.set(sequence, { resolve, reject }));
            void result.catch(() => {});
            send({
              kind: "port",
              invocation: message.invocation,
              frame: message.frame,
              sequence,
              portId: port.portId,
              action: { ...port.action, ...(port.key === undefined ? {} : { key: port.key }) },
              input: port.input,
            });
            const value = await result;
            if (!call.open) throw fixedError("EXPIRED_CONTEXT");
            const encoded = JSON.stringify(value);
            if (Buffer.byteLength(encoded) > limits.maxBytes) throw fixedError("RESTRICTED_MESSAGE_LIMIT");
            response.writeHead(200, { "content-type": "application/json" });
            response.end(encoded);
          } catch {
            response.writeHead(409);
            response.end('{"status":"failed","error":{"code":"RESTRICTED_PORT_FAILED"}}');
          }
        });
        await new Promise((resolve, reject) => {
          call.server.once("error", reject);
          call.server.listen(broker, resolve);
        });
        await writeFile(join(call.directory, "package.mjs"), packageSource);
        await writeFile(join(call.directory, "entry.mjs"), workerSource);
        const modules = [
          { name: "entry.mjs", esModule: 'embed "entry.mjs"' },
          { name: "package.mjs", esModule: 'embed "package.mjs"' },
        ];
        for (const [index, module] of sdk.entries()) {
          const file = `sdk-${index}.mjs`;
          await writeFile(join(call.directory, file), module.source);
          modules.push({ name: module.name, esModule: "embed " + JSON.stringify(file) });
        }
        const setting = { registryExport: config.registryExport, exports: config.exports, maxBytes: limits.maxBytes };
        const configuration = `using Workerd = import "/workerd/workerd.capnp"; const config :Workerd.Config = (
        services = [(name="main", worker=(modules=[${modules.map((m) => `(name=${JSON.stringify(m.name)}, esModule=${m.esModule})`).join(",")}],
          compatibilityDate="2026-10-05", globalOutbound="deny",
          bindings=[(name="BROKER", service="broker"),(name="CONFIG",json=${JSON.stringify(JSON.stringify(setting))})])),
          (name="broker",external=(address=${JSON.stringify("unix:" + broker)},http=())), (name="deny",network=(allow=[]))],
        sockets=[(name="frame",address=${JSON.stringify("unix:" + engine)},http=(),service="main")]);`;
        const file = join(call.directory, "config.capnp");
        await writeFile(file, configuration);
        if (!call.open) return;
        call.child = spawn(config.workerd, ["serve", file], {
          detached: true,
          env: { TZ: "UTC", LANG: "C.UTF-8" },
          cwd: call.directory,
          stdio: ["ignore", "ignore", "pipe"],
        });
        call.exited = new Promise((resolve) => {
          call.child.once("exit", resolve);
          call.child.once("error", resolve);
        });
        let diagnostics = 0;
        call.child.stderr.on("data", (chunk) => {
          if ((diagnostics += chunk.length) > 65536) terminate("RESTRICTED_DIAGNOSTIC_LIMIT");
        });
        call.child.once("error", () => terminate("RESTRICTED_START_FAILED"));
        void call.exited.then(() => {
          if (call.open) terminate("RESTRICTED_ENGINE_EXITED");
        });
      })();
      await call.preparing;
      if (!call.open) return;
      const result = await new Promise((resolve, reject) => {
        const attempt = () => {
          if (!call.open) {
            reject(fixedError("EXPIRED_CONTEXT"));
            return;
          }
          const request = httpRequest({ socketPath: engine, path: "/", method: "POST" }, async (response) => {
            try {
              const result = await body(response);
              if (response.statusCode !== 200) throw fixedError("RESTRICTED_PACKAGE_FAILED");
              resolve(result);
            } catch (error) {
              reject(error);
            }
          });
          request.once("error", (error) => {
            if (call.open && ["ENOENT", "ECONNREFUSED"].includes(error.code)) setTimeout(attempt, 5);
            else reject(error);
          });
          request.end(JSON.stringify(message));
        };
        attempt();
      });
      if (!call.open) return;
      if (
        !object(result) ||
        !["returned", "failed"].includes(result.status) ||
        Object.keys(result).some(
          (key) => !(result.status === "returned" ? ["status", "output"] : ["status", "error"]).includes(key),
        ) ||
        (result.status === "returned"
          ? !Object.hasOwn(result, "output")
          : !object(result.error) || !/^[A-Z0-9_]{1,128}$/.test(result.error.code))
      )
        throw fixedError("RESTRICTED_PACKAGE_FAILED");
      // Authored code can alter serialization prototypes. Treat even the
      // platform worker's JSON as untrusted; never accept routing/ABI fields.
      const reply =
        result.status === "returned"
          ? { invocation: message.invocation, status: "returned", output: result.output }
          : {
              invocation: message.invocation,
              status: "failed",
              error: { code: result.error.code, message: result.error.code },
            };
      await finish(call, fixedError("INVOCATION_ENDED"), reply);
    } catch (error) {
      if (call.open) {
        const code = /^RESTRICTED_[A-Z_]+$/.test(error.code ?? "") ? error.code : "RESTRICTED_PACKAGE_FAILED";
        await finish(call, error, { invocation: message.invocation, status: "failed", error: { code, message: code } });
      }
    }
  };
  const receive = (message) => {
    if (!object(message) || message.abi !== abi) throw fixedError("RESTRICTED_ABI_MISMATCH");
    if (message.kind === "shutdown") {
      void fail(fixedError("SERVICE_SHUTDOWN"));
      return;
    }
    if (message.kind === "invoke") {
      if (
        !config.exports.includes(message.exportId) ||
        !["snapshot", "transaction"].includes(message.effect) ||
        typeof message.invocation !== "string" ||
        !message.invocation ||
        message.invocation.length > 4096 ||
        typeof message.frame !== "string" ||
        !message.frame ||
        message.frame.length > 4096
      )
        throw fixedError("INVALID_RESTRICTED_INVOCATION");
      if ([active, ...queued].some((call) => call?.message.invocation === message.invocation))
        throw fixedError("RESTRICTED_DUPLICATE_INVOCATION");
      if (active && queued.length >= limits.maxQueued) {
        send({
          invocation: message.invocation,
          status: "failed",
          error: { code: "RESTRICTED_QUEUE_FULL", message: "RESTRICTED_QUEUE_FULL" },
        });
        return;
      }
      const call = { message, open: true, pending: new Map(), sequence: 0 };
      // Deadline includes queue wait: overload cannot extend a transaction frame.
      call.timer = setTimeout(() => {
        if (!call.open) return;
        void finish(call, fixedError("RESTRICTED_DEADLINE"), {
          invocation: message.invocation,
          status: "failed",
          error: { code: "RESTRICTED_DEADLINE", message: "RESTRICTED_DEADLINE" },
        }).catch(fail);
      }, limits.deadlineMs);
      queued.push(call);
      pump();
      return;
    }
    const waiting = queued.find(
      (call) => call.message.invocation === message.invocation && call.message.frame === message.frame,
    );
    if (waiting && message.kind === "cancel") {
      void finish(waiting, fixedError("INVOCATION_CANCELLED"), {
        kind: "cancelled",
        invocation: message.invocation,
        frame: message.frame,
      }).catch(fail);
      return;
    }
    const call = active;
    if (!call || message.invocation !== call.message.invocation || message.frame !== call.message.frame) {
      if (message.kind === "cancel") send({ kind: "cancelled", invocation: message.invocation, frame: message.frame });
      else if (message.kind !== "port-result") throw fixedError("RESTRICTED_UNKNOWN_INVOCATION");
      return;
    }
    if (message.kind === "cancel") {
      void finish(call, fixedError("INVOCATION_CANCELLED"), {
        kind: "cancelled",
        invocation: message.invocation,
        frame: message.frame,
      }).catch(fail);
      return;
    }
    if (message.kind !== "port-result" || !["returned", "failed"].includes(message.status))
      throw fixedError("INVALID_RESTRICTED_RESULT");
    const pending = call.pending.get(message.sequence);
    if (!pending) throw fixedError("RESTRICTED_UNKNOWN_PORT");
    call.pending.delete(message.sequence);
    pending.resolve(message);
  };
  io.stdin.on("data", (chunk) => {
    if (stopped) return;
    try {
      let offset = 0;
      while (offset < chunk.length) {
        const newline = chunk.indexOf(10, offset),
          end = newline < 0 ? chunk.length : newline;
        if (buffered.length + end - offset > limits.maxBytes) throw fixedError("RESTRICTED_MESSAGE_LIMIT");
        buffered = Buffer.concat([buffered, chunk.subarray(offset, end)]);
        if (newline < 0) break;
        receive(JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(buffered)));
        buffered = Buffer.alloc(0);
        offset = newline + 1;
        if (stopped) break;
      }
    } catch (error) {
      void fail(error);
    }
  });
  io.stdin.on("end", () => void fail(fixedError("COORDINATOR_DISCONNECTED")));
  io.stdin.on("error", (error) => void fail(error));
  io.stdout.on("error", (error) => void fail(error));
  send({ kind: "ready", exports: config.exports });
  return { close: () => fail(fixedError("SERVICE_SHUTDOWN")) };
}

if (process.argv[1] === new URL(import.meta.url).pathname) {
  const config = JSON.parse(await readFile(process.argv[2], "utf8"));
  await serveRestrictedRegistry({
    ...config,
    workerd: process.env.QUIXOS_RESTRICTED_WORKERD,
    worker: process.env.QUIXOS_RESTRICTED_WORKER,
    module: join(dirname(process.argv[2]), config.module),
  });
}