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

// ============================================================================
// Prime Agent variant of pi-openai-codex-status.
//
// Unlike the pi edition (which renders a floating `ctx.ui.custom` overlay), this
// build targets Prime Agent's daemon-backed extension-UI bridge, where
// arbitrary custom overlays are NOT supported. Instead it:
//   * posts /status as an in-chat report via sendMessage(), and
//   * shows live weekly Codex usage in a small widget below the editor,
//     refreshed on session start, agent end, and model selection.
// The widget is provider-scoped and only appears while an openai-codex model
// is active. Only daemon-bridged UI primitives are used.
// ============================================================================

type RawRateLimitWindowSnapshot = {
	used_percent?: number | string | null;
	limit_window_seconds?: number | null;
	reset_after_seconds?: number | null;
	reset_at?: number | null;
};

type RawRateLimitStatusDetails = {
	allowed?: boolean;
	limit_reached?: boolean;
	primary_window?: RawRateLimitWindowSnapshot | null;
	secondary_window?: RawRateLimitWindowSnapshot | null;
};

type RawCreditStatusDetails = {
	has_credits?: boolean;
	unlimited?: boolean;
	balance?: string | number | null;
};

type RawAdditionalRateLimitDetails = {
	limit_name?: string;
	metered_feature?: string;
	rate_limit?: RawRateLimitStatusDetails | null;
};

type RawRateLimitStatusPayload = {
	plan_type?: string;
	rate_limit?: RawRateLimitStatusDetails | null;
	credits?: RawCreditStatusDetails | null;
	additional_rate_limits?: RawAdditionalRateLimitDetails[] | null;
};

type LimitWindow = {
	usedPercent: number;
	windowSeconds?: number;
	resetsAt?: number;
};

type AdditionalLimit = {
	id: string;
	name?: string;
	primary?: LimitWindow;
	secondary?: LimitWindow;
};

type StatusSnapshot = {
	planType?: string;
	allowed?: boolean;
	limitReached?: boolean;
	primary?: LimitWindow;
	secondary?: LimitWindow;
	credits?: {
		hasCredits: boolean;
		unlimited: boolean;
		balance?: string;
	};
	additional: AdditionalLimit[];
	fetchedAt: number;
};

const PROVIDER_ID = "openai-codex";
const USAGE_URL = "https://chatgpt.com/backend-api/wham/usage";
const USAGE_PAGE_URL = "https://chatgpt.com/codex/settings/usage";
const CACHE_TTL_MS = 60_000;
const WEEKLY_STATUS_KEY = "codex-weekly";
const REPORT_TYPE = "openai-status";
const WEEKLY_WINDOW_SECONDS = 7 * 24 * 3600; // 7 days

let cache: StatusSnapshot | null = null;
let inFlight: Promise<StatusSnapshot> | null = null;

// ---------------------------------------------------------------------------
// Parsing / helpers
// ---------------------------------------------------------------------------

function clampPercent(value: unknown): number {
	const n =
		typeof value === "string"
			? Number(value)
			: typeof value === "number"
				? value
				: 0;
	if (!Number.isFinite(n)) return 0;
	return Math.max(0, Math.min(100, Math.round(n)));
}

function parseJwtPayload(token: string): Record<string, unknown> | null {
	const parts = token.split(".");
	const encodedPayload = parts[1];
	if (!encodedPayload) return null;
	try {
		const payload = encodedPayload.replace(/-/g, "+").replace(/_/g, "/");
		const padLength = (4 - (payload.length % 4)) % 4;
		const base64 = payload + "=".repeat(padLength);
		const json = Buffer.from(base64, "base64").toString("utf8");
		const parsed = JSON.parse(json);
		return parsed && typeof parsed === "object"
			? (parsed as Record<string, unknown>)
			: null;
	} catch {
		return null;
	}
}

function getAccountIdFromToken(token: string): string | undefined {
	const payload = parseJwtPayload(token);
	if (!payload) return undefined;

	const direct = payload["https://api.openai.com/auth.chatgpt_account_id"];
	if (typeof direct === "string" && direct.length > 0) return direct;

	const auth = payload["https://api.openai.com/auth"];
	if (auth && typeof auth === "object") {
		const nested = (auth as Record<string, unknown>).chatgpt_account_id;
		if (typeof nested === "string" && nested.length > 0) return nested;
	}

	return undefined;
}

function normalizeWindow(
	window: RawRateLimitWindowSnapshot | null | undefined,
): LimitWindow | undefined {
	if (!window) return undefined;
	const windowSeconds =
		typeof window.limit_window_seconds === "number"
			? window.limit_window_seconds
			: undefined;
	const resetAtSeconds =
		typeof window.reset_at === "number" ? window.reset_at : undefined;
	const resetAfterSeconds =
		typeof window.reset_after_seconds === "number"
			? window.reset_after_seconds
			: undefined;
	return {
		usedPercent: clampPercent(window.used_percent),
		windowSeconds,
		resetsAt:
			typeof resetAtSeconds === "number"
				? resetAtSeconds * 1000
				: typeof resetAfterSeconds === "number"
					? Date.now() + resetAfterSeconds * 1000
					: undefined,
	};
}

function normalizePayload(payload: RawRateLimitStatusPayload): StatusSnapshot {
	const additional = Array.isArray(payload.additional_rate_limits)
		? payload.additional_rate_limits.map((limit) => ({
				id: limit.metered_feature || limit.limit_name || "additional",
				name: limit.limit_name,
				primary: normalizeWindow(limit.rate_limit?.primary_window),
				secondary: normalizeWindow(limit.rate_limit?.secondary_window),
			}))
		: [];

	return {
		planType: payload.plan_type,
		allowed: payload.rate_limit?.allowed,
		limitReached: payload.rate_limit?.limit_reached,
		primary: normalizeWindow(payload.rate_limit?.primary_window),
		secondary: normalizeWindow(payload.rate_limit?.secondary_window),
		credits: payload.credits
			? {
					hasCredits: !!payload.credits.has_credits,
					unlimited: !!payload.credits.unlimited,
					balance:
						payload.credits.balance === null ||
						payload.credits.balance === undefined
							? undefined
							: String(payload.credits.balance),
				}
			: undefined,
		additional,
		fetchedAt: Date.now(),
	};
}

async function fetchUsageSnapshot(
	modelRegistry: ModelRegistry,
	signal?: AbortSignal,
): Promise<StatusSnapshot> {
	const accessToken = await modelRegistry.getApiKeyForProvider(PROVIDER_ID);
	if (!accessToken) {
		throw new Error(
			"You are not logged into ChatGPT Plus/Pro Codex in prime-agent. Run /login and choose OpenAI Codex.",
		);
	}

	const accountId = getAccountIdFromToken(accessToken);
	if (!accountId) {
		throw new Error("OpenAI Codex credential is missing ChatGPT account ID.");
	}

	const response = await fetch(USAGE_URL, {
		method: "GET",
		headers: {
			Authorization: `Bearer ${accessToken}`,
			"ChatGPT-Account-Id": accountId,
			"User-Agent": "prime-openai-status",
		},
		signal,
	});

	if (!response.ok) {
		let details = "";
		try {
			details = await response.text();
		} catch {
			// ignore body parse errors
		}
		const suffix = details
			? ` ${details.replace(/\s+/g, " ").trim().slice(0, 160)}`
			: "";
		throw new Error(
			`OpenAI usage request failed (${response.status} ${response.statusText}).${suffix}`,
		);
	}

	const payload = (await response.json()) as RawRateLimitStatusPayload;
	return normalizePayload(payload);
}

/** Cache + single-flight snapshot fetch so rapid loops never stack requests. */
function getSnapshot(
	modelRegistry: ModelRegistry,
	signal?: AbortSignal,
): Promise<StatusSnapshot> {
	if (cache && Date.now() - cache.fetchedAt < CACHE_TTL_MS) {
		return Promise.resolve(cache);
	}
	if (inFlight) return inFlight;
	inFlight = fetchUsageSnapshot(modelRegistry, signal)
		.then((snapshot) => {
			cache = snapshot;
			return snapshot;
		})
		.finally(() => {
			inFlight = null;
		});
	return inFlight;
}

// ---------------------------------------------------------------------------
// Formatting
// ---------------------------------------------------------------------------

function describeWindow(seconds?: number): string {
	if (!seconds || seconds <= 0) return "limit";
	const hours = seconds / 3600;
	const days = seconds / 86_400;
	if (hours <= 24) {
		const rounded = Math.max(1, Math.round(hours));
		return `${rounded}h limit`;
	}
	if (days <= 7.5) return "weekly limit";
	if (days <= 31) return "monthly limit";
	return "annual limit";
}

function formatAge(ms: number): string {
	if (ms < 5_000) return "just now";
	if (ms < 60_000) return `${Math.round(ms / 1000)}s ago`;
	if (ms < 3_600_000) return `${Math.round(ms / 60_000)}m ago`;
	return `${Math.round(ms / 3_600_000)}h ago`;
}

function formatDuration(ms: number): string {
	const totalSeconds = Math.max(0, Math.round(ms / 1000));
	if (totalSeconds < 60) return `${totalSeconds}s`;
	const hours = Math.floor(totalSeconds / 3600);
	const minutes = Math.floor((totalSeconds % 3600) / 60);
	const days = Math.floor(hours / 24);
	if (days > 0) {
		const remHours = hours % 24;
		return remHours > 0 ? `${days}d ${remHours}h` : `${days}d`;
	}
	if (hours > 0) return minutes > 0 ? `${hours}h ${minutes}m` : `${hours}h`;
	return `${minutes}m`;
}

function formatResetTime(timestamp?: number): string {
	if (!timestamp) return "reset unknown";
	const diff = timestamp - Date.now();
	if (diff <= 0) return "resetting now";
	if (diff < 36 * 3600 * 1000) return `resets in ${formatDuration(diff)}`;
	return `resets ${new Date(timestamp).toLocaleString(undefined, {
		weekday: "short",
		hour: "2-digit",
		minute: "2-digit",
		month: "short",
		day: "numeric",
	})}`;
}

// ---------------------------------------------------------------------------
// Weekly window selection + widget/status updater
// ---------------------------------------------------------------------------

function pickWeeklyWindow(snapshot: StatusSnapshot): LimitWindow | undefined {
	const candidates = [
		snapshot.primary,
		snapshot.secondary,
		...snapshot.additional.flatMap((limit) => [limit.primary, limit.secondary]),
	];
	// Prefer the window closest to exactly 7 days.
	let best: LimitWindow | undefined;
	let bestDelta = Infinity;
	for (const w of candidates) {
		if (!w?.windowSeconds || w.windowSeconds <= 0) continue;
		const delta = Math.abs(w.windowSeconds - WEEKLY_WINDOW_SECONDS);
		if (delta < bestDelta) {
			bestDelta = delta;
			best = w;
		}
	}
	return best;
}

function formatWeeklyStatus(snapshot: StatusSnapshot): string {
	const weekly = pickWeeklyWindow(snapshot);
	if (!weekly) return "Codex status unavailable";
	const left = Math.max(0, 100 - weekly.usedPercent);
	return `Codex: ${left}% left · ${formatResetTime(weekly.resetsAt)}`;
}

type ExtensionContextLike = {
	hasUI: boolean;
	modelRegistry: ModelRegistry;
	model?: { provider: string; id: string };
	ui: {
		setStatus(key: string, text: string | undefined): void;
		setWidget(
			key: string,
			content: string[] | undefined,
			options?: { placement?: "aboveEditor" | "belowEditor" },
		): void;
		notify(message: string, type?: "info" | "warning" | "error"): void;
	};
};

function clearWeeklyStatus(ctx: ExtensionContextLike): void {
	try {
		if (ctx.hasUI) {
			ctx.ui.setWidget(WEEKLY_STATUS_KEY, undefined);
			ctx.ui.setStatus(WEEKLY_STATUS_KEY, undefined);
		}
	} catch {
		// ignore
	}
}

function isCodexProviderActive(ctx: ExtensionContextLike): boolean {
	return ctx.model?.provider === PROVIDER_ID;
}

function pushWeeklyStatus(
	ctx: ExtensionContextLike,
	snapshot: StatusSnapshot,
): void {
	try {
		if (ctx.hasUI) {
			if (isCodexProviderActive(ctx)) {
				const status = formatWeeklyStatus(snapshot);
				ctx.ui.setWidget(WEEKLY_STATUS_KEY, [status], {
					placement: "belowEditor",
				});
				ctx.ui.setStatus(WEEKLY_STATUS_KEY, status);
			} else {
				clearWeeklyStatus(ctx);
			}
		}
	} catch {
		// ignore
	}
}

function refreshWeeklyStatus(ctx: ExtensionContextLike): void {
	if (!ctx.hasUI) return;
	// Provider-scoped: only fetch/show while an openai-codex model is active.
	if (!isCodexProviderActive(ctx)) {
		clearWeeklyStatus(ctx);
		return;
	}
	void getSnapshot(ctx.modelRegistry)
		.then((snapshot) => pushWeeklyStatus(ctx, snapshot))
		.catch(() => clearWeeklyStatus(ctx));
}

// ---------------------------------------------------------------------------
// In-chat report rendering (plain text; no ascii frame or progress bars)
// ---------------------------------------------------------------------------

function windowLine(label: string, window: LimitWindow | undefined): string {
	if (!window) return `${label}: unavailable`;
	const leftPercent = Math.max(0, 100 - window.usedPercent);
	return `${label}: ${leftPercent}% left (${formatResetTime(window.resetsAt)})`;
}

function renderReport(snapshot: StatusSnapshot, modelLabel: string): string {
	const lines: string[] = [];
	lines.push("OpenAI Codex Status");
	lines.push(`model: ${modelLabel}`);
	if (snapshot.planType) {
		const statusBits: string[] = [];
		if (typeof snapshot.allowed === "boolean")
			statusBits.push(snapshot.allowed ? "allowed" : "blocked");
		if (snapshot.limitReached) statusBits.push("limit reached");
		lines.push(
			`plan: ${snapshot.planType}${statusBits.length > 0 ? ` (${statusBits.join(", ")})` : ""}`,
		);
	}
	lines.push(`Updated ${formatAge(Date.now() - snapshot.fetchedAt)}`);
	lines.push("");
	lines.push(
		windowLine(
			describeWindow(snapshot.primary?.windowSeconds),
			snapshot.primary,
		),
	);
	lines.push(
		windowLine(
			describeWindow(snapshot.secondary?.windowSeconds),
			snapshot.secondary,
		),
	);
	if (snapshot.credits?.hasCredits) {
		const creditsText = snapshot.credits.unlimited
			? "unlimited"
			: snapshot.credits.balance
				? String(snapshot.credits.balance)
				: "available";
		lines.push(`credits: ${creditsText}`);
	}
	if (snapshot.additional.length > 0) {
		lines.push("");
		lines.push("Additional rate limits:");
		for (const limit of snapshot.additional.slice(0, 3)) {
			const label = limit.name || limit.id;
			lines.push(`  - ${windowLine(label, limit.primary ?? limit.secondary)}`);
		}
	}
	lines.push("");
	lines.push(`Visit ${USAGE_PAGE_URL} for up-to-date rate limits and credits.`);
	return lines.join("\n");
}

// ---------------------------------------------------------------------------
// Extension factory
// ---------------------------------------------------------------------------

// Note: no registerMessageRenderer here. A message renderer must return a TUI
// Component (not plain text) and is only used in the in-process TUI, not through
// the daemon bridge. sendMessage() carries the formatted `content`, which every
// daemon/TUI renders as the custom message text.

export default function openAIStatusExtension(pi: ExtensionAPI) {
	pi.registerCommand("status", {
		description: "Show OpenAI Codex subscription usage status",
		handler: async (_args, ctx) => {
			const authStatus = ctx.modelRegistry.getProviderAuthStatus(PROVIDER_ID);
			if (!authStatus.configured) {
				ctx.ui.notify(
					"Not logged into OpenAI Codex. Run /login and choose OpenAI Codex.",
					"error",
				);
				clearWeeklyStatus(ctx);
				return;
			}

			let snapshot: StatusSnapshot;
			try {
				snapshot = await getSnapshot(ctx.modelRegistry);
			} catch (error) {
				const message = error instanceof Error ? error.message : String(error);
				ctx.ui.notify(message, "error");
				clearWeeklyStatus(ctx);
				return;
			}

			const currentModel = ctx.model
				? `${ctx.model.provider}/${ctx.model.id}`
				: "no active model";
			pi.sendMessage({
				customType: REPORT_TYPE,
				content: renderReport(snapshot, currentModel),
				display: true,
			});

			// Refresh the persistent weekly widget too.
			pushWeeklyStatus(ctx, snapshot);
		},
	});

	// Populate the widget when a session starts.
	pi.on("session_start", (_event, ctx) => {
		refreshWeeklyStatus(ctx);
	});

	// Refresh after every completed agent loop.
	pi.on("agent_end", (_event, ctx) => {
		refreshWeeklyStatus(ctx);
	});

	// Re-evaluate when the model/provider changes: hide the widget when leaving
	// openai-codex, show it when switching to an openai-codex model.
	pi.on("model_select", (_event, ctx) => {
		refreshWeeklyStatus(ctx);
	});

	// Clear the widget and status when the session ends.
	pi.on("session_shutdown", (_event, ctx) => {
		clearWeeklyStatus(ctx);
	});
}