import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { PROVIDER_ID } from "./config.ts";
import { latestPhysicalResponse } from "./session-response.ts";
import type { GatewayDetails, GatewayMessage } from "./gateway.ts";

function usd(value: number): string {
	return `$${value.toFixed(6)}`;
}

export function gatewayStatus(details: GatewayDetails): string {
	const changed = (details.ignoredParameters?.length ?? 0) + (details.omittedSampling?.length ?? 0);
	return [
		details.provider,
		details.chargedUsd !== undefined ? `charged ${usd(details.chargedUsd)}` : undefined,
		details.isByok ? "BYOK" : undefined,
		details.upstreamUsd !== undefined ? `upstream ${usd(details.upstreamUsd)}` : undefined,
		changed ? `⚠ ${changed} parameters` : undefined,
	].filter(Boolean).join(" · ");
}

export function formatGatewayDetails(message: GatewayMessage | undefined): string {
	if (!message?.explabs) return "No Experiential Labs request details on this session branch yet.";
	const details = message.explabs;
	return [
		`Experiential Labs: ${message.model}`,
		`Request: ${details.requestId ?? "not reported"}`,
		`Route: ${details.provider ?? "not reported"}`,
		`Route depth: ${details.routeDepth ?? "not reported"}`,
		`Route reason: ${details.routeReason ?? "not reported"}`,
		`Charged to wallet: ${details.chargedUsd !== undefined ? usd(details.chargedUsd) : "not reported"}`,
		`BYOK: ${details.isByok === undefined ? "not reported" : details.isByok ? "yes" : "no"}`,
		`Upstream cost: ${details.upstreamUsd !== undefined ? usd(details.upstreamUsd) : "not reported"}`,
		`Pi token-cost estimate: ${usd(message.usage.cost.total)}`,
		`Zero data retention: ${details.zdr ?? "not reported"}`,
		...(details.replayRepair ? [`Replay repair: ${details.replayRepair}`] : []),
		...(details.omittedSampling?.length ? [`Sampling parameters omitted by extension: ${details.omittedSampling.join(", ")}`] : []),
		...(details.ignoredParameters?.length ? ["Parameter changes reported by gateway:", ...details.ignoredParameters.map((item) => `- ${item}`)] : []),
	].join("\n");
}

export function registerGatewayDiagnostics(pi: ExtensionAPI): void {
	function refresh(ctx: ExtensionContext, model = ctx.model): void {
		if (!ctx.hasUI) return;
		const latest = latestPhysicalResponse(ctx) as GatewayMessage | undefined;
		const active = model?.provider === PROVIDER_ID || (model?.api === "pi-virtual" && latest?.provider === PROVIDER_ID);
		ctx.ui.setStatus("explabs-request", active && latest?.provider === PROVIDER_ID && latest.explabs
			? gatewayStatus(latest.explabs) || undefined : undefined);
	}
	pi.registerCommand("explabs", {
		description: "Show Experiential Labs route, actual charge, and parameter changes for the latest response",
		handler: async (_args, ctx) => {
			const latest = latestPhysicalResponse(ctx) as GatewayMessage | undefined;
			pi.sendMessage({
				customType: "explabs-diagnostics",
				content: formatGatewayDetails(latest?.provider === PROVIDER_ID ? latest : undefined),
				display: true,
			}, { triggerTurn: false });
		},
	});
	pi.on("session_start", (_event, ctx) => refresh(ctx));
	pi.on("session_tree", (_event, ctx) => refresh(ctx));
	pi.on("agent_settled", (_event, ctx) => refresh(ctx));
	pi.on("model_select", (event, ctx) => refresh(ctx, event.model));
	pi.on("session_shutdown", (_event, ctx) => {
		if (ctx.hasUI) ctx.ui.setStatus("explabs-request", undefined);
	});
}