// node --test modules/vm-mcp/test.mjs  (after `cargo wasm` here and in ../atlibs-iroh, and the pinned modules
// vendored in compositions/ (README: Test); ../atlibs: its JS host)
// vm-mcp served as MCP (atlibs actionsServer) over the requester's fake market: through the composition
// (iroh) and alone with its hook defaults (webrtc, as the agent page runs it). The model's tools take a
// short VM id; the kept list is one record, read at its key.
import assert from "node:assert/strict";
import { readFileSync } from "node:fs";
import { test } from "node:test";

import { bindParameters, loadComposition, runCompose } from "../../../atlibs/javascript/src/compose/index.js";
import { actionsServer } from "../../../atlibs/javascript/src/mcp/index.js";
import { did } from "../../../atlibs/rust/src/harness.mjs";
import { ENDPOINT, HAT7, HAT9, market } from "../requester/fake-market.mjs";

const KEPT = "computer.socialweb.vm.kept";
// A module: this repo's modules/, or atlibs-iroh's; the requester's policy default is policy-open here
// (tangled-vouch's own tests cover it).
const lib = (key) => {
	const name = key === "policy-requester-tangled-vouch" ? "policy-open" : key;
	const dir = name === "iroh-user-data" ? "../../../atlibs-iroh/" : "../../";
	const spec = JSON.parse(readFileSync(new URL(`${dir}modules/${name}/module.json`, import.meta.url), "utf8"));
	return { ...spec, name: key, module: readFileSync(new URL(`${dir}target/wasm32-unknown-unknown/release/${name.replaceAll("-", "_")}.wasm`, import.meta.url)) };
};
// A module as the composition pins it: the published (container) build `atlibs add` vendored in compositions/.
const vendored = (key) => {
	const at = (ext) => new URL(`../../compositions/modules/${key}.${ext}`, import.meta.url);
	return { ...JSON.parse(readFileSync(at("json"), "utf8")), name: key, module: readFileSync(at("wasm")) };
};
const instances = new Map();
const 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));
		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))));
	},
};

// The user's PDS (the kept record, putRecord compare-and-swap on its CID) over the fake market, and the
// guest agent's MCP shell behind the host's iroh: and webrtc: fetchers. put(input) may answer a putRecord first.
function world({ bidders = [HAT7], put, refuses = false } = {}) {
	const m = market({ bidders, refuses });
	const kept = new Map();
	let n = 0;
	const calls = [];
	const respond = (call) => {
		calls.push(call);
		if (call.params?.collection === KEPT || call.input?.collection === KEPT) {
			if (call.nsid === "com.atproto.repo.getRecord") {
				const r = kept.get(call.params.rkey);
				return r ? { ok: true, output: { uri: `at://${did}/${KEPT}/${call.params.rkey}`, ...r } } : { ok: false, error: "RecordNotFound: Could not locate record" };
			}
			assert.equal(call.nsid, "com.atproto.repo.putRecord");
			const early = put?.(call.input);
			if (early) return early;
			if ((call.input.swapRecord ?? null) !== (kept.get(call.input.rkey)?.cid ?? null)) return { ok: false, error: "InvalidSwap: Record was at bafy-other" };
			kept.set(call.input.rkey, { cid: `bafykept${++n}`, value: call.input.record });
			return { ok: true, output: { uri: `at://${did}/${KEPT}/${call.input.rkey}`, cid: `bafykept${n}` } };
		}
		const out = m.respond(call);
		// a live receipt names the bidder's submitEvent: where a kept VM's release goes
		if (out.ok && call.params?.collection?.endsWith("market.receipt")) return { ok: true, output: { ...out.output, value: { ...out.output.value, submitEvent: ENDPOINT } } };
		return out;
	};
	const call = async (nsid, request) => {
		const out = await respond({ nsid, ...request });
		if (out?.ok === false) throw new Error(out.error ?? "call failed");
		return out?.output;
	};
	const shell = [];
	const fetch = async ({ url, body }) => {
		const { params } = JSON.parse(body);
		shell.push({ url, ...params });
		assert.equal(params.name, "exec");
		const out = { exitCode: 0, stdout: `ran ${params.arguments.command}\n`, stderr: "" };
		return { status: 200, headers: {}, body: JSON.stringify({ jsonrpc: "2.0", id: 1, result: { content: [{ type: "text", text: JSON.stringify(out) }], structuredContent: out } }) };
	};
	const records = new Map();
	const 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) };
	const ctx = { wasm, storage: { managed_records } };
	return { ctx, call, calls, kept, shell, fetchers: { "iroh:": fetch, "webrtc:": fetch } };
}

const http = { "iroh://*": {}, "webrtc://*": {} };
const budget = { calls: Infinity, wallMs: Infinity, events: Infinity };

// The composition, served as MCP: each tools/call one runCompose of instance vm.
async function viaComposition(w, inputs = {}) {
	const manifest = JSON.parse(readFileSync(new URL("../../compositions/socialweb-computer-vm-mcp.composition.json", import.meta.url), "utf8"));
	// fast polls; the vouch policy is policy-open here (unpinned): tangled-vouch's own tests judge bidders
	for (const i of manifest.instances) i.input = { ...i.input, ...(i.name === "machine" && { poll_ms: 1 }) };
	Object.assign(manifest.instances.find((i) => i.name === "vouch"), { module: "policy-open", pin: undefined });
	const bound = bindParameters(manifest, { inputs, host: { sshPublicKey: "ssh-ed25519 AAAA test" } });
	const comp = await loadComposition(w.ctx, bound, { load: async (key) => (key === "policy-open" ? lib(key) : vendored(key)), provides: ["com.fedfork.exec.iroh.connect"] });
	const run = (mod, entry, args) => runCompose(w.ctx, comp, "vm", entry, { ...args, did }, { callFor: () => w.call, fetchers: w.fetchers, http, budget });
	return actionsServer({ name: "vm-mcp", modules: [vendored("vm-mcp")], run, confirm: async () => true });
}

// vm-mcp alone, its hooks served by their defaults (requester, process-mcp), as the agent page runs a lib.
function viaDefaults(w, input = {}) {
	const host = { "vm-mcp": { call: w.call, callFor: () => w.call, load: async (key) => lib(key), fetchers: w.fetchers, http, budget, input: { did, via: "webrtc", ...input } } };
	return actionsServer({ name: "vm-mcp", modules: [lib("vm-mcp")], ctx: w.ctx, host, confirm: async () => true });
}

const tool = async (server, name, args = {}) => {
	const reply = await server.message({ jsonrpc: "2.0", id: 1, method: "tools/call", params: { name, arguments: args } });
	const r = reply.result;
	return r.isError ? { error: r.content[0].text } : r.structuredContent;
};

test("vm-mcp: four tools, typed for a model; vm_list alone is read-only", async () => {
	const server = await viaComposition(world());
	const { tools } = (await server.message({ jsonrpc: "2.0", id: 1, method: "tools/list" })).result;
	assert.deepEqual(tools.map((t) => [t.name, t.annotations.readOnlyHint]), [["vm_create", false], ["vm_list", true], ["vm_exec", false], ["vm_release", false], ["vm_forget", false]]);
	const exec = tools.find((t) => t.name === "vm_exec").inputSchema;
	assert.deepEqual([Object.keys(exec.properties), exec.required], [["vm", "command", "stdin", "timeout_s"], ["vm", "command"]]);
	assert.deepEqual(tools.find((t) => t.name === "vm_create").inputSchema.required, undefined, "vm_create needs no argument");
});

for (const [how, serve, scheme] of [["the composition (iroh)", viaComposition, "iroh:"], ["its hook defaults (webrtc)", viaDefaults, "webrtc:"]]) {
	test(`vm-mcp through ${how}: create, list, exec, release; one record at its key, never listRecords`, async () => {
		const w = world();
		const server = await serve(w);
		const vm = await tool(server, "vm_create", { cpus: 2, hold_minutes: 30 });
		assert.equal(vm.error, undefined, vm.error);
		assert.match(vm.vm, /^vm-[0-9a-f]{8}$/);
		assert.equal(vm.bidder, HAT7);
		assert.ok(vm.shell.startsWith(`${scheme}///`), vm.shell);
		assert.equal(Date.parse(vm.expires) - Date.parse(vm.created), 30 * 60_000);
		assert.deepEqual(Object.keys(vm).sort(), ["bidder", "cpus", "created", "disk", "expired", "expires", "mem", "shell", "vm"]);
		const spec = w.calls.find((c) => c.nsid === "com.atproto.repo.createRecord" && c.input.collection.endsWith("compute.vm")).input.record;
		assert.equal(spec.cpus, 2);
		assert.equal(spec.expiresAt, vm.expires, "the bidder ends it at compute.vm expiresAt");

		const listed = await tool(server, "vm_list");
		assert.deepEqual(listed.vms.map((v) => [v.vm, v.expired]), [[vm.vm, false]]);
		const stored = w.kept.get("self").value.vms[0];
		assert.ok(stored.receipt.uri && stored.records.length, "the record keeps what releases it");

		const ran = await tool(server, "vm_exec", { vm: vm.vm, command: "uname -a", timeout_s: 5 });
		assert.deepEqual(ran, { exit: 0, stdout: "ran uname -a\n", stderr: "" });
		assert.deepEqual([w.shell[0].url, w.shell[0].arguments], [vm.shell, { command: "uname -a", timeoutMs: 5000 }]);
		assert.match((await tool(server, "vm_exec", { vm: "vm-nope", command: "id" })).error, /no VM vm-nope here \(yours: vm-/);

		const released = await tool(server, "vm_release", { vm: vm.vm });
		assert.deepEqual(released, { vm: vm.vm, released: true, terminated: true });
		assert.ok(w.calls.some((c) => c.nsid.endsWith("submitEvent") && c.service === ENDPOINT), "vm.delete went to the bidder");
		assert.ok(w.calls.some((c) => c.nsid === "com.atproto.repo.deleteRecord" && c.input.collection.endsWith("market.rfp")), "the rent's records were deleted");
		assert.deepEqual((await tool(server, "vm_list")).vms, []);
		assert.ok(!w.calls.some((c) => c.nsid === "com.atproto.repo.listRecords"), "never listRecords");
	});
}

test("vm-mcp: no bid keeps nothing and says what to change; the host's limits bound a create", async () => {
	const w = world();
	const server = viaDefaults(w, { max_vms: 1, max_cpus: 4 });
	const none = await tool(server, "vm_create", { bidders: [HAT9] });
	assert.match(none.error, /^no VM \(.*no acceptable bid.*\); nothing is kept\. Try vm_create again/);
	assert.deepEqual(w.kept.get("self").value.vms, [], "nothing kept");
	assert.match((await tool(server, "vm_create", { cpus: 5 })).error, /cpus must be an integer from 1 to 4/);
	assert.match((await tool(server, "vm_create", { bidders: ["hat-7"] })).error, /bidders must be 1 to 10 bidder DIDs/);
	assert.equal((await tool(server, "vm_create")).error, undefined, "the default bidders: hat-7 bids");
	assert.match((await tool(server, "vm_create")).error, /you keep 1 VMs, at most 1: vm_release one first/);
});

test("vm-mcp: a lost putRecord race re-reads; a VM it cannot record is released", async () => {
	let races = 2; // the probe (a lost swap still proves the write), then the record
	const w = world({ put: () => (races-- > 0 ? { ok: false, error: "InvalidSwap: Record was at bafy-other" } : undefined) });
	const vm = await tool(viaDefaults(w), "vm_create");
	assert.equal(vm.error, undefined, vm.error);
	assert.equal(w.kept.get("self").value.vms.length, 1);

	const scopeless = world({ put: () => ({ ok: false, error: "missing scope repo:computer.socialweb.vm.kept" }) });
	assert.match((await tool(viaDefaults(scopeless), "vm_create")).error, /^cannot keep your VM list \(missing scope .*\); nothing was rented\. Sign in again/);
	assert.ok(!scopeless.calls.some((c) => c.nsid.endsWith("submitRfp")), "no rent without the scope");

	let puts = 0;
	const w2 = world({ put: () => (puts++ > 0 ? { ok: false, error: "AuthRequired: scope" } : undefined) });
	const lost = await tool(viaDefaults(w2), "vm_create");
	assert.match(lost.error, /^the VM came up but could not be recorded \(AuthRequired: scope\); it was released$/);
	assert.ok(w2.calls.some((c) => c.nsid.endsWith("submitEvent")));
});

test("vm-mcp: a VM whose bidder refuses the release stays kept; vm_forget drops it after one try and says so", async () => {
	const w = world({ refuses: true });
	const server = viaDefaults(w);
	const vm = await tool(server, "vm_create");
	assert.equal(vm.error, undefined, vm.error);
	assert.match((await tool(server, "vm_release", { vm: vm.vm })).error, new RegExp(`^${vm.vm} is still kept: .*submitEvent.*vm_forget drops it`));
	assert.deepEqual((await tool(server, "vm_list")).vms.map((v) => v.vm), [vm.vm]);
	const forgot = await tool(server, "vm_forget", { vm: vm.vm });
	assert.deepEqual([forgot.vm, forgot.forgotten, forgot.released, forgot.terminated], [vm.vm, true, false, false]);
	assert.match(forgot.note, /did not go through .*until its expiry/);
	assert.deepEqual((await tool(server, "vm_list")).vms, []);
	// a release that goes through: forget is a release
	const w2 = world();
	const s2 = viaDefaults(w2);
	const vm2 = await tool(s2, "vm_create");
	assert.deepEqual(await tool(s2, "vm_forget", { vm: vm2.vm }), { vm: vm2.vm, released: true, terminated: true, forgotten: true });
});