import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";

const WIDGET_KEY = "tokenin-usage";
const REFRESH_MS = 60_000;
const DEFAULT_BASE_URL = "https://lite.andlet.me/v1";

type Balance = {
	available_usd: number;
	fixed_available_usd?: number;
	payg_available_usd?: number;
	debt_usd?: number;
};

type LegacyUsage = {
	spend: number;
	max_budget?: number;
};

function proxyRoot(baseUrl: string): string {
	return baseUrl.replace(/\/+$/, "").replace(/\/v1$/, "");
}

function money(value: number): string {
	return `$${value.toFixed(2)}`;
}

async function requestJson(url: string, apiKey: string): Promise<{ status: number; body?: unknown }> {
	const response = await fetch(url, {
		headers: { Authorization: `Bearer ${apiKey}`, Accept: "application/json" },
		signal: AbortSignal.timeout(5_000),
	});
	let body: unknown;
	try {
		body = await response.json();
	} catch {
		// The status still lets the caller distinguish a missing balance route.
	}
	return { status: response.status, body };
}

function isRecord(value: unknown): value is Record<string, unknown> {
	return !!value && typeof value === "object" && !Array.isArray(value);
}

function readNumber(record: Record<string, unknown>, key: string): number | undefined {
	const value = record[key];
	return typeof value === "number" && Number.isFinite(value) ? value : undefined;
}

async function getUsageLines(ctx: ExtensionContext): Promise<string[]> {
	const resolved = await ctx.modelRegistry.getProviderAuth("tokenin");
	const apiKey = resolved?.auth.apiKey;
	if (!apiKey) return ["Token-In: not connected (run /tokenin add)"];

	const provider = ctx.modelRegistry.getProvider("tokenin");
	const baseUrl = resolved.auth.baseUrl ?? provider?.baseUrl ?? DEFAULT_BASE_URL;
	const root = proxyRoot(baseUrl);
	const balanceResponse = await requestJson(`${root}/tokenin/account/me/balance`, apiKey);

	if (balanceResponse.status >= 200 && balanceResponse.status < 300 && isRecord(balanceResponse.body)) {
		const available = readNumber(balanceResponse.body, "available_usd");
		if (available !== undefined) {
			const balance: Balance = {
				available_usd: available,
				fixed_available_usd: readNumber(balanceResponse.body, "fixed_available_usd"),
				payg_available_usd: readNumber(balanceResponse.body, "payg_available_usd"),
				debt_usd: readNumber(balanceResponse.body, "debt_usd"),
			};
			const debt = balance.debt_usd && balance.debt_usd > 0 ? ` · debt ${money(balance.debt_usd)}` : "";
			return [
				`Token-In available: ${money(balance.available_usd)}`,
				`Plan ${money(balance.fixed_available_usd ?? 0)} · PAYG ${money(balance.payg_available_usd ?? 0)}${debt}`,
			];
		}
	}

	// Match /tokenin usage: fall back only when the account balance is absent or unusable.
	if (balanceResponse.status !== 404 && (balanceResponse.status < 200 || balanceResponse.status >= 300)) {
		throw new Error(`balance endpoint returned HTTP ${balanceResponse.status}`);
	}

	const usageResponse = await requestJson(`${root}/key/info`, apiKey);
	if (usageResponse.status < 200 || usageResponse.status >= 300 || !isRecord(usageResponse.body)) {
		throw new Error(`usage endpoint returned HTTP ${usageResponse.status}`);
	}
	const info = isRecord(usageResponse.body.info) ? usageResponse.body.info : usageResponse.body;
	const spend = readNumber(info, "spend");
	if (spend === undefined) throw new Error("usage endpoint returned an unexpected payload");
	const usage: LegacyUsage = { spend, max_budget: readNumber(info, "max_budget") };
	const remaining = usage.max_budget === undefined ? "no budget limit" : `${money(Math.max(0, usage.max_budget - spend))} remaining`;
	return [`Token-In spent: ${money(spend)}`, remaining];
}

export default function tokenInUsageWidget(pi: ExtensionAPI) {
	let timer: ReturnType<typeof setInterval> | undefined;
	let refreshing = false;

	pi.on("session_start", async (_event, ctx) => {
		if (!ctx.hasUI) return;
		if (timer) clearInterval(timer);

		const refresh = async () => {
			if (refreshing) return;
			refreshing = true;
			try {
				ctx.ui.setWidget(WIDGET_KEY, await getUsageLines(ctx), { placement: "belowEditor" });
			} catch (error) {
				const message = error instanceof Error ? error.message : String(error);
				ctx.ui.setWidget(WIDGET_KEY, [`Token-In usage unavailable: ${message}`], { placement: "belowEditor" });
			} finally {
				refreshing = false;
			}
		};

		await refresh();
		timer = setInterval(() => void refresh(), REFRESH_MS);
	});

	pi.on("session_shutdown", (_event, ctx) => {
		if (timer) clearInterval(timer);
		timer = undefined;
		if (ctx.hasUI) ctx.ui.setWidget(WIDGET_KEY, undefined);
	});
}