import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { PROVIDER_ID, USER_AGENT } from "./config.ts";
import { creditsUrl } from "./urls.ts";

const CREDITS_FETCH_TIMEOUT_MS = 30_000;
const CREDITS_RETRY_DELAY_MS_INITIAL = 5_000;
const CREDITS_RETRY_DELAY_MS_MAX = 5 * 60_000;
const CREDITS_RETRY_EXPONENT_MAX = 6;

export function parseCreditsRemaining(payload: unknown): number {
	const record = asRecord(payload);
	if (!record) {
		throw new Error("Experiential Labs credits response is invalid");
	}
	const data = asRecord(record.data) ?? record;
	const remaining = asNumber(data.remaining);
	if (remaining !== undefined) {
		return remaining;
	}
	const total = asNumber(data.total_credits);
	if (total !== undefined) {
		return total - (asNumber(data.total_usage) ?? 0);
	}
	const balance = asNumber(data.balance) ?? asNumber(data.credits);
	if (balance !== undefined) {
		return balance;
	}
	throw new Error("Experiential Labs credits response is missing a balance");
}

export function formatCreditsUsd(amount: number): string {
	return amount.toLocaleString("en-US", {
		style: "currency",
		currency: "USD",
		minimumFractionDigits: 2,
		maximumFractionDigits: 2,
	});
}

function asRecord(value: unknown): Record<string, unknown> | undefined {
	return value !== null && typeof value === "object" && !Array.isArray(value)
		? (value as Record<string, unknown>)
		: undefined;
}

function asNumber(value: unknown): number | undefined {
	return typeof value === "number" && Number.isFinite(value) ? value : undefined;
}

function isExplabsModel(model: ExtensionContext["model"]): model is NonNullable<ExtensionContext["model"]> {
	return model?.provider === PROVIDER_ID;
}

function isTransientCreditError(error: unknown): boolean {
	if (error instanceof DOMException && error.name === "TimeoutError") {
		return true;
	}
	if (error instanceof TypeError) {
		return true;
	}
	const status = (error as { status?: number } | undefined)?.status;
	return status === 408 || status === 429 || (typeof status === "number" && status >= 500 && status <= 599);
}

async function fetchRemaining(baseUrl: string, apiKey: string, signal: AbortSignal): Promise<number> {
	const timeout = AbortSignal.timeout(CREDITS_FETCH_TIMEOUT_MS);
	const response = await fetch(creditsUrl(baseUrl), {
		headers: {
			Accept: "application/json",
			Authorization: `Bearer ${apiKey}`,
			"User-Agent": USER_AGENT,
		},
		signal: AbortSignal.any([signal, timeout]),
	});
	if (!response.ok) {
		const error = new Error(`Experiential Labs credits failed: HTTP ${response.status}`) as Error & {
			status: number;
		};
		error.status = response.status;
		throw error;
	}
	return parseCreditsRemaining(await response.json());
}

export interface CreditStatusRuntime {
	refresh(ctx: ExtensionContext, selectedModel: ExtensionContext["model"]): Promise<void>;
	dispose(): void;
}

export function createCreditStatusRuntime(baseUrl: string): CreditStatusRuntime {
	let invocationSequence = 0;
	let committedInvocation = 0;
	let credentialEpoch = 0;
	let consecutiveFailures = 0;
	let retryAtMs = 0;
	let currentApiKey: string | undefined;
	let cachedRemaining: number | undefined;
	let inFlight:
		| { apiKey: string; credentialEpoch: number; controller: AbortController; operation: Promise<void> }
		| undefined;
	let disposed = false;
	type CredentialLease = { apiKey: string; credentialEpoch: number };

	function clearCreditState(): void {
		cachedRemaining = undefined;
		consecutiveFailures = 0;
		retryAtMs = 0;
	}

	function invalidateCredential(): void {
		credentialEpoch += 1;
		inFlight?.controller.abort();
	}

	function renderStatus(ctx: ExtensionContext): void {
		if (cachedRemaining === undefined) {
			ctx.ui.setStatus(PROVIDER_ID, undefined);
			return;
		}
		ctx.ui.setStatus(PROVIDER_ID, formatCreditsUsd(cachedRemaining));
	}

	function ownsCredential(lease: CredentialLease): boolean {
		return lease.apiKey === currentApiKey && lease.credentialEpoch === credentialEpoch;
	}

	async function fetchAndCache(lease: CredentialLease, signal: AbortSignal): Promise<void> {
		try {
			const remaining = await fetchRemaining(baseUrl, lease.apiKey, signal);
			if (disposed || !ownsCredential(lease)) {
				return;
			}
			cachedRemaining = remaining;
			consecutiveFailures = 0;
			retryAtMs = 0;
		} catch (error) {
			if (disposed || !ownsCredential(lease)) {
				throw error;
			}
			if (!isTransientCreditError(error)) {
				clearCreditState();
				throw error;
			}
			consecutiveFailures += 1;
			retryAtMs = Date.now() + Math.min(
				CREDITS_RETRY_DELAY_MS_INITIAL * 2 ** Math.min(consecutiveFailures - 1, CREDITS_RETRY_EXPONENT_MAX),
				CREDITS_RETRY_DELAY_MS_MAX,
			);
			throw error;
		}
	}

	async function shareFetch(apiKey: string, expectedCredentialEpoch: number): Promise<void> {
		const active = inFlight;
		if (active?.apiKey === apiKey && active.credentialEpoch === expectedCredentialEpoch) {
			await active.operation;
			return;
		}

		const controller = new AbortController();
		const operation = fetchAndCache({ apiKey, credentialEpoch: expectedCredentialEpoch }, controller.signal);
		const started = { apiKey, credentialEpoch: expectedCredentialEpoch, controller, operation };
		inFlight = started;
		try {
			await operation;
		} finally {
			if (inFlight === started) {
				inFlight = undefined;
			}
		}
	}

	return {
		async refresh(ctx, selectedModel) {
			if (disposed || !ctx.hasUI) {
				return;
			}
			const invocation = invocationSequence + 1;
			invocationSequence = invocation;
			if (!isExplabsModel(selectedModel)) {
				committedInvocation = invocation;
				invalidateCredential();
				ctx.ui.setStatus(PROVIDER_ID, undefined);
				return;
			}

			renderStatus(ctx);
			try {
				const auth = await ctx.modelRegistry.getApiKeyAndHeaders(selectedModel);
				if (disposed || invocation < committedInvocation) {
					return;
				}
				if (!auth.ok || !auth.apiKey) {
					committedInvocation = invocation;
					invalidateCredential();
					currentApiKey = undefined;
					clearCreditState();
					renderStatus(ctx);
					return;
				}
				committedInvocation = invocation;
				if (currentApiKey === auth.apiKey && Date.now() < retryAtMs) {
					return;
				}
				if (currentApiKey !== auth.apiKey) {
					currentApiKey = auth.apiKey;
					invalidateCredential();
					clearCreditState();
					renderStatus(ctx);
				}
				const lease = { apiKey: auth.apiKey, credentialEpoch };
				await shareFetch(auth.apiKey, lease.credentialEpoch);
				if (disposed || !ownsCredential(lease) || invocation !== invocationSequence) {
					return;
				}
				renderStatus(ctx);
			} catch {
				if (disposed || invocation !== invocationSequence) {
					return;
				}
				renderStatus(ctx);
			}
		},

		dispose() {
			if (disposed) {
				return;
			}
			disposed = true;
			committedInvocation = ++invocationSequence;
			invalidateCredential();
		},
	};
}