import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { isCanonicalCodexSubscriptionModel, type CodexQuotaClient } from "./client.ts";
import { getCodexModels } from "../provider/models.ts";
import { isTerminalRateLimitError } from "../provider/errors.ts";

import { CODEX_RESERVE_MODEL } from "./reserve-policy.ts";

const RETURN_ENTRY = "codex-reserve-return";
type ThinkingLevel = ReturnType<ExtensionAPI["getThinkingLevel"]>;
interface ReserveReturn {
	accountKey: string;
	provider: string;
	modelId: string;
	thinkingLevel: ThinkingLevel;
}

function readReturn(ctx: ExtensionContext): ReserveReturn | undefined {
	const entry = [...ctx.sessionManager.getBranch()]
		.reverse()
		.find((item) => item.type === "custom" && item.customType === RETURN_ENTRY);
	if (entry?.type !== "custom") return undefined;
	// type-boundary: saved Reserve return record; validate identity and model fields below.
	type SavedReturn = unknown;
	const value: SavedReturn = entry.data;
	if (!value || typeof value !== "object") return undefined;
	if (
		!("accountKey" in value) ||
		typeof value.accountKey !== "string" ||
		!("provider" in value) ||
		typeof value.provider !== "string" ||
		!("modelId" in value) ||
		typeof value.modelId !== "string" ||
		!("thinkingLevel" in value)
	)
		return undefined;
	const thinkingLevel = value.thinkingLevel;
	if (
		thinkingLevel !== "off" &&
		thinkingLevel !== "minimal" &&
		thinkingLevel !== "low" &&
		thinkingLevel !== "medium" &&
		thinkingLevel !== "high" &&
		thinkingLevel !== "xhigh" &&
		thinkingLevel !== "max"
	)
		return undefined;
	return { accountKey: value.accountKey, provider: value.provider, modelId: value.modelId, thinkingLevel };
}

export function createCodexReserveController(pi: ExtensionAPI, client: CodexQuotaClient) {
	let switching = false;
	return {
		modelSelected(ctx: ExtensionContext) {
			if (!switching && ctx.model?.id !== CODEX_RESERVE_MODEL && readReturn(ctx)) pi.appendEntry(RETURN_ENTRY, null);
		},
		async beforeTurn(ctx: ExtensionContext): Promise<void> {
			if (ctx.model?.id !== CODEX_RESERVE_MODEL || !isCanonicalCodexSubscriptionModel(ctx.model)) return;
			const previous = readReturn(ctx);
			if (!previous) return;
			const model = ctx.model;
			const leaf = ctx.sessionManager.getLeafId();
			try {
				const status = await client.reserve(ctx);
				if (ctx.model !== model || ctx.sessionManager.getLeafId() !== leaf) return;
				if (!status || status.accountKey !== previous.accountKey) {
					ctx.ui.notify(
						"Cannot verify the original Luna Reserve account; automatic model return is paused. Use /model to choose another model.",
						"warning",
					);
					return;
				}
				if (!status.ordinaryUsageRecovered) return;
				const target = ctx.modelRegistry.find(previous.provider, previous.modelId);
				if (!target) {
					ctx.ui.notify(
						"Ordinary Codex usage is available again, but the original model is unavailable. Select a model with /model.",
						"warning",
					);
					return;
				}
				switching = true;
				if (!(await pi.setModel(target))) throw new Error("Original model authentication is unavailable");
				pi.setThinkingLevel(previous.thinkingLevel);
				pi.appendEntry(RETURN_ENTRY, null);
				ctx.ui.notify(`Ordinary Codex usage recovered. Restored ${target.name} and its reasoning level.`, "info");
			} catch (error) {
				ctx.ui.notify(
					`Could not complete ordinary Codex model recovery: ${error instanceof Error ? error.message : String(error)}. Use /model to check or switch manually.`,
					"warning",
				);
			} finally {
				switching = false;
			}
		},
		async settled(ctx: ExtensionContext): Promise<boolean> {
			const model = ctx.model;
			if (!model || !isCanonicalCodexSubscriptionModel(model)) return false;
			const branch = ctx.sessionManager.getBranch();
			const last = [...branch]
				.reverse()
				.find((entry) => entry.type === "message" && entry.message.role === "assistant");
			if (last?.type !== "message" || last.message.role !== "assistant" || last.message.stopReason !== "error")
				return false;
			if (last.message.model !== model.id || last.message.provider !== model.provider) return false;
			const errorText = last.message.errorMessage ?? "";
			if (
				!isTerminalRateLimitError(errorText) &&
				!/Codex usage limit reached|hit your ChatGPT usage limit/i.test(errorText)
			)
				return false;
			if (model.id === CODEX_RESERVE_MODEL) return true;
			const leaf = ctx.sessionManager.getLeafId();
			try {
				const status = await client.reserve(ctx);
				if (ctx.model !== model || ctx.sessionManager.getLeafId() !== leaf || !status?.entryAllowed) return true;
				const luna = getCodexModels().find((candidate) => candidate.id === "gpt-5.6-luna");
				const reserve = luna ? { ...luna, id: CODEX_RESERVE_MODEL, name: "Luna Reserve" } : undefined;
				if (!reserve || !isCanonicalCodexSubscriptionModel(reserve)) {
					ctx.ui.notify(
						"Codex quota exhausted; Luna Reserve is offered but unavailable in this provider’s model catalog.",
						"warning",
					);
					return true;
				}
				const previous: ReserveReturn = {
					accountKey: status.accountKey,
					provider: model.provider,
					modelId: model.id,
					thinkingLevel: pi.getThinkingLevel(),
				};
				pi.appendEntry(RETURN_ENTRY, previous);
				switching = true;
				if (!(await pi.setModel(reserve))) throw new Error("Luna Reserve authentication is unavailable");
				pi.setThinkingLevel(previous.thinkingLevel);
				ctx.ui.notify(
					"Codex quota exhausted. Switched to Luna Reserve (separate, limited allowance). Send ‘continue’ if you want to use Luna, or choose another model with /model. Nothing was retried.",
					"warning",
				);
			} catch (error) {
				ctx.ui.notify(
					`Codex quota exhausted; Luna Reserve fallback could not complete: ${error instanceof Error ? error.message : String(error)}. Check the current model with /model.`,
					"warning",
				);
			} finally {
				switching = false;
			}
			return true;
		},
	};
}