import { BorderedLoader, getAgentDir, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
import type { OAuthCredentials } from "@earendil-works/pi-ai/compat";
import { mkdir, readFile, rename, rmdir, writeFile } from "node:fs/promises";
import { join } from "node:path";
import { extractWeeklyQuota, type WeeklyQuota } from "./usage.ts";

const CONFIG_FILE = join(getAgentDir(), "codex-accounts.json");
const AUTH_FILE = join(getAgentDir(), "auth.json");
const AUTH_LOCK = `${AUTH_FILE}.lock`;
const LEGACY_PROVIDER_PREFIX = "openai-codex-account-";
const NATIVE_PROVIDER = "openai-codex";
const STATUS_KEY = "codex-accounts";
const USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
const REFRESH_INTERVAL_MS = 5 * 60 * 1000;
const OPENAI_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann";
const OPENAI_TOKEN_URL = "https://auth.openai.com/oauth/token";
const OPENAI_DEVICE_USER_CODE_URL = "https://auth.openai.com/api/accounts/deviceauth/usercode";
const OPENAI_DEVICE_TOKEN_URL = "https://auth.openai.com/api/accounts/deviceauth/token";
const OPENAI_DEVICE_REDIRECT_URI = "https://auth.openai.com/deviceauth/callback";
const OPENAI_DEVICE_VERIFICATION_URI = "https://auth.openai.com/codex/device";

type StoredCredential = OAuthCredentials & { accountId?: string };
type AuthFile = Record<string, { type?: unknown; access?: unknown; refresh?: unknown; expires?: unknown; accountId?: unknown }>;

type AccountsConfig = {
	accounts: Map<string, StoredCredential>;
	activeAlias?: string;
};

type DeviceAuthInfo = {
	deviceAuthId: string;
	userCode: string;
	intervalSeconds: number;
};

function isValidAlias(alias: string): boolean {
	return /^[a-z0-9][a-z0-9-]{0,31}$/.test(alias);
}

function isStoredCredential(value: unknown): value is StoredCredential {
	if (typeof value !== "object" || value === null) return false;
	const credential = value as { access?: unknown; refresh?: unknown; expires?: unknown };
	return typeof credential.access === "string" && typeof credential.refresh === "string" && typeof credential.expires === "number";
}

async function readJson<T>(path: string, fallback: T): Promise<T> {
	try {
		return JSON.parse(await readFile(path, "utf8")) as T;
	} catch (error) {
		if (typeof error === "object" && error !== null && "code" in error && error.code === "ENOENT") return fallback;
		throw error;
	}
}

async function writeJson(path: string, value: unknown): Promise<void> {
	await mkdir(getAgentDir(), { recursive: true, mode: 0o700 });
	const temp = `${path}.${process.pid}.${Date.now()}.tmp`;
	await writeFile(temp, `${JSON.stringify(value, null, 2)}\n`, { encoding: "utf8", mode: 0o600 });
	await rename(temp, path);
}

async function readAccountsConfig(): Promise<AccountsConfig> {
	const raw = await readJson<{ accounts?: unknown; activeAlias?: unknown }>(CONFIG_FILE, {});
	const accounts = new Map<string, StoredCredential>();
	if (typeof raw.accounts === "object" && raw.accounts !== null && !Array.isArray(raw.accounts)) {
		for (const [alias, credential] of Object.entries(raw.accounts)) {
			if (isValidAlias(alias) && isStoredCredential(credential)) accounts.set(alias, credential);
		}
	}
	return { accounts, activeAlias: typeof raw.activeAlias === "string" && isValidAlias(raw.activeAlias) ? raw.activeAlias : undefined };
}

async function saveAccountsConfig(config: AccountsConfig): Promise<void> {
	await writeJson(CONFIG_FILE, {
		accounts: Object.fromEntries(config.accounts),
		...(config.activeAlias ? { activeAlias: config.activeAlias } : {}),
	});
}

async function withAuthLock<T>(operation: () => Promise<T>): Promise<T> {
	await mkdir(getAgentDir(), { recursive: true, mode: 0o700 });
	for (let attempt = 0; attempt < 100; attempt++) {
		try {
			await mkdir(AUTH_LOCK);
			try {
				return await operation();
			} finally {
				await rmdir(AUTH_LOCK).catch(() => {});
			}
		} catch (error) {
			if (typeof error !== "object" || error === null || !("code" in error) || error.code !== "EEXIST") throw error;
			await new Promise((resolve) => setTimeout(resolve, 50));
		}
	}
	throw new Error("等待 Pi auth.json 锁超时");
}

async function updateAuthFile(mutator: (auth: AuthFile) => void): Promise<void> {
	await withAuthLock(async () => {
		const auth = await readJson<AuthFile>(AUTH_FILE, {});
		mutator(auth);
		await writeJson(AUTH_FILE, auth);
	});
}

async function migrateLegacyAccounts(config: AccountsConfig): Promise<void> {
	let migrated = false;
	await updateAuthFile((auth) => {
		for (const [provider, credential] of Object.entries(auth)) {
			if (!provider.startsWith(LEGACY_PROVIDER_PREFIX) || !isStoredCredential(credential)) continue;
			const alias = provider.slice(LEGACY_PROVIDER_PREFIX.length);
			if (isValidAlias(alias) && !config.accounts.has(alias)) {
				config.accounts.set(alias, credential);
				migrated = true;
			}
			delete auth[provider];
		}
	});
	if (migrated) await saveAccountsConfig(config);
}

async function activateNativeCredential(credential: StoredCredential): Promise<void> {
	await updateAuthFile((auth) => {
		for (const provider of Object.keys(auth)) {
			if (provider.startsWith(LEGACY_PROVIDER_PREFIX)) delete auth[provider];
		}
		auth[NATIVE_PROVIDER] = { type: "oauth", ...credential };
	});
}

async function readNativeCredential(): Promise<StoredCredential | undefined> {
	const auth = await readJson<AuthFile>(AUTH_FILE, {});
	const entry = auth[NATIVE_PROVIDER];
	if (!entry || !isStoredCredential(entry)) return undefined;
	return {
		access: entry.access,
		refresh: entry.refresh,
		expires: entry.expires,
		...(typeof entry.accountId === "string" ? { accountId: entry.accountId } : {}),
	};
}

function getAccountId(accessToken: string): string | undefined {
	try {
		const payload = JSON.parse(Buffer.from(accessToken.split(".")[1] ?? "", "base64url").toString("utf8")) as {
			"https://api.openai.com/auth"?: { chatgpt_account_id?: unknown };
		};
		const accountId = payload["https://api.openai.com/auth"]?.chatgpt_account_id;
		return typeof accountId === "string" && accountId.length > 0 ? accountId : undefined;
	} catch {
		return undefined;
	}
}

async function readTokenResponse(response: Response, operation: "exchange" | "refresh"): Promise<StoredCredential> {
	if (!response.ok) throw new Error(`OpenAI Codex ${operation} token failed (${response.status})`);
	const token = (await response.json()) as { access_token?: unknown; refresh_token?: unknown; expires_in?: unknown };
	if (typeof token.access_token !== "string" || typeof token.refresh_token !== "string" || typeof token.expires_in !== "number") {
		throw new Error("OpenAI Codex token response missing required fields");
	}
	const accountId = getAccountId(token.access_token);
	return {
		access: token.access_token,
		refresh: token.refresh_token,
		expires: Date.now() + token.expires_in * 1000,
		...(accountId ? { accountId } : {}),
	};
}

function waitForDevicePoll(intervalSeconds: number): Promise<void> {
	return new Promise((resolve) => setTimeout(resolve, intervalSeconds * 1000));
}

async function refreshCredential(credential: StoredCredential): Promise<StoredCredential> {
	const response = await fetch(OPENAI_TOKEN_URL, {
		method: "POST",
		headers: { "Content-Type": "application/x-www-form-urlencoded" },
		body: new URLSearchParams({
			grant_type: "refresh_token",
			client_id: OPENAI_CLIENT_ID,
			refresh_token: credential.refresh,
		}),
	});
	return readTokenResponse(response, "refresh");
}

async function startDeviceAuthorization(): Promise<DeviceAuthInfo> {
	const response = await fetch(OPENAI_DEVICE_USER_CODE_URL, {
		method: "POST",
		headers: { "Content-Type": "application/json" },
		body: JSON.stringify({ client_id: OPENAI_CLIENT_ID }),
	});
	if (!response.ok) throw new Error(`OpenAI Codex device authorization failed (${response.status})`);
	const data = (await response.json()) as { device_auth_id?: unknown; user_code?: unknown; interval?: unknown };
	const intervalSeconds = typeof data.interval === "string" ? Number(data.interval) : data.interval;
	if (typeof data.device_auth_id !== "string" || typeof data.user_code !== "string" || typeof intervalSeconds !== "number") {
		throw new Error("OpenAI Codex device authorization response is invalid");
	}
	return { deviceAuthId: data.device_auth_id, userCode: data.user_code, intervalSeconds };
}

async function pollDeviceAuthorization(device: DeviceAuthInfo): Promise<{ code: string; verifier: string }> {
	const deadline = Date.now() + 15 * 60 * 1000;
	while (Date.now() < deadline) {
		await waitForDevicePoll(device.intervalSeconds);
		const response = await fetch(OPENAI_DEVICE_TOKEN_URL, {
			method: "POST",
			headers: { "Content-Type": "application/json" },
			body: JSON.stringify({ device_auth_id: device.deviceAuthId, user_code: device.userCode }),
		});
		if (response.ok) {
			const data = (await response.json()) as { authorization_code?: unknown; code_verifier?: unknown };
			if (typeof data.authorization_code !== "string" || typeof data.code_verifier !== "string") throw new Error("OpenAI Codex device token response is invalid");
			return { code: data.authorization_code, verifier: data.code_verifier };
		}
		if (response.status !== 403 && response.status !== 404) throw new Error(`OpenAI Codex device authorization failed (${response.status})`);
	}
	throw new Error("OpenAI Codex device authorization timed out");
}

async function loginWithDeviceCode(ctx: { hasUI: boolean; ui: { input(title: string, placeholder?: string): Promise<string | undefined> } }): Promise<StoredCredential> {
	if (!ctx.hasUI) throw new Error("/codex-login requires interactive mode");
	const device = await startDeviceAuthorization();
	const confirmed = await ctx.ui.input(
		`Open ${OPENAI_DEVICE_VERIFICATION_URI}, enter code ${device.userCode}, then press Enter:`,
		"Press Enter after authorizing",
	);
	if (confirmed === undefined) throw new Error("Login cancelled");
	const result = await pollDeviceAuthorization(device);
	const response = await fetch(OPENAI_TOKEN_URL, {
		method: "POST",
		headers: { "Content-Type": "application/x-www-form-urlencoded" },
		body: new URLSearchParams({
			grant_type: "authorization_code",
			client_id: OPENAI_CLIENT_ID,
			code: result.code,
			code_verifier: result.verifier,
			redirect_uri: OPENAI_DEVICE_REDIRECT_URI,
		}),
	});
	return readTokenResponse(response, "exchange");
}

async function fetchWeeklyQuota(credential: StoredCredential): Promise<WeeklyQuota> {
	const response = await fetch(USAGE_URL, {
		headers: {
			Authorization: `Bearer ${credential.access}`,
			...(credential.accountId ? { "chatgpt-account-id": credential.accountId } : {}),
		},
	});
	if (!response.ok) throw new Error(`Weekly quota request failed (${response.status})`);
	const quota = extractWeeklyQuota(await response.json());
	if (!quota) throw new Error("OpenAI did not return a weekly quota");
	return quota;
}

function isAuthFailure(error: unknown): boolean {
	return error instanceof Error && /\(401\)/.test(error.message);
}

function formatQuota(quota: WeeklyQuota | undefined): string {
	if (!quota) return "Weekly loading...";
	const timestamp = quota.resetsAt ? (quota.resetsAt < 1_000_000_000_000 ? quota.resetsAt * 1000 : quota.resetsAt) : undefined;
	if (!timestamp) return `Weekly ${quota.remainingPercent.toFixed(0)}% left`;
	const reset = new Date(timestamp);
	const pad = (value: number) => `${value}`.padStart(2, "0");
	return `Weekly ${quota.remainingPercent.toFixed(0)}% left · resets ${reset.getFullYear()}-${pad(reset.getMonth() + 1)}-${pad(reset.getDate())} ${pad(reset.getHours())}:${pad(reset.getMinutes())}`;
}

export default async function codexAccounts(pi: ExtensionAPI) {
	const config = await readAccountsConfig();
	await migrateLegacyAccounts(config);
	const quotas = new Map<string, WeeklyQuota>();
	const unavailableAliases = new Set<string>();
	const refreshing = new Map<string, Promise<void>>();
	let refreshTimer: ReturnType<typeof setInterval> | undefined;
	let disposed = false;

	function formatActiveQuota(alias: string): string {
		return unavailableAliases.has(alias) ? "Weekly unavailable" : formatQuota(quotas.get(alias));
	}

	function updateStatus(ctx: { model?: { provider: string }; ui: { setStatus(key: string, value?: string): void; theme: { fg(color: string, text: string): string } } }) {
		if (ctx.model?.provider !== NATIVE_PROVIDER || !config.activeAlias) {
			ctx.ui.setStatus(STATUS_KEY, undefined);
			return;
		}
		ctx.ui.setStatus(
			STATUS_KEY,
			ctx.ui.theme.fg("accent", `Codex: ${config.activeAlias}`) + ctx.ui.theme.fg("dim", ` · ${formatActiveQuota(config.activeAlias)}`),
		);
	}

	async function persistCredential(alias: string, credential: StoredCredential): Promise<void> {
		const existing = config.accounts.get(alias);
		if (!existing || existing.access !== credential.access || existing.refresh !== credential.refresh || existing.expires !== credential.expires) {
			config.accounts.set(alias, credential);
			await saveAccountsConfig(config);
		}
		if (alias !== config.activeAlias) return;
		await updateAuthFile((auth) => {
			const current = auth[NATIVE_PROVIDER];
			if (current && current.access === credential.access && current.refresh === credential.refresh) return;
			auth[NATIVE_PROVIDER] = { type: "oauth", ...credential };
		});
	}

	/**
	 * Stored alias credentials go stale because Pi refreshes (and rotates) the
	 * native auth.json credential independently. Prefer the freshest candidate
	 * between our copy and auth.json, and fall back to an OAuth refresh.
	 */
	async function ensureFreshCredential(alias: string, stored: StoredCredential): Promise<StoredCredential> {
		const candidates: StoredCredential[] = [stored];
		if (alias === config.activeAlias) {
			try {
				const native = await readNativeCredential();
				if (native && native.access !== stored.access) candidates.push(native);
			} catch {
				// auth.json unreadable; keep the stored credential only
			}
		}
		candidates.sort((a, b) => b.expires - a.expires);

		const now = Date.now();
		for (const candidate of candidates) {
			if (candidate.expires > now + 60_000) {
				await persistCredential(alias, candidate);
				return candidate;
			}
		}

		let lastError: unknown;
		for (const candidate of candidates) {
			try {
				const fresh = await refreshCredential(candidate);
				await persistCredential(alias, fresh);
				return fresh;
			} catch (error) {
				lastError = error;
			}
		}
		throw lastError instanceof Error ? lastError : new Error("No usable Codex credential");
	}

	async function refreshQuota(
		alias: string,
		storedCredential: StoredCredential,
		ctx: { model?: { provider: string }; ui: { setStatus(key: string, value?: string): void; theme: { fg(color: string, text: string): string } } },
	): Promise<void> {
		const existing = refreshing.get(alias);
		if (existing) return existing;

		const refresh = (async () => {
			try {
				let credential = await ensureFreshCredential(alias, storedCredential);
				try {
					quotas.set(alias, await fetchWeeklyQuota(credential));
				} catch (error) {
					if (!isAuthFailure(error)) throw error;
					credential = await refreshCredential(credential);
					await persistCredential(alias, credential);
					quotas.set(alias, await fetchWeeklyQuota(credential));
				}
				unavailableAliases.delete(alias);
			} catch {
				unavailableAliases.add(alias);
			} finally {
				refreshing.delete(alias);
				if (!disposed) updateStatus(ctx);
			}
		})();
		refreshing.set(alias, refresh);
		return refresh;
	}

	function refreshActiveQuota(ctx: { model?: { provider: string }; ui: { setStatus(key: string, value?: string): void; theme: { fg(color: string, text: string): string } } }) {
		if (disposed) return;
		const alias = config.activeAlias;
		const credential = alias ? config.accounts.get(alias) : undefined;
		if (!alias || !credential || ctx.model?.provider !== NATIVE_PROVIDER) return;
		void refreshQuota(alias, credential, ctx);
	}

	pi.registerCommand("codex-login", {
		description: "Log in a Codex alias and activate Pi's native openai-codex credential",
		handler: async (args, ctx) => {
			const alias = args.trim().toLowerCase();
			if (!isValidAlias(alias)) {
				ctx.ui.notify("Alias must contain 1-32 lowercase letters, numbers, or hyphens.", "error");
				return;
			}
			try {
				const credential = await loginWithDeviceCode(ctx);
				config.accounts.set(alias, credential);
				config.activeAlias = alias;
				await saveAccountsConfig(config);
				await activateNativeCredential(credential);
				await ctx.reload();
				return;
			} catch (error) {
				ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
			}
		},
	});

	pi.registerCommand("codex-use", {
		description: "Activate an alias by replacing Pi's native openai-codex credential",
		getArgumentCompletions: (prefix) => {
			const aliases = [...config.accounts.keys()].filter((alias) => alias.startsWith(prefix));
			return aliases.length > 0 ? aliases.map((alias) => ({ value: alias, label: alias })) : null;
		},
		handler: async (args, ctx) => {
			const alias = args.trim().toLowerCase();
			const credential = config.accounts.get(alias);
			if (!credential) {
				ctx.ui.notify(`Unknown alias “${alias}”. Run /codex-login ${alias} first.`, "error");
				return;
			}
			config.activeAlias = alias;
			await saveAccountsConfig(config);
			await activateNativeCredential(credential);
			await ctx.reload();
		},
	});

	pi.registerCommand("codex-accounts", {
		description: "Refresh and list saved Codex account aliases",
		handler: async (_args, ctx) => {
			const aliases = [...config.accounts.keys()].sort();
			if (aliases.length === 0) {
				ctx.ui.notify("No saved Codex accounts. Use /codex-login <alias>.", "info");
				return;
			}

			const refreshAll = Promise.all(
				aliases.map((alias) => refreshQuota(alias, config.accounts.get(alias)!, ctx)),
			).then(() => undefined);

			if (ctx.mode === "tui") {
				const error = await ctx.ui.custom<unknown>((tui, theme, _keybindings, done) => {
					const loader = new BorderedLoader(
						tui,
						theme,
						`Loading usage for ${aliases.length} Codex account${aliases.length === 1 ? "" : "s"}...`,
						{ cancellable: false },
					);
					void refreshAll.then(() => done(undefined), done);
					return loader;
				});
				if (error !== undefined) {
					ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
					return;
				}
			} else {
				try {
					await refreshAll;
				} catch (error) {
					ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
					return;
				}
			}

			ctx.ui.notify(
				aliases.map((alias) => `${alias}${alias === config.activeAlias ? " (active)" : ""}: ${formatActiveQuota(alias)}`).join("\n"),
				"info",
			);
		},
	});

	pi.on("session_start", (_event, ctx) => {
		disposed = false;
		updateStatus(ctx);
		refreshActiveQuota(ctx);
		if (ctx.hasUI && !refreshTimer) refreshTimer = setInterval(() => refreshActiveQuota(ctx), REFRESH_INTERVAL_MS);
	});

	pi.on("model_select", async (_event, ctx) => {
		updateStatus(ctx);
		refreshActiveQuota(ctx);
	});

	pi.on("agent_settled", async (_event, ctx) => {
		refreshActiveQuota(ctx);
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		disposed = true;
		if (refreshTimer) clearInterval(refreshTimer);
		refreshTimer = undefined;
		ctx.ui.setStatus(STATUS_KEY, undefined);
	});
}