// node --test modules/agents/test.mjs   (needs `cargo wasm` here and in the sibling checkouts; no network)
// The agents module composed with the real requester (composition.json) against a fake PDS,
// atirc and two bidders (stub.mjs), and serve.json's ingress forwarding the XRPC methods to handlers.
import assert from "node:assert/strict";
import { execFileSync } from "node:child_process";
import { access, readFile } from "node:fs/promises";
import { test } from "node:test";

import { did, now } from "../../../atlibs/rust/src/harness.mjs";
import { loadComposition, runCompose } from "../../../atlibs/javascript/src/compose/index.js";
import { moduleOwner, runModule } from "../../../atlibs/javascript/src/run-module/index.js";
import { HAT7, HAT9, LAUNCH, T, world } from "./stub.mjs";

const wasm = await readFile(new URL("../../target/wasm32-unknown-unknown/release/agents.wasm", import.meta.url));

// ctx.wasm.run for any module bytes in process, now fixed, and an in-memory rkey ledger.
function memoryCtx() {
	const instances = new Map();
	const records = new Map();
	return {
		wasm: {
			run: async ({ module, entry, input }) => {
				if (!instances.has(module)) instances.set(module, new WebAssembly.Instance(new WebAssembly.Module(module), {}).exports);
				const x = instances.get(module);
				const data = new TextEncoder().encode(JSON.stringify({ ...input, now }));
				const ptr = x.alloc(data.length);
				new Uint8Array(x.memory.buffer, ptr, data.length).set(data);
				const packed = BigInt.asUintN(64, x[entry](ptr, data.length));
				return JSON.parse(new TextDecoder().decode(new Uint8Array(x.memory.buffer, Number(packed >> 32n), Number(packed & 0xffffffffn))));
			},
		},
		storage: { managed_records: { get: async (k) => records.get(k) ?? null, put: async (k, v) => void records.set(k, v), delete: async (k) => void records.delete(k) } },
	};
}

// The checkouts a module key may be in, each with its own cargo wasm build: this repo, then the siblings that
// own the modules the compositions borrow (requester and policy-open; ingress; iroh-user-data, a hook default).
const ROOTS = ["../../", "../../../atlibs-socialweb-computer-requester/", "../../../atlibs-fedproxy/", "../../../atlibs-iroh/"].map((r) => new URL(r, import.meta.url));
const exists = (url) => access(url).then(() => true, () => false);
async function rootOf(key) {
	for (const root of ROOTS) if (await exists(new URL(`modules/${key}/module.json`, root))) return root;
	throw new Error(`${key}: no modules/${key}/module.json in ${ROOTS.map((r) => r.pathname).join(", ")}`);
}

// composition.json with each instance's module from its checkout's cargo wasm build.
async function loadAgents(ctx, file = "./composition.json") {
	const manifest = JSON.parse(await readFile(new URL(file, import.meta.url), "utf8"));
	const load = async (key) => {
		const root = await rootOf(key);
		return {
			...JSON.parse(await readFile(new URL(`modules/${key}/module.json`, root), "utf8")),
			module: new Uint8Array(await readFile(new URL(`target/wasm32-unknown-unknown/release/${key.replaceAll("-", "_")}.wasm`, root))),
		};
	};
	return loadComposition(ctx, manifest, { load });
}

const ctx = memoryCtx();
const comp = await loadAgents(ctx);
const run = (w, entry, input, options) => runCompose(ctx, comp, "agents", entry, { did, ...input }, { callFor: () => w.call, ...options });
const LAUNCH_INPUT = { bidders: [HAT7, HAT9], user_data: "#cloud-config\n{}\n", poll_ms: 1, name: "tester", agent: "atirc pub 'atlibs-agents test: hi' && exec sleep infinity" };
const b64 = (s) => Buffer.from(s).toString("base64");
// The user's records but the attestation key binding (badgeBlueKeys), which the requester keeps by design.
const left = (w) => [...w.repo.keys()].filter((u) => !u.includes(`/${T}.badgeBlueKeys/`));

test("module has no imports and exports its entries", () => {
	const module = new WebAssembly.Module(wasm);
	assert.equal(WebAssembly.Module.imports(module).length, 0);
	const names = WebAssembly.Module.exports(module).map((e) => e.name);
	for (const name of ["memory", "alloc", "launch", "status", "shutdown", "allow_bidders"]) assert.ok(names.includes(name), name);
});

test("launch: record, mint via the PDS, requester contract keeping the VM, agent started over ssh, record running", async () => {
	const w = world({ did });
	const early = [];
	const ran = await run(w, "launch", LAUNCH_INPUT, { onEvent: (_e, out, inst) => inst === "agents" && out.result && early.push(out.result) });
	assert.equal(ran.error, undefined, ran.error);
	const r = ran.result;
	assert.equal(r.status, "running");
	assert.match(r.rkey, /^3mzz/);
	assert.equal(r.subhandle, `tester-${r.rkey}.bob.test`);
	assert.equal(r.receipt.uri.startsWith(`at://${HAT7}/${T}.market.receipt/receipt-`), true);
	assert.deepEqual([early[0].rkey, early[0].uri, early[0].status, early[0].receipt], [r.rkey, `at://${did}/${LAUNCH}/${r.rkey}`, "launching", null], "the rkey is known before the VM");

	const mint = w.calls.find((c) => c.nsid === "chat.atirc.mint");
	assert.deepEqual(mint.input, { subhandle: r.subhandle, exp: Math.floor(Date.parse(now) / 1000) + 7200 });
	const ssh = w.calls.find((c) => c.nsid === "com.fedfork.exec.iroh.connect");
	assert.equal(ssh.input.ticket, "iroh://endpointaabm6");
	const command = ssh.input.command;
	assert.ok(command.includes(`IFS= read -r t; printf 'ATIRC_TOKEN=%s\\n' "$t" > /etc/atirc/atirc.env; unset t;`));
	assert.equal(ssh.input.stdin, "eyJhbGciOi.tok\n", "the token goes over the ssh session's stdin");
	assert.ok(!command.includes("eyJhbGciOi"), "never in the command (ssh's argv on the host, the shell's on the guest)");
	assert.ok(command.includes(b64(LAUNCH_INPUT.agent)), "the agent command, base64");
	assert.ok(command.includes("systemctl restart atirc-agent.service"));

	const record = w.repo.get(r.uri).value;
	assert.equal(record.status, "running");
	assert.equal(record.nick, `tester-${r.rkey}`);
	assert.equal(record.bidder, HAT7);
	assert.deepEqual(Object.keys(record.records).sort(), ["accept", "rfp", "vm"]);
	assert.ok(!JSON.stringify([...w.repo.values()]).includes("eyJhbGciOi"), "the token is in no record");
	assert.ok(!JSON.stringify(ran.result).includes("eyJhbGciOi"), "nor in the result");
	assert.ok(ran.graph.some((g) => g.policy === "allow-bidders" && g.ok));
	assert.ok(ran.graph.some((g) => g.nsid === "com.fedfork.market.contract" && g.to === "requester" && g.ok));
	assert.equal([...w.repo.keys()].filter((u) => u.includes(`/${T}.compute.vm/`)).length, 1, "VM kept: no release yet");
	// The record grew in steps: subhandle after the mint, receipt from the requester's progress call
	// (a callback through runCompose), then running.
	const puts = w.calls.filter((c) => c.nsid === "com.atproto.repo.putRecord" && c.input.collection === LAUNCH).map((c) => c.input.record);
	assert.deepEqual(puts.map((p) => [p.status, Boolean(p.subhandle), Boolean(p.receipt)]), [["launching", true, false], ["launching", true, true], ["running", true, true]]);
	assert.ok(ran.graph.some((g) => g.from === "requester" && g.nsid === "cap" && g.ok), "the requester called the progress callback");
});

test("the default command is atirc agent run", async () => {
	const w = world({ did });
	const { agent, ...input } = LAUNCH_INPUT;
	const ran = await run(w, "launch", input);
	assert.equal(ran.error, undefined, ran.error);
	assert.ok(w.calls.find((c) => c.nsid === "com.fedfork.exec.iroh.connect").input.command.includes(b64("atirc agent run")));
	assert.equal(w.repo.get(ran.result.uri).value.agent, "atirc agent run");
});

test("status reads the record and one poll of the guest; shutdown revokes, releases, waits, deletes everything", async () => {
	const w = world({ did });
	const { result } = await run(w, "launch", LAUNCH_INPUT);
	const status = await run(w, "status", { rkey: result.rkey });
	assert.equal(status.result.record.status, "running");
	assert.equal(status.result.guest.terminated, false);
	assert.equal(status.result.guest.iroh, "iroh://endpointaabm6");

	const ran = await run(w, "shutdown", { rkey: result.rkey, delete: true, poll_ms: 1 });
	assert.equal(ran.error, undefined, ran.error);
	const s = ran.result;
	assert.equal(s.status, "stopped");
	assert.equal(s.revoked, true);
	assert.equal(s.terminated, true);
	assert.equal(s.deleted.length, 6, "vm, rfp, accept, vm.delete, event and the launch record");
	assert.deepEqual(s.errors, []);
	assert.equal(w.calls.find((c) => c.nsid === "chat.atirc.revoke").input.subhandle, result.subhandle);
	assert.deepEqual(left(w), [], "nothing left in the repo");
});

test("shutdown without delete marks the record stopped; a second shutdown has nothing to release", async () => {
	const w = world({ did });
	const { result } = await run(w, "launch", LAUNCH_INPUT);
	await run(w, "shutdown", { rkey: result.rkey, poll_ms: 1 });
	assert.equal(w.repo.get(result.uri).value.status, "stopped");
	assert.equal(left(w).length, 1);
});

test("the allow-bidders policy refuses other bidders: the launch fails, nothing is requested, the record says why", async () => {
	const w = world({ did });
	const ran = await run(w, "launch", { ...LAUNCH_INPUT, bidders: ["did:plc:someoneelse"] });
	assert.equal(ran.result.status, "failed");
	const record = w.repo.get(`at://${did}/${LAUNCH}/${ran.result.rkey}`).value;
	assert.equal(record.status, "failed");
	assert.match(record.error, /policy allow-bidders denied: bidders not allowed: did:plc:someoneelse/);
	assert.equal(w.calls.filter((c) => c.input?.collection === `${T}.compute.vm`).length, 0);
	assert.ok(w.calls.some((c) => c.nsid === "chat.atirc.revoke"), "the minted token is revoked");
});

test("an agent that does not start releases the VM and cleans up; a refused mint never asks for a VM", async () => {
	const w = world({ did, shellOutput: "Failed to start atirc-agent.service\n" });
	const ran = await run(w, "launch", LAUNCH_INPUT);
	const record = w.repo.get(`at://${did}/${LAUNCH}/${ran.result.rkey}`).value;
	assert.equal(record.status, "failed");
	assert.match(record.error, /starting the agent on the guest/);
	assert.equal(left(w).length, 1, "the round's records are deleted; the failed launch record stays");
	assert.ok(w.calls.some((c) => c.input?.collection === `${T}.market.event`), "released");

	const refused = world({ did, mint: false });
	const r2 = await run(refused, "launch", LAUNCH_INPUT);
	assert.match(refused.repo.get(`at://${did}/${LAUNCH}/${r2.result.rkey}`).value.error, /chat.atirc.mint: 403/);
	assert.equal(refused.calls.filter((c) => c.nsid.startsWith(T)).length, 0);
});

test("a token with shell metacharacters is refused before it reaches the guest", async () => {
	const w = world({ did, token: "x'; rm -rf / #" });
	const ran = await run(w, "launch", LAUNCH_INPUT);
	assert.match(w.repo.get(`at://${did}/${LAUNCH}/${ran.result.rkey}`).value.error, /unexpected characters/);
	assert.equal(w.calls.filter((c) => c.nsid === "com.fedfork.exec.iroh.connect").length, 0);
});

test("a long token leaves the command as it is (the token is stdin); webrtc goes through webrtc.connect", async () => {
	const w = world({ did, token: "a".repeat(1500) });
	const ran = await run(w, "launch", { ...LAUNCH_INPUT, via: "webrtc", tag: "ssh" });
	assert.equal(ran.result.status, "running", JSON.stringify(ran.result));
	const ssh = w.calls.find((c) => c.nsid === "com.fedfork.exec.webrtc.connect");
	assert.ok(ssh.input.command.length < 4096, `${ssh.input.command.length}`);
	assert.ok(!ssh.input.command.includes("aaaa"));
	assert.equal(ssh.input.stdin, `${"a".repeat(1500)}\n`);
	assert.equal(ssh.input.tag, "ssh");
});

test("bad input fails before any call", async () => {
	for (const [input, error] of [
		[{ ...LAUNCH_INPUT, bidders: [] }, /bidders/],
		[{ ...LAUNCH_INPUT, name: "Bad_Name" }, /name/],
		[{ ...LAUNCH_INPUT, via: "tcp" }, /via/],
		[{ ...LAUNCH_INPUT, user_data: " " }, /user_data/],
		[{ ...LAUNCH_INPUT, exp_s: 5 }, /exp_s/],
		[{ ...LAUNCH_INPUT, exp_s: 86401 }, /exp_s must be 60 to 86400/],
	]) {
		const w = world({ did });
		assert.match((await run(w, "launch", input)).error, error);
		assert.equal(w.calls.length, 0);
	}
	assert.match((await run(world({ did }), "shutdown", { rkey: "../x" })).error, /rkey/);
});

// The agents module alone under runModule, its routed calls answered by `contract` (a stub requester).
async function agentsAlone(w, contract) {
	const call = async (nsid, request) => {
		if (nsid.startsWith("cap.")) return void globalThis.__ack?.(request.input); // the module answering the progress call
		return nsid.startsWith("com.fedfork.market.") ? contract(nsid, request) : w.call(nsid, request);
	};
	return runModule(ctx, { module: new Uint8Array(wasm) }, "launch", { did, ...LAUNCH_INPUT }, call);
}
const launchRecord = (w) => [...w.repo.values()].find((r) => r.uri.includes(`/${LAUNCH}/`)).value;

test("the record has the subhandle right after the mint and the receipt from the requester's progress call, before the contract ends", async () => {
	const w = world({ did });
	const receipt = { uri: `at://${HAT7}/${T}.market.receipt/receipt-1`, cid: "bafyrc", submitEvent: null };
	const round = { vm: { uri: `at://${did}/${T}.compute.vm/v1` }, rfp: { uri: `at://${did}/${T}.market.rfp/r1` }, accept: { uri: `at://${did}/${T}.market.accept/a1` } };
	let atMint, atReceipt;
	const ran = await agentsAlone(w, async (nsid, request) => {
		assert.equal(nsid, "com.fedfork.market.contract");
		assert.deepEqual(request.input.progress, { $cap: "progress" }, "a callback for the requester");
		await new Promise((r) => setTimeout(r, 5));
		atMint = structuredClone(launchRecord(w));
		// The requester reports the receipt; the module answers on the reply only once it is recorded.
		const acked = new Promise((resolve) => (globalThis.__ack = resolve));
		request.emit({ cap: "progress", input: { ...round, receipt }, reply: "cap.reply.1" });
		await acked;
		atReceipt = structuredClone(launchRecord(w));
		return { ...round, receipt, shell: { output: "atlibs-agents: started\n", exitCode: 0 }, errors: [] };
	}).catch((e) => ({ error: e.message }));
	assert.equal(ran.error, undefined, ran.error);
	assert.equal(atMint.subhandle, `tester-${ran.result.rkey}.bob.test`);
	assert.equal(atMint.status, "launching");
	assert.equal(atMint.receipt, undefined);
	assert.deepEqual(atReceipt.receipt, receipt);
	assert.equal(atReceipt.bidder, HAT7);
	assert.deepEqual(atReceipt.records, { vm: round.vm.uri, rfp: round.rfp.uri, accept: round.accept.uri });
	assert.equal(launchRecord(w).status, "running");
});

test("a contract that dies after the receipt was reported still releases the VM and revokes the token", async () => {
	const w = world({ did });
	const receipt = { uri: `at://${HAT7}/${T}.market.receipt/receipt-2`, cid: "bafyrc2", submitEvent: null };
	const ran = await agentsAlone(w, async (nsid, request) => {
		// the requester's release result: it waited for vm.terminated and deleted what it could
		if (nsid === "com.fedfork.market.release") return { release: { terminated: true, payload: { uri: `at://${did}/${T}.compute.events.vm.delete/d` }, event: { uri: `at://${did}/${T}.market.event/e` } }, deleted: [], errors: [] };
		const acked = new Promise((resolve) => (globalThis.__ack = resolve));
		request.emit({ cap: "progress", input: { receipt }, reply: "cap.reply.2" });
		await acked;
		throw new Error("requester run aborted");
	});
	assert.equal(ran.result.status, "failed");
	assert.equal(ran.result.revoked, true);
	assert.equal(ran.result.terminated, true);
	assert.ok(w.calls.some((c) => c.nsid === "chat.atirc.revoke"));
	assert.match(launchRecord(w).error, /contract: requester run aborted/);
});

test("a contract without a receipt: the round's vm and rfp are deleted by teardown (already-gone ones are not errors)", async () => {
	const w = world({ did });
	const vm = `at://${did}/${T}.compute.vm/left1`;
	const rfp = `at://${did}/${T}.market.rfp/gone1`;
	w.repo.set(vm, { uri: vm, value: {} });
	// made in this host by this module; rfp was deleted by the requester (not in the ledger)
	await ctx.storage.managed_records.put(vm, { module: await moduleOwner({ module: new Uint8Array(wasm) }) });
	const ran = await agentsAlone(w, async () => ({ vm: { uri: vm }, rfp: { uri: rfp }, accept: null, receipt: null, errors: ["no acceptable bid"] }));
	assert.equal(ran.result.status, "failed");
	assert.ok(ran.result.deleted.includes(vm));
	assert.ok(!w.repo.has(vm));
	const record = launchRecord(w);
	assert.match(record.error, /no VM: the contract ended without a receipt/);
	assert.doesNotMatch(record.error, /deleteRecord/);
});

test("caller: a host serving many accounts records the launcher; status and shutdown answer only that caller", async () => {
	const w = world({ did });
	const alice = "did:plc:alice000000000000000000";
	const ran = await run(w, "launch", { ...LAUNCH_INPUT, caller: alice });
	assert.equal(w.repo.get(ran.result.uri).value.caller, alice);
	assert.equal((await run(w, "status", { rkey: ran.result.rkey, caller: alice })).result.record.caller, alice);
	for (const entry of ["status", "shutdown"]) {
		const other = await run(w, entry, { rkey: ran.result.rkey, caller: "did:plc:mallory0000000000000000", poll_ms: 1 });
		assert.equal(other.error, `403 Forbidden: launch ${ran.result.rkey} belongs to another caller`);
	}
	assert.equal(w.calls.filter((c) => c.nsid === "chat.atirc.revoke").length, 0, "nothing torn down for another caller");
	assert.match((await run(w, "launch", { ...LAUNCH_INPUT, caller: "alice" })).error, /caller must be a DID/);
	const down = await run(w, "shutdown", { rkey: ran.result.rkey, caller: alice, delete: true, poll_ms: 1 });
	assert.equal(down.result.status, "stopped");
});

test("the agent's run script takes ATIRC_TOKEN from the unit's environment, never from argv", async () => {
	const w = world({ did });
	await run(w, "launch", LAUNCH_INPUT);
	const command = w.calls.find((c) => c.nsid === "com.fedfork.exec.iroh.connect").input.command;
	const script = Buffer.from(command.match(/echo (\S+) \| base64 -d > \/usr\/local\/bin\/atirc-agent-run\.sh/)[1], "base64").toString();
	assert.match(script, /exec runuser -u "\$\{nick\}" -- env HOME=/);
	assert.doesNotMatch(script, /ATIRC_TOKEN=/);
	const unit = Buffer.from(command.match(/echo (\S+) \| base64 -d > \/etc\/systemd\/system\/atirc-agent\.service/)[1], "base64").toString();
	assert.match(unit, /EnvironmentFile=\/etc\/atirc\/atirc\.env/);
	assert.match(command, /umask 077; install -d -m 0700 \/etc\/atirc; IFS= read -r t; printf 'ATIRC_TOKEN=%s\\n'/, "the env file is 0600");
});

// serve.json: the ingress module, a fake relay socket, the three methods forwarded to handlers (here JS;
// the Rust agents-service supplies them as closures and checks service auth in them).
test("serve.json forwards launchAgent / getAgent / shutdownAgent to the host's handlers and answers on the relay", async () => {
	const { createDidKey, didkeyCall } = await import("../../../atlibs/javascript/src/didkey/index.js");
	const SUB = "com.fedproxy.temp.xrpc.subscribe";
	const queue = [];
	const sent = [];
	let wake;
	let closed = false;
	const relay = {
		send: async (data) => {
			sent.push(data);
			if (sent.filter((f) => f.requestId !== "k").length === 4) (closed = true), wake?.();
		},
		close: () => ((closed = true), wake?.()),
		async *[Symbol.asyncIterator]() {
			while (!closed) {
				if (queue.length) yield queue.shift();
				else await new Promise((resolve) => (wake = resolve));
			}
		},
	};
	const push = (frame) => (queue.push({ $type: `${SUB}#${frame.kind}`, ...frame }), wake?.());
	const req = (requestId, method, nsid, extra) => push({ kind: "request", requestId, method, path: `/xrpc/${nsid}`, params: {}, headers: { authorization: "Bearer jwt" }, ...extra });
	const key = await createDidKey();
	const inner = async (nsid, { subscribe }) => {
		assert.ok(subscribe, nsid);
		push({ kind: "registered", subdomain: "did-key-x" });
		req("k", "GET", "com.publicdomainrelay.temp.keepalive.ping");
		req("l", "POST", "com.fedfork.agents.launchAgent", { body: { name: "a" } });
		req("g", "GET", "com.fedfork.agents.getAgent", { params: { rkey: "3m" } });
		req("s", "POST", "com.fedfork.agents.shutdownAgent", { body: { rkey: "3m" } });
		req("n", "POST", "com.fedfork.agents.nope");
		return relay;
	};
	const call = didkeyCall(key, inner, { services: ["https://xrpc.fedproxy.com"], fetch: async () => new Response(JSON.stringify({ nonce: Buffer.from(crypto.getRandomValues(new Uint8Array(64))).toString("base64") })) });
	const seen = [];
	const handler = (name) => async (input, { caller, reason }) => {
		seen.push([name, caller, input.method, input.authorization, input.body ?? null]);
		if (name === "getAgent") throw new Error("401 AuthenticationRequired: issuer did:plc:x is not allowed");
		return { rkey: "3m", status: name === "launchAgent" ? "launching" : "stopped", reason };
	};
	const comp = await loadAgents(ctx, "./serve.json");
	const ran = await runCompose(ctx, comp, "ingress", "serve", { relay: "https://xrpc.fedproxy.com" }, {
		callFor: () => call,
		handlers: { launchAgent: handler("launchAgent"), getAgent: handler("getAgent"), shutdownAgent: handler("shutdownAgent") },
	});
	const answers = Object.fromEntries(sent.filter((f) => f.requestId).map((f) => [f.requestId, [f.status, f.body]]));
	assert.deepEqual(answers.l, [200, { rkey: "3m", status: "launching", reason: "a caller on the internet asks to launch an agent" }]);
	assert.deepEqual(answers.g, [401, { error: "AuthenticationRequired", message: "issuer did:plc:x is not allowed" }]);
	assert.equal(answers.s[0], 200);
	assert.equal(answers.n[0], 404);
	assert.deepEqual(seen.map((s) => s.slice(0, 4)), [
		["launchAgent", "ingress", "POST", "Bearer jwt"],
		["getAgent", "ingress", "GET", "Bearer jwt"],
		["shutdownAgent", "ingress", "POST", "Bearer jwt"],
	]);
	assert.ok(ran.graph.filter((g) => g.handler).every((g) => g.from === "ingress"));
});

test("compose.ts and serve.ts write composition.json and serve.json", async () => {
	for (const [ts, json] of [["compose.ts", "composition.json"], ["serve.ts", "serve.json"]]) {
		assert.equal(execFileSync(process.execPath, [new URL(ts, import.meta.url).pathname], { encoding: "utf8" }), await readFile(new URL(json, import.meta.url), "utf8"), ts);
	}
});