import assert from "node:assert/strict";
import { it } from "node:test";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { registerGatewayDiagnostics } from "../src/diagnostics.ts";
import type { GatewayMessage } from "../src/gateway.ts";

it("restores branch diagnostics, follows virtual routing, and shows details without an inference turn", async () => {
	const handlers = new Map<string, ((event: any, ctx: ExtensionContext) => void)[]>();
	let command: any;
	const sent: any[] = [];
	registerGatewayDiagnostics({
		on: (event: string, handler: any) => handlers.set(event, [...(handlers.get(event) ?? []), handler]),
		registerCommand: (_name: string, options: any) => { command = options; },
		sendMessage: (message: any, options: any) => sent.push({ message, options }),
	} as unknown as ExtensionAPI);
	const message: GatewayMessage = {
		role: "assistant", api: "openai-responses", provider: "explabs", model: "cheap", timestamp: 0, content: [], stopReason: "stop",
		usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
		explabs: { provider: "actual-rung", requestId: "req_1", chargedUsd: 0, isByok: true, upstreamUsd: 0.001 },
	};
	let branch: any[] = [{ type: "message", message: JSON.parse(JSON.stringify(message)) }];
	const statuses = new Map<string, string | undefined>();
	const ctx = {
		hasUI: true, model: { provider: "router", api: "pi-virtual" },
		sessionManager: { getBranch: () => branch },
		ui: { setStatus: (key: string, value: string | undefined) => statuses.set(key, value) },
	} as unknown as ExtensionContext;
	function fire(event: string, data: any = {}) { handlers.get(event)?.forEach(handler => handler(data, ctx)); }
	fire("session_start");
	assert.match(statuses.get("explabs-request") ?? "", /actual-rung.*charged \$0\.000000.*BYOK/);
	await command.handler("", ctx);
	assert.match(sent[0].message.content, /Request: req_1/);
	assert.deepEqual(sent[0].options, { triggerTurn: false });

	branch.push({ type: "message", message: { ...message, provider: "other", explabs: undefined } });
	fire("agent_settled");
	assert.equal(statuses.get("explabs-request"), undefined);
	branch = [{ type: "message", message }, { type: "message", message: { ...message, provider: "other", stopReason: "error" } }];
	fire("session_tree");
	assert.match(statuses.get("explabs-request") ?? "", /actual-rung/);
	fire("model_select", { model: { provider: "other", api: "openai-responses" } });
	assert.equal(statuses.get("explabs-request"), undefined);
	fire("session_shutdown");
	assert.equal(statuses.get("explabs-request"), undefined);
});