import assert from "node:assert/strict";
import { createPublicKey, verify } from "node:crypto";
import { readFile } from "node:fs/promises";
import { test } from "node:test";

import { createDidKey, didkeyCall } from "../../../atlibs/javascript/src/didkey/index.js";
import { runModule } from "../../../atlibs/javascript/src/run-module/index.js";
import { load, now } from "../../../atlibs/rust/src/harness.mjs";

const { module, run } = await load(new URL("../../target/wasm32-unknown-unknown/release/ingress.wasm", import.meta.url));
const manifest = JSON.parse(await readFile(new URL("module.json", import.meta.url), "utf8"));
const RELAY = "https://xrpc.fedproxy.com";
const T = "com.fedproxy.temp.xrpc.subscribe";
const did = "did:key:zDnaeExample";
const exp = Math.floor(Date.parse(now) / 1000);

// One handler call: the output, with state carried by the caller.
// abi::Handle leaves out empty parts; fill them in so tests can index them.
const step = (state, event, input = {}) => ({ calls: [], send: [], timers: [], cancel: [], done: false, ...run("serve", { now, state, event, ...input }) });
const ok = (id, output) => ({ type: "result", id, ok: true, output });
const frame = (kind, fields) => ({ type: "message", id: "relay", data: { $type: `${T}#${kind}`, ...fields } });

function connect(input) {
	let out = step(null, { type: "start" }, input);
	const starts = out;
	out = step(out.state, ok("auth-nonce", { did, token: "t1" }), input);
	const nonceCall = out.calls[0];
	out = step(out.state, ok("nonce", { nonce: "bm9uY2U=" }), input);
	const signCall = out.calls[0];
	out = step(out.state, ok("sign", { did, signature: "c2ln" }), input);
	assert.deepEqual(out.calls, [], "waits for the subscribe token");
	out = step(out.state, ok("auth-sub", { did, token: "t2" }), input);
	return { starts, nonceCall, signCall, subscribe: out.calls[0], state: out.state };
}

test("module has no imports and exports serve", () => {
	assert.equal(WebAssembly.Module.imports(module).length, 0);
	assert.ok(WebAssembly.Module.exports(module).some((e) => e.name === "serve"));
});

test("registration: two tokens in parallel, nonce, signature, then one subscription", () => {
	const { starts, nonceCall, signCall, subscribe } = connect({ lifetimeMs: "120000" });
	assert.deepEqual(starts.calls, [
		{ id: "auth-nonce", nsid: "com.fedfork.exec.didkey.getServiceAuth", params: { aud: "did:web:xrpc.fedproxy.com", lxm: "com.fedproxy.temp.xrpc.getRegistrationNonce" } },
		{ id: "auth-sub", nsid: "com.fedfork.exec.didkey.getServiceAuth", params: { aud: "did:web:xrpc.fedproxy.com", lxm: T, exp: String(exp + 180) } },
	]);
	assert.deepEqual(starts.timers, [{ id: "lifetime", ms: 120000 }]);
	assert.deepEqual(nonceCall, { id: "nonce", nsid: "com.fedproxy.temp.xrpc.getRegistrationNonce", service: RELAY, input: { key: did, signatures: [] } });
	assert.deepEqual(signCall, { id: "sign", nsid: "com.fedfork.exec.didkey.sign", input: { data: "bm9uY2U=" } });
	assert.equal(subscribe.subscribe, true);
	assert.equal(subscribe.service, RELAY);
	assert.equal(subscribe.overflow, "close", "a dropped relay frame would be a lost request");
	assert.deepEqual(JSON.parse(subscribe.params.registration), {
		$type: "com.fedproxy.temp.xrpc.registration",
		key: did,
		nonce: "bm9uY2U=",
		signatures: [{ key: did, signature: "c2ln" }],
	});
	assert.deepEqual([subscribe.params.did, subscribe.params.service_auth], [did, "t2"]);
	for (const call of [...starts.calls, nonceCall, signCall, subscribe]) assert.ok(manifest.calls.includes(call.nsid), call.nsid);
});

test("without a lifetime the subscribe token asks for the hour the host allows", () => {
	assert.equal(step(null, { type: "start" }).calls[1].params.exp, String(exp + 3600));
});

test("a failed registration step ends the run with its error", () => {
	const out = step(step(null, { type: "start" }).state, { type: "result", id: "auth-nonce", ok: false, error: "no key" });
	assert.equal(out.error, "auth-nonce: no key");
});

test("requests are answered on the relay socket; reads run concurrently and time out", () => {
	let { state } = connect();
	let out = step(state, frame("registered", { subdomain: "did-key-zdnaeexample" }));
	assert.equal(out.result.url, "https://did-key-zdnaeexample.xrpc.fedproxy.com/");
	const request = (requestId, path, params = {}) => (out = step(out.state, frame("request", { requestId, method: "GET", path, params, headers: {} }), { deadlineMs: 500 }));
	const sent = () => out.send.map(({ id, data }) => (assert.equal(id, "relay"), [data.requestId, data.status, data.body]));

	request("k", "/xrpc/com.publicdomainrelay.temp.keepalive.ping");
	assert.deepEqual(sent(), [["k", 200, {}]]);
	request("h", "/xrpc/com.example.ingress.hello", { name: "bob" });
	assert.deepEqual(sent(), [["h", 200, { message: "hello bob", did, served: 1, now }]]);
	request("a", "/xrpc/com.example.ingress.describeRepo", { repo: "did:plc:aaa" });
	assert.deepEqual(out.calls, [{ id: "req:a", nsid: "com.atproto.repo.describeRepo", params: { repo: "did:plc:aaa" }, service: "did:plc:aaa" }]);
	assert.deepEqual(out.timers, [{ id: "deadline:a", ms: 500 }]);
	request("b", "/xrpc/com.example.ingress.describeRepo", { repo: "did:web:b.example" });
	assert.deepEqual(out.state.pending, ["a", "b"]);
	request("c", "/xrpc/com.example.ingress.describeRepo", { repo: "alice.example" });
	assert.equal(sent()[0][1], 400);
	request("d", "/xrpc/nope.nope.nope");
	assert.equal(sent()[0][1], 404);
	request("e", "/");
	assert.match(sent()[0][2], /Served by a wasm module/);
	assert.equal(out.send[0].data.contentType, "text/html; charset=utf-8");

	out = step(out.state, ok("req:b", { did: "did:web:b.example", handle: "b.example", collections: ["x.y.z"] }));
	assert.deepEqual(sent(), [["b", 200, { did: "did:web:b.example", handle: "b.example", collections: ["x.y.z"], servedBy: did }]]);
	assert.deepEqual(out.cancel, ["deadline:b"]);
	out = step(out.state, { type: "timer", id: "deadline:a" });
	assert.deepEqual(out.cancel, ["req:a"]);
	assert.equal(sent()[0][1], 504);
	out = step(out.state, ok("req:a", {}));
	assert.deepEqual(out.send, [], "a late result after the deadline is dropped");
	assert.equal(out.result.served, 6);
});

test("routes: /xrpc/<nsid> in routes becomes a call with the request; output is the 200 body, a coded error its status", () => {
	const routes = ["com.fedfork.agents.launchAgent", "com.fedfork.agents.getAgent"];
	let { state } = connect({ routes });
	let out = step(state, frame("registered", { subdomain: "did-key-zdnaeexample" }), { routes });
	const request = (requestId, method, path, extra = {}) =>
		(out = step(out.state, frame("request", { requestId, method, path, params: {}, headers: { authorization: "Bearer jwt" }, ...extra }), { routes: routes.join(" ") }));
	request("l", "POST", "/xrpc/com.fedfork.agents.launchAgent", { body: { name: "a" } });
	assert.deepEqual(out.calls, [{ id: "req:l", nsid: "com.fedfork.agents.launchAgent", input: { method: "POST", params: {}, body: { name: "a" }, authorization: "Bearer jwt" } }]);
	assert.deepEqual(out.timers, [{ id: "deadline:l", ms: 8000 }]);
	request("g", "GET", "/xrpc/com.fedfork.agents.getAgent", { params: { rkey: "r" } });
	request("x", "POST", "/xrpc/com.fedfork.agents.shutdownAgent");
	assert.equal(out.send[0].data.status, 404, "not in routes");
	const result = (id, nsid, answer) => (out = step(out.state, { type: "result", id, nsid, ...answer }, { routes }));
	result("req:l", routes[0], { ok: true, output: { rkey: "3m", status: "launching" } });
	assert.deepEqual([out.send[0].data.status, out.send[0].data.body], [200, { rkey: "3m", status: "launching" }]);
	result("req:g", routes[1], { ok: false, error: "401 AuthenticationRequired: issuer did:plc:x is not allowed" });
	assert.deepEqual([out.send[0].data.status, out.send[0].data.body], [401, { error: "AuthenticationRequired", message: "issuer did:plc:x is not allowed" }]);
	({ state } = connect({ routes }));
	out = step(state, frame("request", { requestId: "z", method: "GET", path: "/xrpc/com.fedfork.agents.getAgent", params: {}, headers: {} }), { routes });
	out = step(out.state, { type: "result", id: "req:z", nsid: routes[1], ok: false, error: "boom" }, { routes });
	assert.deepEqual([out.send[0].data.status, out.send[0].data.body.error], [502, "UpstreamFailure"]);
});

test("forward: every request but keepalives is one call to the host; its output is the response", () => {
	const input = { forward: "true" };
	let out = step(connect(input).state, frame("registered", { subdomain: "did-key-zdnaeexample" }), input);
	const request = (requestId, method, path, extra = {}) =>
		(out = step(out.state, frame("request", { requestId, method, path, params: {}, headers: { accept: "text/html" }, callerDid: "did:plc:caller", ...extra }), input));
	request("k", "GET", "/xrpc/com.publicdomainrelay.temp.keepalive.ping");
	assert.deepEqual([out.calls, out.send[0].data.status], [[], 200], "keepalives are still answered here");
	request("p", "GET", "/about", { params: { q: "1" } });
	assert.deepEqual(out.calls, [{ id: "req:p", nsid: "com.fedfork.exec.ingress.request", input: { method: "GET", path: "/about", params: { q: "1" }, headers: { accept: "text/html" }, body: null, callerDid: "did:plc:caller" } }]);
	request("x", "POST", "/xrpc/com.example.ingress.hello", { body: { a: 1 } });
	assert.equal(out.calls[0].input.path, "/xrpc/com.example.ingress.hello", "fixed paths forward too");
	const result = (id, answer) => (out = step(out.state, { type: "result", id, nsid: "com.fedfork.exec.ingress.request", ...answer }, input));
	result("req:p", { ok: true, output: { status: 201, body: "<h1>hi</h1>", contentType: "text/html; charset=utf-8" } });
	assert.deepEqual(out.send[0].data, { $type: `${T}#response`, requestId: "p", status: 201, body: "<h1>hi</h1>", contentType: "text/html; charset=utf-8" });
	result("req:x", { ok: false, error: "403 Forbidden: no" });
	assert.deepEqual([out.send[0].data.status, out.send[0].data.body], [403, { error: "Forbidden", message: "no" }]);
	assert.ok(manifest.calls.includes("com.fedfork.exec.ingress.request"));

	request("c", "GET", "/old");
	const headers = { location: "/new", "Access-Control-Allow-Origin": "*", "set-cookie": ["a=1", "b=2"], "x-num": 5, "x-mixed": ["a", 1], "": "x" };
	result("req:c", { ok: true, output: { status: 302, body: "", contentType: "text/plain", headers } });
	assert.deepEqual(out.send[0].data.headers, { location: "/new", "Access-Control-Allow-Origin": "*", "set-cookie": ["a=1", "b=2"] }, "strings and string lists pass, the rest drop");
	request("n", "GET", "/n");
	result("req:n", { ok: true, output: { status: 200, body: "x", headers: "nope" } });
	assert.equal("headers" in out.send[0].data, false, "no headers object, no field");
});

test("aliases: one alias frame after every registered, the relay's #result in the run result", () => {
	const input = { aliases: "Mush.at, www.mush.at" };
	let out = step(connect(input).state, frame("registered", { subdomain: "did-key-zdnaeexample" }), input);
	assert.deepEqual(out.send, [{ id: "relay", data: { $type: "com.fedproxy.temp.xrpc.alias", hosts: ["mush.at", "www.mush.at"] } }]);
	assert.equal(out.result.aliases, undefined, "no answer yet");
	const answer = { accepted: ["mush.at"], rejected: { "www.mush.at": "no matching _xrpc.www.mush.at TXT record" } };
	out = step(out.state, { type: "message", id: "relay", data: { $type: "com.fedproxy.temp.xrpc.alias#result", ...answer } }, input);
	assert.deepEqual([out.send, out.result.aliases], [[], answer]);
	assert.deepEqual(out.timers, [{ id: "alias-retry", ms: 60000 }], "a rejected host is asked for again");
	out = step(out.state, { type: "timer", id: "alias-retry" }, input);
	assert.deepEqual(out.send, [{ id: "relay", data: { $type: "com.fedproxy.temp.xrpc.alias", hosts: ["mush.at", "www.mush.at"] } }]);
	const all = { accepted: ["mush.at", "www.mush.at"], rejected: {} };
	out = step(out.state, { type: "message", id: "relay", data: { $type: "com.fedproxy.temp.xrpc.alias#result", ...all } }, input);
	assert.deepEqual([out.timers, out.result.aliases], [[], all], "all accepted: no more retries");
	out = step(out.state, { type: "closed", id: "relay", error: "1006" }, input);
	assert.equal(out.result.aliases, undefined, "the relay drops aliases with the connection");
	let pending = step(connect(input).state, frame("registered", { subdomain: "did-key-zdnaeexample" }), input);
	pending = step(pending.state, { type: "message", id: "relay", data: { $type: "com.fedproxy.temp.xrpc.alias#result", ...answer } }, input);
	pending = step(pending.state, { type: "closed", id: "relay", error: "1006" }, input);
	assert.ok(pending.cancel.includes("alias-retry"), "a pending retry goes with the connection");
	assert.deepEqual(step(pending.state, { type: "timer", id: "alias-retry" }, input).send, [], "no retry while disconnected");
	out = step(out.state, frame("registered", { subdomain: "did-key-zdnaeexample" }), { aliases: ["mush.at"] });
	assert.deepEqual(out.send[0].data.hosts, ["mush.at"], "sent again on the reconnect");
	assert.deepEqual(step(connect().state, frame("registered", { subdomain: "x" })).send, [], "no aliases, no frame");
});

test("a caller subscription on our host gets ticks, then a close; unknown ones are refused", () => {
	let out = step(connect().state, frame("subscribe", { subscriptionId: "s1", nsid: "com.example.ingress.ticks", params: { n: "2", ms: "250" } }));
	assert.deepEqual(out.send.map((s) => s.data.$type), [`${T}#subscriptionOpen`]);
	assert.deepEqual(out.timers, [{ id: "tick:s1", ms: 250 }]);
	out = step(out.state, { type: "timer", id: "tick:s1" });
	assert.deepEqual(out.send[0].data.message.left, 1);
	assert.deepEqual(out.timers, [{ id: "tick:s1", ms: 250 }]);
	out = step(out.state, { type: "timer", id: "tick:s1" });
	assert.deepEqual(out.send.map((s) => [s.data.$type, s.data.code]), [[`${T}#subscriptionEvent`, undefined], [`${T}#subscriptionClose`, 1000]]);
	assert.deepEqual(out.state.ticks, {});

	out = step(out.state, frame("subscribe", { subscriptionId: "s2", nsid: "com.example.ingress.ticks", params: {} }));
	out = step(out.state, frame("subscriptionCancel", { subscriptionId: "s2" }));
	assert.deepEqual(out.cancel, ["tick:s2"]);
	out = step(out.state, frame("subscribe", { subscriptionId: "t", nsid: "com.fedproxy.temp.xrpc.tunnel", params: {} }));
	assert.equal(out.send[0].data.code, 1003);
});

test("a closed relay re-registers with backoff, three times, then fails", () => {
	let out = step(connect().state, frame("request", { requestId: "a", path: "/xrpc/com.example.ingress.describeRepo", params: { repo: "did:plc:aaa" } }));
	out = step(out.state, { type: "closed", id: "relay", error: "1011 unresponsive" });
	assert.deepEqual(out.cancel, ["req:a", "deadline:a"]);
	assert.deepEqual(out.timers, [{ id: "reconnect", ms: 2000 }]);
	out = step(out.state, { type: "timer", id: "reconnect" });
	assert.deepEqual(out.calls.map((c) => c.id), ["auth-nonce", "auth-sub"]);
	out = step(out.state, { type: "closed", id: "relay" });
	out = step(out.state, { type: "closed", id: "relay" });
	assert.deepEqual(out.timers, [{ id: "reconnect", ms: 8000 }]);
	assert.match(step(out.state, { type: "closed", id: "relay", error: "gone" }).error, /relay closed 3 times: gone/);
});

test("quietMs leaves keepalives unanswered on the first connection only, so the relay evicts it", () => {
	const input = { quietMs: "120000" };
	const ping = (state) => step(state, frame("request", { requestId: "k", path: "/xrpc/com.publicdomainrelay.temp.keepalive.ping", params: {} }), input);
	let out = step(connect().state, frame("registered", { subdomain: "did-key-x" }), input);
	assert.equal(out.state.quietUntil, exp + 120);
	assert.deepEqual(ping(out.state).send, [], "quiet: no pong");
	assert.equal(step(out.state, frame("request", { requestId: "h", path: "/xrpc/com.example.ingress.hello", params: {} }), input).send[0].data.status, 200, "requests are still served");
	out = step(out.state, { type: "closed", id: "relay", error: "1011 subscriber unresponsive for 90s" }, input);
	assert.equal(out.result.drops, 1);
	out = step(out.state, frame("registered", { subdomain: "did-key-x" }), input);
	assert.equal(ping(out.state).send[0].data.status, 200, "after a drop it answers again");
	assert.equal(out.state.reconnects, 0);
});

test("maxRequests (a UI string) ends the run once nothing is in flight", () => {
	let out = step(connect().state, frame("request", { requestId: "a", path: "/xrpc/com.example.ingress.describeRepo", params: { repo: "did:plc:aaa" } }), { maxRequests: "2" });
	out = step(out.state, frame("request", { requestId: "h", path: "/xrpc/com.example.ingress.hello", params: {} }), { maxRequests: "2" });
	assert.equal(out.done, false, "the read is still in flight");
	out = step(out.state, ok("req:a", {}), { maxRequests: "2" });
	assert.equal(out.done, true);
});

// The real reactor and didkeyCall, with a fake relay socket: the module, the host did:key and
// concurrency end to end, with no network.
function fakeRelay() {
	const queue = [];
	let wake;
	let closed = false;
	const sent = [];
	const stream = {
		sent,
		push: (data) => (queue.push(data), wake?.()),
		send: async (data) => void sent.push(data),
		close: () => ((closed = true), wake?.()),
		async *[Symbol.asyncIterator]() {
			while (!closed) {
				if (queue.length) yield queue.shift();
				else await new Promise((resolve) => (wake = resolve));
			}
		},
	};
	return stream;
}

test("runModule + didkeyCall: registers, serves concurrent requests, cancels a slow read, ends at maxRequests", async () => {
	const key = await createDidKey();
	const pub = createPublicKey({ key: { kty: "EC", crv: "P-256", x: key.jwk.x, y: key.jwk.y }, format: "jwk" });
	const relay = fakeRelay();
	const reads = {};
	const aborted = [];
	const nonce = Buffer.from(crypto.getRandomValues(new Uint8Array(64))).toString("base64");
	const inner = async (nsid, { params, service, subscribe, signal }) => {
		if (subscribe) {
			assert.equal(`${service}|${nsid}`, `${RELAY}|${T}`);
			const registration = JSON.parse(params.registration);
			assert.equal(registration.key, key.did);
			assert.ok(verify("sha256", Buffer.from(nonce, "base64"), { key: pub, dsaEncoding: "ieee-p1363" }, Buffer.from(registration.signatures[0].signature, "base64")));
			const claims = JSON.parse(Buffer.from(params.service_auth.split(".")[1], "base64url"));
			assert.deepEqual([claims.iss, claims.aud, claims.lxm], [key.did, "did:web:xrpc.fedproxy.com", T]);
			const req = (requestId, path, query = {}) => relay.push({ $type: `${T}#request`, requestId, method: "GET", path, params: query, headers: {} });
			relay.push({ $type: `${T}#registered`, subdomain: key.did.replaceAll(":", "-").toLowerCase() });
			req("slow", "/xrpc/com.example.ingress.describeRepo", { repo: "did:plc:slow" });
			req("fast", "/xrpc/com.example.ingress.describeRepo", { repo: "did:plc:fast" });
			req("k", "/xrpc/com.publicdomainrelay.temp.keepalive.ping");
			req("hi", "/xrpc/com.example.ingress.hello", { name: "web" });
			relay.push({ $type: `${T}#subscribe`, subscriptionId: "s", nsid: "com.example.ingress.ticks", params: { n: "2", ms: "200" } });
			return relay;
		}
		assert.equal(nsid, "com.atproto.repo.describeRepo");
		reads[params.repo] = true;
		return new Promise((resolve, reject) => {
			let settled = false;
			signal.addEventListener("abort", () => settled || (aborted.push(params.repo), reject(signal.reason)));
			if (params.repo === "did:plc:fast") setTimeout(() => ((settled = true), resolve({ did: params.repo, handle: "fast.example", collections: [] })), 20);
		});
	};
	const call = didkeyCall(key, inner, {
		services: manifest.didkey.services,
		fetch: async (url, init) => {
			assert.equal(url, `${RELAY}/xrpc/com.fedproxy.temp.xrpc.getRegistrationNonce`);
			assert.match(init.headers.authorization, /^Bearer /);
			return new Response(JSON.stringify({ nonce, signatures: [] }));
		},
	});
	const ctx = { wasm: { run: async ({ entry, input }) => run(entry, input) } };
	const ran = await runModule(ctx, { module, calls: manifest.calls }, "serve", { maxRequests: 3, deadlineMs: 300 }, call, { budget: { wallMs: 10_000 } });

	assert.equal(ran.error, undefined);
	assert.deepEqual(ran.result, { did: key.did, url: `https://${key.did.replaceAll(":", "-").toLowerCase()}.xrpc.fedproxy.com/`, served: 3, drops: 0 });
	assert.deepEqual(Object.keys(reads).sort(), ["did:plc:fast", "did:plc:slow"], "both reads were in flight together");
	assert.deepEqual(aborted, ["did:plc:slow"]);
	const responses = Object.fromEntries(relay.sent.filter((f) => f.$type === `${T}#response`).map((f) => [f.requestId, f.status]));
	assert.deepEqual(responses, { k: 200, hi: 200, fast: 200, slow: 504 });
	assert.deepEqual(
		relay.sent.filter((f) => f.subscriptionId === "s").map((f) => f.$type.split("#")[1]),
		["subscriptionOpen", "subscriptionEvent", "subscriptionEvent", "subscriptionClose"],
	);
	assert.ok(ran.calls.some((c) => c.id === "req:slow" && c.cancelled));
});

test("runModule: aliases are sent up the relay socket and the #result reaches the run result", async () => {
	const relay = fakeRelay();
	const sub = "did-key-zdnaeexample";
	const call = async (nsid, { subscribe }) => {
		if (nsid === "com.fedfork.exec.didkey.getServiceAuth") return { did, token: "t" };
		if (nsid === "com.fedproxy.temp.xrpc.getRegistrationNonce") return { nonce: "bm9uY2U=" };
		if (nsid === "com.fedfork.exec.didkey.sign") return { signature: "c2ln" };
		assert.ok(subscribe);
		relay.push({ $type: `${T}#registered`, subdomain: sub });
		relay.send = async (data) => {
			relay.sent.push(data);
			if (data.$type === "com.fedproxy.temp.xrpc.alias") {
				relay.push({ $type: "com.fedproxy.temp.xrpc.alias#result", accepted: data.hosts, rejected: {} });
				relay.push({ $type: `${T}#request`, requestId: "r", method: "GET", path: "/xrpc/com.example.ingress.hello", params: {}, headers: {} });
			}
		};
		return relay;
	};
	const ctx = { wasm: { run: async ({ entry, input }) => run(entry, input) } };
	const ran = await runModule(ctx, { module, calls: manifest.calls }, "serve", { maxRequests: 1, aliases: ["mush.at"] }, call, { budget: { wallMs: 5000 } });
	assert.equal(ran.error, undefined);
	assert.deepEqual(ran.result.aliases, { accepted: ["mush.at"], rejected: {} });
	assert.deepEqual(relay.sent.map((f) => f.$type), ["com.fedproxy.temp.xrpc.alias", `${T}#response`]);
});