import { spawn } from "node:child_process";
import { createServer, request as httpRequest } from "node:http";
import { mkdtemp, readFile, writeFile, rm, open } from "node:fs/promises";
import { constants } from "node:fs";
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("..");
async function boundedSource(path, maximum, code) {
  const handle = await open(path, constants.O_RDONLY | constants.O_NOFOLLOW | constants.O_NONBLOCK);
  try {
    const stat = await handle.stat();
    if (!stat.isFile() || stat.size > maximum) throw fixedError(code);
    const bytes = Buffer.alloc(stat.size + 1);
    let offset = 0;
    while (offset < bytes.length) {
      const read = await handle.read(bytes, offset, bytes.length - offset, offset);
      if (!read.bytesRead) break;
      offset += read.bytesRead;
    }
    if (offset !== stat.size) throw fixedError(code);
    return new TextDecoder("utf-8", { fatal: true }).decode(bytes.subarray(0, offset));
  } finally {
    await handle.close();
  }
}

/** 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,
    maxActive: config.maxActive ?? 2,
    maxNested: config.maxNested ?? 4,
  };
  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 ||
    !Number.isSafeInteger(limits.maxActive) ||
    limits.maxActive < 1 ||
    limits.maxActive > 16 ||
    !Number.isSafeInteger(limits.maxNested) ||
    limits.maxNested < 1 ||
    limits.maxNested > 16
  )
    throw fixedError("INVALID_RESTRICTED_LIMITS");
  const packageSource = await boundedSource(config.module, 1024 * 1024, "RESTRICTED_SOURCE_LIMIT");
  const workerSource = await boundedSource(config.worker, 65536, "RESTRICTED_ENTRY_LIMIT");
  const sdk = [];
  for (const m of config.sdk)
    sdk.push({ ...m, source: await boundedSource(m.path, 2 * 1024 * 1024, "RESTRICTED_SDK_LIMIT") });
  const active = new Map();
  let 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);
  };
  // Linux USER_HZ is 100 on the supported NixOS architectures. Read process
  // CPU outside the engine; authored code cannot forge these usage messages.
  const sampleCPU = async (call) => {
    if (!call.child?.pid) return;
    let text;
    try {
      text = await readFile(`/proc/${call.child.pid}/stat`, "utf8");
    } catch (error) {
      if (error.code === "ENOENT" || error.code === "ESRCH") return;
      throw error;
    }
    const fields = text.slice(text.lastIndexOf(")") + 2).split(" ");
    const cpuMicroseconds = (Number(fields[11]) + Number(fields[12])) * 10000;
    if (!Number.isSafeInteger(cpuMicroseconds) || cpuMicroseconds < (call.cpuMicroseconds ?? 0))
      throw fixedError("RESTRICTED_CPU_ACCOUNTING");
    if (cpuMicroseconds !== (call.cpuMicroseconds ?? 0)) {
      call.cpuMicroseconds = cpuMicroseconds;
      send({ kind: "usage", invocation: call.message.invocation, frame: call.message.frame, cpuMicroseconds });
    }
  };
  const watchCPU = (call) => {
    call.cpuTimer = setTimeout(() => {
      call.cpuRead = sampleCPU(call).then(() => {
        if (call.open) watchCPU(call);
      });
      void call.cpuRead.catch(fail);
    }, 20);
  };
  const retire = async (call, reason) => {
    if (call.retiring) return call.retiring;
    call.open = false;
    call.retiring = (async () => {
      clearTimeout(call.timer);
      clearTimeout(call.cpuTimer);
      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(() => {});
      clearTimeout(call.cpuTimer);
      await call.cpuRead?.catch(() => {});
      let usageFailure;
      try {
        await sampleCPU(call);
      } catch (error) {
        usageFailure = error;
      }
      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 });
      active.delete(call.message.invocation);
      if (usageFailure) throw usageFailure;
    })();
    return call.retiring;
  };
  const fail = async (error) => {
    if (stopped) return;
    stopped = true;
    io.stdin.destroy();
    for (const call of [...queued]) await retire(call, error);
    await Promise.all([...active.values()].map((call) => retire(call, 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 available = (call) =>
    [...active.values()].filter((entry) => entry.nested === call.nested).length <
    (call.nested ? limits.maxNested : limits.maxActive);
  const pump = () => {
    while (!stopped) {
      const index = queued.findIndex(available);
      if (index < 0) return;
      void invoke(queued.splice(index, 1)[0]).catch(fail);
    }
  };
  const finish = async (call, reason, reply) => {
    if (reply.kind === "cancelled") call.cancellationReply = reply;
    return (call.finished ??= (async () => {
      await retire(call, reason);
      send(call.cancellationReply ?? reply);
      pump();
    })());
  };
  const invoke = async (call) => {
    const { message } = call;
    active.set(message.invocation, 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 > Math.min(limits.maxPortCalls, message.budget?.calls ?? 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);
        });
        watchCPU(call);
        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", agent: false },
            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 (
        message.budget !== undefined &&
        (!object(message.budget) ||
          Object.keys(message.budget).sort().join(",") !== "calls,deadlineMs,depth,root" ||
          typeof message.budget.root !== "string" ||
          !message.budget.root ||
          message.budget.root.length > 4096 ||
          !Number.isSafeInteger(message.budget.depth) ||
          message.budget.depth < 0 ||
          message.budget.depth > 32 ||
          !Number.isSafeInteger(message.budget.calls) ||
          message.budget.calls < 1 ||
          message.budget.calls > 1000000 ||
          !Number.isSafeInteger(message.budget.deadlineMs) ||
          message.budget.deadlineMs < 1 ||
          message.budget.deadlineMs > 60000)
      )
        throw fixedError("INVALID_RESTRICTED_BUDGET");
      if (active.has(message.invocation) || queued.some((call) => call.message.invocation === message.invocation))
        throw fixedError("RESTRICTED_DUPLICATE_INVOCATION");
      const call = { message, open: true, pending: new Map(), sequence: 0, nested: (message.budget?.depth ?? 0) > 1 };
      // Descendants cannot wait behind ancestors that hold all nested capacity.
      // Reject excess fan-out/depth rather than deadlock an admitted call tree.
      if (!available(call) && (call.nested || queued.length >= limits.maxQueued)) {
        send({
          invocation: message.invocation,
          status: "failed",
          error: { code: "RESTRICTED_QUEUE_FULL", message: "RESTRICTED_QUEUE_FULL" },
        });
        return;
      }
      // 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);
        },
        Math.min(limits.deadlineMs, message.budget?.deadlineMs ?? 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.get(message.invocation);
    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),
  });
}