import assert from "node:assert/strict";
import http from "node:http";
import { after, before, describe, test } from "node:test";
import { once } from "node:events";
import WebSocket, { WebSocketServer } from "ws";
import {
	createPreviewRuntime,
	fetchPreviewUpstream,
	runtimeConfigFromEnv,
	validatePreviewResolution,
} from "./server.mjs";

const CAPABILITY = "A".repeat(43);
const OTHER_CAPABILITY = "B".repeat(43);
const SANDBOX_ID = "job-123";
const SANDBOX_TOKEN = "a".repeat(64);
const JOBS_TOKEN = `hf_${"j".repeat(40)}`;

function discardedResponse(onCancel) {
	return new Response(
		new ReadableStream({
			start(controller) {
				controller.enqueue(new TextEncoder().encode("transient edge miss"));
			},
			cancel() {
				onCancel();
			},
		}),
		{ status: 404 }
	);
}

test("retries transient upstream 404s for safe GET and HEAD requests", async (context) => {
	for (const method of ["GET", "HEAD"]) {
		await context.test(method, async () => {
			let calls = 0;
			let cancellations = 0;
			const fetchImpl = async () => {
				calls += 1;
				return calls === 1
					? discardedResponse(() => {
							cancellations += 1;
						})
					: new Response(method === "HEAD" ? null : "preview ready", { status: 200 });
			};

			const response = await fetchPreviewUpstream(
				fetchImpl,
				"https://job-123--49983.hf.jobs/v1/proxy/3000/",
				{ method },
				[0]
			);

			assert.equal(response.status, 200);
			assert.equal(calls, 2);
			assert.equal(cancellations, 1);
			if (method === "GET") assert.equal(await response.text(), "preview ready");
		});
	}
});

test("bounds persistent upstream 404 retries and returns the final application response", async () => {
	let calls = 0;
	let cancellations = 0;
	const fetchImpl = async () => {
		calls += 1;
		return calls < 3
			? discardedResponse(() => {
					cancellations += 1;
				})
			: new Response("application route not found", { status: 404 });
	};

	const response = await fetchPreviewUpstream(
		fetchImpl,
		"https://job-123--49983.hf.jobs/v1/proxy/3000/missing",
		{ method: "GET" },
		[0, 0]
	);

	assert.equal(response.status, 404);
	assert.equal(await response.text(), "application route not found");
	assert.equal(calls, 3);
	assert.equal(cancellations, 2);
});

test("does not retry mutation requests or non-404 upstream responses", async (context) => {
	for (const { method, status } of [
		{ method: "POST", status: 404 },
		{ method: "GET", status: 503 },
	]) {
		await context.test(`${method} ${status}`, async () => {
			let calls = 0;
			const fetchImpl = async () => {
				calls += 1;
				return new Response("unchanged", { status });
			};

			const response = await fetchPreviewUpstream(
				fetchImpl,
				"https://job-123--49983.hf.jobs/v1/proxy/3000/submit",
				{ method, body: method === "POST" ? "once" : undefined },
				[0]
			);

			assert.equal(response.status, status);
			assert.equal(await response.text(), "unchanged");
			assert.equal(calls, 1);
		});
	}
});

test("aborts during the retry delay without issuing another upstream request", async () => {
	const controller = new AbortController();
	const reason = new Error("preview request cancelled");
	let calls = 0;
	let cancellations = 0;
	const fetchImpl = async () => {
		calls += 1;
		return discardedResponse(() => {
			cancellations += 1;
			controller.abort(reason);
		});
	};

	await assert.rejects(
		fetchPreviewUpstream(
			fetchImpl,
			"https://job-123--49983.hf.jobs/v1/proxy/3000/",
			{ method: "GET", signal: controller.signal },
			[1_000]
		),
		reason
	);
	assert.equal(calls, 1);
	assert.equal(cancellations, 1);
});

describe("isolated Agent preview runtime", () => {
	let upstream;
	let upstreamWs;
	let runtime;
	let upstreamOrigin;
	let runtimeOrigin;
	let incarnation = "generation-1";
	let capabilityExpiresAt = new Date(Date.now() + 60_000).toISOString();

	before(async () => {
		upstream = http.createServer(async (request, response) => {
			const url = new URL(request.url ?? "/", "http://upstream.invalid");
			if (url.pathname === "/slow") {
				response.writeHead(200, { "content-type": "text/event-stream" });
				response.write("data: ready\n\n");
				const keepalive = setInterval(() => response.write(": keepalive\n\n"), 1_000);
				request.once("close", () => clearInterval(keepalive));
				return;
			}
			if (url.pathname === "/") {
				response.writeHead(200, {
					"content-type": "text/html; charset=utf-8",
					"content-security-policy": "default-src 'none'",
				});
				response.end(
					`<!doctype html><html><head><title>App</title></head><body><script>fetch('/api/data')</script><form action="/submit" method="post"></form></body></html>`
				);
				return;
			}
			if (url.pathname === "/headers") {
				response.setHeader("content-type", "application/json");
				response.end(JSON.stringify(request.headers));
				return;
			}
			if (url.pathname === "/cookie") {
				response.setHeader("set-cookie", [
					"preview-session=one; Path=/; HttpOnly",
					"preference=two; Domain=upstream.invalid; Path=/account",
				]);
				response.end("cookie");
				return;
			}
			if (url.pathname === "/local-redirect") {
				response.writeHead(302, { location: "/next?ok=1" });
				response.end();
				return;
			}
			let body = "";
			for await (const chunk of request) body += chunk;
			response.setHeader("content-type", "application/json");
			response.end(JSON.stringify({ path: url.pathname, method: request.method, body }));
		});
		upstreamWs = new WebSocketServer({ noServer: true });
		upstream.on("upgrade", (request, socket, head) => {
			if (request.url !== "/socket") return socket.destroy();
			upstreamWs.handleUpgrade(request, socket, head, (ws) => {
				ws.on("message", (message) => ws.send(`echo:${message.toString()}`));
			});
		});
		upstream.listen(0, "127.0.0.1");
		await once(upstream, "listening");
		const upstreamAddress = upstream.address();
		assert(upstreamAddress && typeof upstreamAddress !== "string");
		upstreamOrigin = `http://127.0.0.1:${upstreamAddress.port}`;

		const config = {
			chatOrigin: "https://chat.example",
			publicOrigin: "http://127.0.0.1",
			resolverSecret: "s".repeat(32),
			jobsToken: JOBS_TOKEN,
			resolverUrl: "https://chat.example/api/agent/preview/resolve",
			port: 0,
		};
		const resolution = () => ({
			sandboxId: SANDBOX_ID,
			sandboxGeneration: incarnation,
			processId: "p-42",
			baseUrl: `https://${SANDBOX_ID}--49983.hf.jobs`,
			port: 3000,
			sandboxToken: SANDBOX_TOKEN,
			expiresAt: capabilityExpiresAt,
		});
		const resolveCapability = async (capability) => {
			if (capability !== CAPABILITY || incarnation !== "generation-1") {
				throw new Error("stale capability");
			}
			return resolution();
		};
		const fetchImpl = (target, init) => {
			const url = new URL(target);
			const marker = "/v1/proxy/3000";
			assert.ok(url.pathname.startsWith(marker));
			const edgeHeaders = new Headers(init.headers);
			assert.equal(edgeHeaders.get("authorization"), `Bearer ${JOBS_TOKEN}`);
			assert.equal(edgeHeaders.get("x-sandbox-token"), SANDBOX_TOKEN);
			// The real HF edge and sbx-server consume both credentials before the
			// request reaches the untrusted inner preview process.
			edgeHeaders.delete("authorization");
			edgeHeaders.delete("x-sandbox-token");
			return fetch(`${upstreamOrigin}${url.pathname.slice(marker.length) || "/"}${url.search}`, {
				...init,
				headers: edgeHeaders,
			});
		};
		class MappedWebSocket extends WebSocket {
			constructor(target, protocols, options) {
				const url = new URL(target);
				const marker = "/v1/proxy/3000";
				const mapped = `${upstreamOrigin.replace("http:", "ws:")}${url.pathname.slice(marker.length) || "/"}${url.search}`;
				const headers = { ...options.headers };
				assert.equal(headers.authorization, `Bearer ${JOBS_TOKEN}`);
				assert.equal(headers["x-sandbox-token"], SANDBOX_TOKEN);
				delete headers.authorization;
				delete headers["x-sandbox-token"];
				super(mapped, protocols, { ...options, headers });
			}
		}
		runtime = createPreviewRuntime({
			config,
			resolveCapability,
			fetchImpl,
			WebSocketImpl: MappedWebSocket,
		});
		runtime.listen(0, "127.0.0.1");
		await once(runtime, "listening");
		const runtimeAddress = runtime.address();
		assert(runtimeAddress && typeof runtimeAddress !== "string");
		runtimeOrigin = `http://127.0.0.1:${runtimeAddress.port}`;
		config.publicOrigin = runtimeOrigin;
	});

	after(async () => {
		for (const client of upstreamWs.clients) client.terminate();
		upstreamWs.close();
		await Promise.all([
			new Promise((resolve) => runtime.close(resolve)),
			new Promise((resolve) => upstream.close(resolve)),
		]);
	});

	test("injects the scoping shim inside an opaque-origin sandbox", async () => {
		const response = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/`);
		assert.equal(response.status, 200);
		const html = await response.text();
		assert.match(html, /window\.fetch/);
		assert.match(html, /XMLHttpRequest/);
		assert.match(html, /WebSocket/);
		assert.match(html, /EventSource/);
		assert.match(html, /form instanceof HTMLFormElement/);
		assert.match(html, /pushState/);
		assert.ok(html.indexOf("window.fetch") < html.indexOf("fetch('/api/data')"));
		const csp = response.headers.get("content-security-policy") ?? "";
		assert.match(csp, /sandbox[^;]*allow-scripts/);
		assert.doesNotMatch(csp, /allow-same-origin/);
		assert.match(csp, /worker-src 'none'/);
		assert.match(csp, /frame-ancestors https:\/\/chat\.example/);
		assert.equal(response.headers.get("referrer-policy"), "same-origin");
		assert.equal(response.headers.get("access-control-allow-origin"), "null");
		assert.equal(response.headers.get("access-control-allow-credentials"), "true");
		assert.equal(response.headers.get("cross-origin-resource-policy"), "cross-origin");
		assert.doesNotMatch(html, new RegExp(JOBS_TOKEN));
	});

	test("answers opaque-origin CORS preflights without exposing edge credentials", async () => {
		const response = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/api/data`, {
			method: "OPTIONS",
			headers: {
				origin: "null",
				"access-control-request-method": "POST",
				"access-control-request-headers": "content-type, x-preview-value",
			},
		});
		assert.equal(response.status, 204);
		assert.equal(response.headers.get("access-control-allow-origin"), "null");
		assert.equal(response.headers.get("access-control-allow-methods"), "POST");
		assert.equal(
			response.headers.get("access-control-allow-headers"),
			"content-type, x-preview-value"
		);
	});

	test("recovers root-relative assets and native form POSTs with a same-origin 307", async () => {
		const referer = `${runtimeOrigin}/p/${CAPABILITY}/nested/page`;
		const asset = await fetch(`${runtimeOrigin}/assets/app.js?version=1`, {
			headers: { referer },
			redirect: "manual",
		});
		assert.equal(asset.status, 307);
		assert.equal(asset.headers.get("location"), `/p/${CAPABILITY}/assets/app.js?version=1`);

		const form = await fetch(`${runtimeOrigin}/submit`, {
			method: "POST",
			headers: { referer, "content-type": "application/x-www-form-urlencoded" },
			body: "answer=42",
			redirect: "manual",
		});
		assert.equal(form.status, 307);
		const submitted = await fetch(`${runtimeOrigin}${form.headers.get("location")}`, {
			method: "POST",
			headers: { "content-type": "application/x-www-form-urlencoded" },
			body: "answer=42",
		});
		assert.deepEqual(await submitted.json(), {
			path: "/submit",
			method: "POST",
			body: "answer=42",
		});
	});

	test("strips browser and HF edge credentials before the inner app", async () => {
		const response = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/headers`, {
			headers: {
				authorization: "Bearer chat-or-hf-secret",
				cookie: "chat-session=secret",
				forwarded: "for=attacker",
				"x-forwarded-host": "attacker.example",
				"x-sandbox-token": "browser-controlled",
				"x-hf-token": "hf_browser_secret",
				"x-safe-preview-header": "kept",
			},
		});
		const headers = await response.json();
		assert.equal(headers.authorization, undefined);
		assert.equal(headers.cookie, undefined);
		assert.equal(headers.forwarded, undefined);
		assert.equal(headers["x-forwarded-host"], undefined);
		assert.equal(headers["x-hf-token"], undefined);
		assert.equal(headers["x-sandbox-token"], undefined);
		assert.equal(headers["x-safe-preview-header"], "kept");
	});

	test("scopes upstream cookies and local redirects to one capability", async () => {
		const cookie = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/cookie`);
		const cookies = cookie.headers.getSetCookie();
		assert.equal(cookies.length, 2);
		for (const value of cookies) {
			assert.match(value, new RegExp(`Path=/p/${CAPABILITY}/`));
			assert.doesNotMatch(value, /Domain=/i);
			assert.match(value, /Secure/i);
		}
		const redirect = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/local-redirect`, {
			redirect: "manual",
		});
		assert.equal(redirect.status, 302);
		assert.equal(redirect.headers.get("location"), `/p/${CAPABILITY}/next?ok=1`);
	});

	test("blocks service-worker script requests before capability resolution", async () => {
		const response = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/sw.js`, {
			headers: { "service-worker": "script" },
		});
		assert.equal(response.status, 403);
		assert.match(await response.text(), /Service workers are disabled/);
	});

	test("relays WebSocket frames through the exact scoped preview target", async () => {
		const ws = new WebSocket(`${runtimeOrigin.replace("http:", "ws:")}/p/${CAPABILITY}/socket`);
		await once(ws, "open");
		ws.send("hello");
		const [message] = await once(ws, "message");
		assert.equal(message.toString(), "echo:hello");
		ws.close();
		await once(ws, "close");
	});

	test("cuts off a long HTTP stream when its capability expires", async () => {
		capabilityExpiresAt = new Date(Date.now() + 120).toISOString();
		try {
			const response = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/slow`);
			assert.equal(response.status, 200);
			const reader = response.body.getReader();
			const first = await reader.read();
			assert.equal(first.done, false);
			await new Promise((resolve) => setTimeout(resolve, 180));
			let closed = false;
			try {
				closed = (await reader.read()).done;
			} catch {
				closed = true;
			}
			assert.equal(closed, true);
		} finally {
			capabilityExpiresAt = new Date(Date.now() + 60_000).toISOString();
		}
	});

	test("rejects cross-capability, expired, and pre-restart capabilities", async () => {
		const crossCapability = await fetch(`${runtimeOrigin}/p/${OTHER_CAPABILITY}/`);
		assert.notEqual(crossCapability.status, 200);

		assert.throws(() =>
			validatePreviewResolution({
				sandboxId: SANDBOX_ID,
				sandboxGeneration: "generation-1",
				processId: "p-42",
				baseUrl: `https://${SANDBOX_ID}--49983.hf.jobs`,
				port: 3000,
				sandboxToken: SANDBOX_TOKEN,
				expiresAt: new Date(Date.now() - 1).toISOString(),
			})
		);

		incarnation = "generation-2";
		const stale = await fetch(`${runtimeOrigin}/p/${CAPABILITY}/`);
		assert.notEqual(stale.status, 200);
		incarnation = "generation-1";
	});
});

test("validates resolver targets against the exact hf.jobs sandbox origin", () => {
	const valid = {
		sandboxId: SANDBOX_ID,
		sandboxGeneration: "generation-1",
		processId: "p-42",
		baseUrl: `https://${SANDBOX_ID}--49983.hf.jobs`,
		port: 3000,
		sandboxToken: SANDBOX_TOKEN,
		expiresAt: new Date(Date.now() + 60_000).toISOString(),
	};
	assert.deepEqual(validatePreviewResolution(valid), valid);
	assert.throws(() =>
		validatePreviewResolution({ ...valid, baseUrl: "https://metadata.internal.example" })
	);
	assert.throws(() => validatePreviewResolution({ ...valid, processId: 42 }));
	assert.throws(() => validatePreviewResolution({ ...valid, processId: "" }));
	assert.throws(() => validatePreviewResolution({ ...valid, processId: "p\n7" }));
	assert.throws(() => validatePreviewResolution({ ...valid, processId: "p".repeat(257) }));
	assert.throws(() => validatePreviewResolution({ ...valid, port: 49_983 }));
	assert.throws(() => validatePreviewResolution({ ...valid, sandboxToken: "hf_raw_token" }));
});

test("requires distinct HTTPS origins and a high-entropy resolver secret", () => {
	const valid = runtimeConfigFromEnv({
		CHAT_UI_ORIGIN: "https://chat.example",
		AGENT_PREVIEW_ORIGIN: "https://preview.example",
		AGENT_PREVIEW_RESOLVER_SECRET: "s".repeat(32),
		AGENT_HF_TOKEN: JOBS_TOKEN,
		PORT: "7860",
	});
	assert.equal(valid.resolverUrl, "https://chat.example/api/agent/preview/resolve");

	const basePath = runtimeConfigFromEnv({
		CHAT_UI_ORIGIN: "https://chat.example/chat/",
		AGENT_PREVIEW_ORIGIN: "https://preview.example",
		AGENT_PREVIEW_RESOLVER_SECRET: "r".repeat(32),
		AGENT_HF_TOKEN: "hf_server_token",
		PORT: "7860",
	});
	assert.equal(basePath.chatOrigin, "https://chat.example");
	assert.equal(basePath.resolverUrl, "https://chat.example/chat/api/agent/preview/resolve");
	assert.equal(valid.jobsToken, JOBS_TOKEN);
	assert.throws(() =>
		runtimeConfigFromEnv({
			CHAT_UI_ORIGIN: "https://chat.example",
			AGENT_PREVIEW_ORIGIN: "https://chat.example",
			AGENT_PREVIEW_RESOLVER_SECRET: "s".repeat(32),
			AGENT_HF_TOKEN: JOBS_TOKEN,
		})
	);
	assert.throws(() =>
		runtimeConfigFromEnv({
			CHAT_UI_ORIGIN: "https://chat.example",
			AGENT_PREVIEW_ORIGIN: "http://preview.example",
			AGENT_PREVIEW_RESOLVER_SECRET: "short",
			AGENT_HF_TOKEN: JOBS_TOKEN,
		})
	);
	assert.throws(() =>
		runtimeConfigFromEnv({
			CHAT_UI_ORIGIN: "https://chat.example",
			AGENT_PREVIEW_ORIGIN: "https://preview.example",
			AGENT_PREVIEW_RESOLVER_SECRET: "s".repeat(32),
		})
	);
});