// atlibs-code-mode: an MCP server behind one host NSID (com.fedfork.exec.mcp.message). Its one tool, run_assemblyscript, compiles an
// agent's AssemblyScript with asc (browser and Node alike) and runs it as an atlibs module: runModule on
// the host's ctx (no imports, 64 MiB memory, 5 s per event, the run's budget) and the host's own call,
// held to `allow`. node: only behind a dynamic import (no Worker global), so pages load it as is.
import { runModule } from "@publicdomainrelay/atlibs/run-module/index.js";

// one MCP JSON-RPC message in, its response out (null for a notification)
export const MESSAGE = "com.fedfork.exec.mcp.message";
export const PROVIDES = [MESSAGE];
// newest first: initialize answers the client's version when it is one of these, else the newest
export const PROTOCOL_VERSIONS = ["2025-11-25", "2025-06-18", "2025-03-26"];
export const VERSION = "0.2.0";
export const TOOL = "run_assemblyscript";
export const ASC_VERSION = "0.28.20";
// the compiled module's limits are the host's (atlibs WASM_LIMITS); these bound what reaches the compiler
export const SOURCE_MAX = 65536;
export const COMPILE_MS = 20000;
// a page maps these in its import map (node_modules/assemblyscript/dist/{asc,assemblyscript}.js, binaryen/index.js, long/index.js)
export const ASC_IMPORTS = ["assemblyscript/asc", "assemblyscript", "binaryen", "long"];
export const COMPILE_HEAP_MB = 512;

/**
 * AssemblyScript source -> { binary } or { error } (the compiler's log). The compiler runs in a worker
 * (compile-worker.js) that is terminated after `timeoutMs`: a crash or a runaway compile costs that worker.
 * It reads only the prelude, the entry and `code`, and refuses a module that declares imports.
 */
export async function compile(code, { timeoutMs = COMPILE_MS } = {}) {
	if (typeof code !== "string" || !code.trim()) return { error: "code must be a non-empty string" };
	if (code.length > SOURCE_MAX) return { error: `code exceeds ${SOURCE_MAX} characters` };
	const worker = await spawn(new URL("./compile-worker.js", import.meta.url));
	try {
		return await new Promise((resolve) => {
			const timer = setTimeout(() => resolve({ error: `compile ran longer than ${timeoutMs} ms` }), timeoutMs);
			const done = (result) => (clearTimeout(timer), resolve(result));
			worker.on(done, () => done({ error: "the compiler crashed (a stack overflow, or out of memory)" }));
			worker.post({ code });
		});
	} finally {
		worker.terminate();
	}
}

// ponytail: a fresh worker per compile (asc loads in ~0.3 s each time); keep a warm one if that shows.
const spawn = async (url) => {
	if (typeof Worker === "function") {
		// a same-origin blob that imports compile-worker.js: a page may load this file from a CDN, and a
		// Worker's script must be the page's own origin (a module worker may import from anywhere with CORS)
		const worker = new Worker(URL.createObjectURL(new Blob([`import ${JSON.stringify(url.href)};`], { type: "text/javascript" })), { type: "module" });
		// the page's import map, which its workers do not see: compile-worker.js loads asc from these
		const urls = Object.fromEntries(ASC_IMPORTS.map((name) => [name, import.meta.resolve(name)]));
		return {
			post: (m) => worker.postMessage({ ...m, urls }),
			on: (message, fail) => ((worker.onmessage = ({ data }) => message(data)), (worker.onerror = fail)),
			terminate: () => worker.terminate(),
		};
	}
	const { Worker: NodeWorker } = await import("node:worker_threads");
	const worker = new NodeWorker(url, { stdout: true, stderr: true, resourceLimits: { maxOldGenerationSizeMb: COMPILE_HEAP_MB } });
	return {
		post: (m) => worker.postMessage(m),
		on: (message, fail) => (worker.once("message", message), worker.once("error", fail), worker.once("exit", fail)),
		terminate: () => void worker.terminate(),
	};
};

const describe = (allow) =>
	`Compile and run AssemblyScript (asc ${ASC_VERSION}, stub runtime) in a sandbox: no imports (no console, Math.random or Date), ` +
	"64 MiB of memory, 5 s per event. Write `export function main(run: Run): void`; Json, Run and print are in scope. " +
	"print(text) writes a line of stdout. run.input is the tool's `input` as Json (get(key), has(key), at(i), " +
	".length (array items, object keys, string chars), .keys (string[]), .items (Json[]), .str, .num (f64), .bool, " +
	".isNull/.isBool/.isNumber/.isString/.isArray/.isObject; Json.object().set(k, Json) (values are Json: Json.string(s), not s), " +
	"push(Json), Json.array([...]), Json.string/number(f64)/bool/null, Json.parse(text), toString()). " +
	"run.finish(value) ends the run with that result. main runs once per event: run.start on the first; " +
	(allow.length
		? `run.call(id, nsid, input) or run.query(id, nsid, params, service?) makes an XRPC call (allowed: ${allow.join(", ")}); ` +
			"its result is the next event (run.id, run.ok, run.output, run.error). Keep what you need in run.state (Json). "
		: "it may make no calls. ") +
	"throw (or any trap) ends the run with an error. Returns JSON { stdout, result } (and error); a run that never calls run.finish returns result null.";

const text = (value, isError) => ({ content: [{ type: "text", text: JSON.stringify(value) }], ...(isError && { isError: true }) });

/**
 * The code-mode MCP server for a host: message(jsonrpc) -> the JSON-RPC response, or null for a
 * notification (or a response sent to us). Transport-agnostic: a hook, hostCall's calls, a Streamable
 * HTTP endpoint and the MCP SDK's transports all carry the same messages.
 *   ctx: the host's context (createPluginContext in a page, createNodeContext in Node): its wasm engine runs the code
 *   call: the host's call (hostCall, policyCall, ...) for the code's XRPC calls; absent: none
 *   allow: the NSIDs the code may call (default none), as a module's `calls`
 *   budget: runModule's budget for each run (atlibs defaults otherwise)
 * Returns message, `calls` for hostCall's calls (a module that lists MESSAGE in its calls), and `handler`
 * for runCompose's handlers (an edge routing MESSAGE to it).
 */
export function codeMode({ ctx, call, allow = [], budget } = {}) {
	if (!ctx?.wasm?.run) throw new Error("codeMode needs the host's ctx (its wasm engine runs the code)");
	allow = [...allow];
	const own = call ?? (async (nsid) => Promise.reject(new Error(`${nsid}: code mode was given no call`)));
	const tools = [{
		name: TOOL,
		title: "Run AssemblyScript",
		description: describe(allow),
		inputSchema: {
			type: "object",
			properties: { code: { type: "string", description: "AssemblyScript that exports main(run: Run): void" }, input: { description: "any JSON: run.input" } },
			required: ["code"],
		},
	}];
	const run = async (args, signal) => {
		const built = await compile(args?.code);
		if (built.error) return text({ error: `compile: ${built.error}` }, true);
		const mod = { name: "code-mode", module: built.binary, calls: allow };
		const ran = await runModule(ctx, mod, "run", { input: args.input ?? null }, own, { budget, signal });
		return text({ stdout: ran.result?.stdout ?? "", result: ran.result?.value ?? null, ...(ran.error !== undefined && { error: ran.error }) }, ran.error !== undefined);
	};
	const methods = {
		initialize: (params) => ({
			protocolVersion: PROTOCOL_VERSIONS.includes(params?.protocolVersion) ? params.protocolVersion : PROTOCOL_VERSIONS[0],
			capabilities: { tools: {} },
			serverInfo: { name: "atlibs-code-mode", version: VERSION },
		}),
		ping: () => ({}),
		"tools/list": () => ({ tools }),
		"tools/call": (params, signal) => {
			if (typeof params?.name !== "string") throw rpcError(-32602, "tools/call needs params.name");
			if (params.name !== TOOL) throw rpcError(-32602, `Unknown tool: ${params.name}`);
			return run(params.arguments ?? {}, signal);
		},
	};
	const message = async (msg, { signal } = {}) => {
		const id = msg?.id;
		if (id === undefined) return null; // a notification (notifications/initialized, cancelled): no response
		if (msg.jsonrpc !== "2.0" || typeof msg.method !== "string") return "result" in msg || "error" in msg ? null : { jsonrpc: "2.0", id, error: { code: -32600, message: "Invalid Request" } };
		try {
			if (!Object.hasOwn(methods, msg.method)) throw rpcError(-32601, `Method not found: ${msg.method}`);
			return { jsonrpc: "2.0", id, result: await methods[msg.method](msg.params, signal) };
		} catch (error) {
			return { jsonrpc: "2.0", id, error: { code: Number.isInteger(error?.code) ? error.code : -32603, message: String(error?.message ?? error) } };
		}
	};
	return {
		provides: PROVIDES,
		message,
		calls: { [MESSAGE]: (request = {}) => message(request.input, request) },
		handler: (input, { signal } = {}) => message(input, { signal }),
	};
}

const rpcError = (code, message) => Object.assign(new Error(message), { code });

/**
 * Streamable HTTP (MCP 2025-11-25), stateless and JSON only, over web Request/Response (pages' service
 * workers, Node, Deno, an ingress module's forward): POST one message, get its response as JSON, or 202
 * for a notification. No GET stream (405): this server never sends unprompted.
 */
export const streamableHttp = (message) => async (request) => {
	if (request.method !== "POST") return new Response(null, { status: 405, headers: { allow: "POST" } });
	let msg;
	try {
		msg = await request.json();
	} catch {
		return Response.json({ jsonrpc: "2.0", id: null, error: { code: -32700, message: "Parse error" } }, { status: 400 });
	}
	const reply = await message(msg, { signal: request.signal });
	return reply === null ? new Response(null, { status: 202 }) : Response.json(reply);
};