import assert from "node:assert/strict";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { afterEach, describe, it } from "node:test";
import { toToolDeclaration } from "@earendil-works/pi-ai";
import registerSubagentExtension from "../../src/extension/index.ts";
import { registerNativeSupervisorClient, TOP_LEVEL_SUPERVISOR_REFUSAL } from "../../src/intercom/native-supervisor-channel.ts";

type Handler = (event: any, context: any) => any;
type Tool = { name: string; exposure?: string; description?: string; promptSnippet?: string; parameters?: unknown; execute?: (...args: any[]) => any };

const runtimes: Array<{ handlers: Map<string, Handler[]>; context: any }> = [];

type RuntimeOptions = { config?: Record<string, unknown>; model?: unknown };

function createRuntime(messages: any[] = [], excluded: string[] = [], missingApis: string[] = [], options: RuntimeOptions = {}) {
	const handlers = new Map<string, Handler[]>();
	const tools = new Map<string, Tool>();
	let activeNames = ["read"];
	const excludedNames = new Set(excluded);
	const missingApiNames = new Set(missingApis);
	const pi = new Proxy({
		events: { on() { return () => {}; }, emit() {} },
		on(name: string, handler: Handler) {
			const registered = handlers.get(name) ?? [];
			registered.push(handler);
			handlers.set(name, registered);
		},
		registerTool(tool: Tool) {
			tools.set(tool.name, tool);
			if (!excludedNames.has(tool.name) && (!tool.exposure || tool.exposure === "direct" || tool.exposure === "model-only")) activeNames = [...new Set([...activeNames, tool.name])];
		},
		getAllTools() {
			return [...tools.values()].filter((tool) => !excludedNames.has(tool.name));
		},
		getActiveTools() { return [...activeNames]; },
		setActiveTools(names: string[]) {
			const available = new Set([...tools.keys(), "read"].filter((name) => !excludedNames.has(name)));
			activeNames = [...new Set(names.filter((name) => available.has(name)))];
		},
		registerCommand() {}, registerShortcut() {}, registerMessageRenderer() {}, sendMessage() {}, getSessionName() {},
	}, { get(target, property) {
		if (missingApiNames.has(String(property))) return undefined;
		return property in target ? target[property as keyof typeof target] : () => undefined;
	} });
	const context = {
		cwd: process.cwd(), isIdle() { return false; }, hasUI: false, model: options.model as any,
		ui: { setWidget() {}, theme: { fg(_name: string, text: string) { return text; }, bg(_name: string, text: string) { return text; }, bold(text: string) { return text; } } },
		sessionManager: {
			getSessionId() { return "activation-session"; }, getSessionFile() { return null; }, getEntries() { return []; },
			buildSessionContext() { return { messages }; },
		},
		modelRegistry: { getAvailable() { return []; } },
	};
	const childEnv = process.env.SELESAI_SUBAGENT_CHILD;
	const priorAgentDir = process.env.SELESAI_CODING_AGENT_DIR;
	const agentDir = fs.mkdtempSync(path.join(os.tmpdir(), "tool-activation-agent-"));
	if (options.config) {
		fs.mkdirSync(path.join(agentDir, "extensions", "subagent"), { recursive: true });
		fs.writeFileSync(path.join(agentDir, "extensions", "subagent", "config.json"), JSON.stringify(options.config));
	}
	process.env.SELESAI_CODING_AGENT_DIR = agentDir;
	delete process.env.SELESAI_SUBAGENT_CHILD;
	try {
		registerSubagentExtension(pi as any);
	} finally {
		if (childEnv === undefined) delete process.env.SELESAI_SUBAGENT_CHILD;
		else process.env.SELESAI_SUBAGENT_CHILD = childEnv;
		if (priorAgentDir === undefined) delete process.env.SELESAI_CODING_AGENT_DIR;
		else process.env.SELESAI_CODING_AGENT_DIR = priorAgentDir;
	}
	runtimes.push({ handlers, context });
	return {
		handlers, tools, context,
		active: () => [...activeNames],
		select: (names: string[]) => { activeNames = [...names]; },
		async emit(name: string, event: any) {
			for (const handler of handlers.get(name) ?? []) await handler(event, context);
		},
	};
}

afterEach(async () => {
	for (const runtime of runtimes.splice(0)) {
		for (const handler of runtime.handlers.get("session_shutdown") ?? []) await handler({ type: "session_shutdown", reason: "quit" }, runtime.context);
	}
});

describe("codemode-first subagent tools", () => {
	it("registers parent tools as codemode without a loader or automatic activation", async () => {
		const runtime = createRuntime();
		await runtime.emit("session_start", { type: "session_start", reason: "startup" });
		assert.equal(runtime.tools.has("subagents_enable"), false);
		for (const name of ["subagent", "bg_wait", "subagent_wait", "subagent_supervisor"]) {
			assert.equal(runtime.tools.get(name)?.exposure, "codemode", name);
		}
		// `contact_supervisor` is model-only: it is declared so a forked child's tool list equals
		// its parent's, and every call in a top-level session returns an error.
		assert.equal(runtime.tools.get("contact_supervisor")?.exposure, "model-only");
		assert.deepEqual(runtime.active(), ["read", "contact_supervisor"]);
	});

	it("does not promote tools on prompts, resume, reload, tree navigation, or model switches", async () => {
		const runtime = createRuntime();
		for (const event of [
			{ type: "session_start", reason: "resume" },
			{ type: "session_start", reason: "reload" },
			{ type: "session_tree", newLeafId: null, oldLeafId: null },
		]) {
			await runtime.emit(event.type, event);
			runtime.context.model = { api: "anthropic-messages", compat: { supportsMidConvoToolChanges: true, supportsMidConvoSystemMessages: true } };
			const prompt = { systemPrompt: "base", systemPromptOptions: { selectedTools: runtime.active(), sections: {} } };
			await runtime.emit("before_agent_start", prompt);
			assert.deepEqual(runtime.active(), ["read", "contact_supervisor"]);
			assert.deepEqual(prompt.systemPromptOptions.selectedTools, ["read", "contact_supervisor"]);
			assert.equal("advertised_subagents" in prompt.systemPromptOptions.sections, false);
		}
	});

	it("ignores obsolete toolActivation configuration without reviving the loader", async () => {
		for (const toolActivation of ["auto", "dynamic", "eager"]) {
			const runtime = createRuntime([], [], [], { config: { toolActivation } });
			await runtime.emit("session_start", { type: "session_start", reason: "startup" });
			assert.equal(runtime.tools.has("subagents_enable"), false);
			assert.deepEqual(runtime.active(), ["read", "contact_supervisor"]);
		}
	});

	// Pi records an extension unload (update, /reload) as a removal of that extension's tools in the
	// transcript, which a continued session replays. The fork has no activation loader and never replays
	// that selection, so its registrations cannot be turned off; codemode tools stay callable either way.
	it("keeps the codemode tools registered when a continued transcript records an extension unload", async () => {
		const declared = (added: Array<{ name: string }>, removed: Array<{ name: string }>) => ({
			role: "system", timestamp: 1, content: "base", sections: {},
			toolsAdded: added.map((tool) => ({ ...tool, description: "recorded", parameters: { type: "object", properties: {} } })),
			toolsRemoved: removed,
		});
		const unloaded = ["subagent", "contact_supervisor", "bg_wait", "subagent_wait", "subagent_supervisor"];
		const transcript = [
			// A pre-0.77 session had selected the loader; the fork must not revive it from the transcript.
			declared([{ name: "read" }, { name: "subagents_enable" }, ...unloaded.map((name) => ({ name }))], []),
			declared([], unloaded.map((name) => ({ name }))),
		];
		for (const reason of ["reload", "resume"]) {
			const runtime = createRuntime(transcript);
			await runtime.emit("session_start", { type: "session_start", reason });
			assert.equal(runtime.tools.has("subagents_enable"), false, `${reason}: no loader`);
			for (const name of ["subagent", "bg_wait", "subagent_wait", "subagent_supervisor"]) {
				assert.equal(runtime.tools.get(name)?.exposure, "codemode", `${reason}: ${name}`);
			}
			assert.equal(runtime.tools.get("contact_supervisor")?.exposure, "model-only", `${reason}: contact_supervisor`);
			// Codemode tools are callable without being selected, so the recorded removal cannot disable them.
			assert.deepEqual(runtime.active(), ["read", "contact_supervisor"], `${reason}: nothing to restore`);
		}
	});
});

describe("contact_supervisor in a top-level session", () => {
	it("declares the definition a child's working tool has, and a call is an error", async () => {
		const runtime = createRuntime();
		await runtime.emit("session_start", { type: "session_start", reason: "startup" });
		let child: Tool | undefined;
		registerNativeSupervisorClient({ registerTool(tool: Tool) { child = tool; }, getAllTools: () => [] } as never, {
			channelDir: fs.mkdtempSync(path.join(os.tmpdir(), "supervisor-channel-")), runId: "run", agent: "worker", childIndex: 0, orchestratorSessionId: "parent",
		} as never);
		const parent = runtime.tools.get("contact_supervisor");

		assert.ok(parent && child);
		assert.deepEqual(toToolDeclaration(parent as never), toToolDeclaration(child as never));
		assert.equal(parent.exposure, child.exposure);
		const result = await parent.execute?.("call-1", { reason: "need_decision", message: "?" }, new AbortController().signal, undefined, runtime.context);
		assert.equal(result?.isError, true);
		assert.equal(result?.content[0]?.text, TOP_LEVEL_SUPERVISOR_REFUSAL);
	});
});