/**
 * pi-airpx — login to the airpx LLM proxy and auto-import its model catalog.
 *
 * The auto-loaded catalog is the proxy's MAIN LIST — the anonymous
 * `GET https://airpx.cc/v1/models` response (= `proxy_public_models`, the
 * canonical models left after we skip alias/pseudo rows), WITH pricing and any
 * global discount already applied. Ids and display names are used verbatim, so
 * pi matches the airpx landing page exactly.
 *
 * IMPORTANT: we fetch the list ANONYMOUSLY (no Authorization header) because a
 * key's role is derived from the OWNER's user role — a normal user key may
 * return the FULL catalog instead. The main list is only obtainable anonymously.
 *
 * Escape hatch: set AIRPX_ALL_MODELS=1 to fetch the FULL catalog WITH the key
 * (still skipping alias/pseudo rows).
 *
 * Auth is an `sk-proxy-...` key entered via `/login airpx` (stored in
 * ~/.pi/agent/auth.json), or the AIRPX_API_KEY env var.
 *
 * Cache semantics: the on-disk cache (~/.pi/agent/airpx-catalog.json) is a
 * LAST-GOOD SNAPSHOT, not a union. A successful fetch REPLACES the cache with
 * exactly the fresh response (so removed/re-priced upstream models don't
 * linger, and AIRPX_ALL_MODELS never permanently pollutes the default list).
 * The cache is read for registration ONLY when a fetch fails, so pi still
 * boots with the last-good list.
 *
 * `/login` additionally validates the entered key with an explicit KEYED probe
 * before saving credentials — a bad/mistyped key is rejected, not silently
 * saved.
 *
 * Usage:
 *   pi -e ./extensions/index.ts        # then: /login airpx  → paste sk-proxy-...
 *   AIRPX_API_KEY=sk-proxy-... pi -e ./extensions/index.ts
 */

import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { dirname, join } from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import type {
	Api,
	Model,
	OAuthCredentials,
	OAuthLoginCallbacks,
} from "@earendil-works/pi-ai";
import {
	type CatalogDeps,
	KeyRejectedError,
	PROVIDER,
	parseCachedModels,
	resolveCatalog,
	resolveKey as resolveKeyPolicy,
} from "../src/catalog.ts";
import { type PiModel, type ProxyModel, mapCatalog, pickerPatterns } from "../src/map.ts";

// PROVIDER is imported from src/catalog.ts: the provider id doubles as the
// auth.json key, so the registration below and the credential lookup there must
// never drift apart.
const DEFAULT_BASE_URL = "https://airpx.cc/v1";
const FETCH_TIMEOUT_MS = 15_000;
const CACHE_PATH = join(homedir(), ".pi", "agent", "airpx-catalog.json");

function baseUrl(): string {
	return process.env.AIRPX_BASE_URL?.replace(/\/$/, "") || DEFAULT_BASE_URL;
}

function useAllModels(): boolean {
	return process.env.AIRPX_ALL_MODELS === "1";
}

// ── injected I/O (policy lives in src/catalog.ts) ───────────────────────────

/** Resolve the credential from auth.json, then the environment. */
function resolveKey(): string | undefined {
	let authBody: string | undefined;
	try {
		authBody = readFileSync(join(homedir(), ".pi", "agent", "auth.json"), "utf8");
	} catch {
		authBody = undefined;
	}
	return resolveKeyPolicy(authBody, process.env);
}

/** Build the real side effects for resolveCatalog(). */
function catalogDeps(key: string | undefined): CatalogDeps {
	return {
		fetchModels: () => fetchCatalog(key),
		readCache: () => parseCachedModels(readFileSync(CACHE_PATH, "utf8")),
		writeCache: (models) => {
			mkdirSync(dirname(CACHE_PATH), { recursive: true });
			// Persist the RAW rows next to the mapped models: picker policy
			// (picker_enabled / policy.state) lives on fields PiModel does not
			// carry, and a cache-only start must still honour server policy.
			writeFileSync(
				CACHE_PATH,
				JSON.stringify(
					{ models, rows: lastFetchedRows, etag: lastEtag, checkedAt: Date.now() },
					null,
					2,
				),
			);
		},
		log: (message) => console.error(`[pi-airpx] ${message}`),
	};
}

// ── fetch ───────────────────────────────────────────────────────────────────

/**
 * Fetch the model catalog and map it.
 *
 * Default (main list): anonymous request, no Authorization header.
 * AIRPX_ALL_MODELS=1: full catalog using the key.
 *
 * Throws KeyRejectedError on 401/403 (only reachable in the keyed path).
 */
async function fetchCatalog(key?: string, external?: AbortSignal): Promise<PiModel[]> {
	const withKey = useAllModels() && Boolean(key);
	const headers: Record<string, string> = withKey
		? { Authorization: `Bearer ${key}` }
		: {};

	// Conditional request, mirroring pi's own models-store.json (which keeps an
	// etag per provider): an unchanged catalog costs a 304 instead of a full
	// re-download. Fail-open — a missing/!ok etag simply means a normal fetch.
	const cachedTag = readCachedEtag();
	if (cachedTag) headers["If-None-Match"] = cachedTag;

	const controller = new AbortController();
	const timer = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS);
	// pi documents that provider callbacks MUST pass its signal to blocking I/O:
	// cancellation has to stop the caller waiting. We keep our own timeout as
	// well, so abort on whichever fires first.
	const onExternalAbort = () => controller.abort();
	external?.addEventListener("abort", onExternalAbort, { once: true });
	if (external?.aborted) controller.abort();
	let res: Response;
	try {
		res = await fetch(`${baseUrl()}/models`, { headers, signal: controller.signal });
	} finally {
		clearTimeout(timer);
		external?.removeEventListener("abort", onExternalAbort);
	}

	if (res.status === 304) {
		// Unchanged: reuse the cached rows so policy still applies this run.
		lastFetchedRows = readCachedRows();
		return parseCachedModels(readFileSync(CACHE_PATH, "utf8")) ?? [];
	}

	if (res.status === 401 || res.status === 403) {
		throw new KeyRejectedError(`key rejected (HTTP ${res.status})`);
	}
	if (!res.ok) {
		throw new Error(
			`airpx /v1/models -> HTTP ${res.status} ${await res.text().catch(() => "")}`.trim(),
		);
	}
	const payload = (await res.json()) as { data?: ProxyModel[] };
	// Keep the RAW rows: picker policy lives on fields toPiModel() does not carry.
	lastFetchedRows = payload.data ?? [];
	lastEtag = res.headers.get("etag") ?? "";
	return mapCatalog(lastFetchedRows);
}

/**
 * Explicit KEYED validation probe used by `/login`: `GET /models` WITH the key
 * (unlike the default anon catalog fetch, which never surfaces auth errors).
 *
 * - 401/403 → throw KeyRejectedError (caller rejects the login; nothing saved).
 * - Non-auth failure (network/timeout/5xx) → resolve without throwing so we do
 *   NOT block login; the caller degrades to the last-good cache.
 */
async function validateKey(k: string): Promise<void> {
	const controller = new AbortController();
	const timer = setTimeout(() => controller.abort(), FETCH_TIMEOUT_MS);
	let res: Response;
	try {
		// NOTE: GET /v1/models is ANONYMOUS (that's how we fetch the main list),
		// so it does NOT validate the key. Use POST /v1/messages/count_tokens —
		// it requires a valid sk-proxy key (401 on bad key) and is FREE (token
		// counting, no upstream inference / no spend).
		//
		// The body carries NO `model`: the proxy authenticates before it looks at
		// the body and treats `model` as optional (unknown/absent → local token
		// estimate, still 200), so naming a model here would be a 4th hardcoded
		// copy of a catalog id that silently rots when that model is retired —
		// exactly the drift this plugin no longer tolerates. Verified live: with a
		// valid key `{}`, a bogus model, and no model all return 200, while a bad
		// key returns 401 in every case.
		res = await fetch(`${baseUrl()}/messages/count_tokens`, {
			method: "POST",
			headers: {
				Authorization: `Bearer ${k}`,
				"Content-Type": "application/json",
			},
			body: JSON.stringify({ messages: [{ role: "user", content: "hi" }] }),
			signal: controller.signal,
		});
	} catch (err) {
		// Network/timeout: don't block login (graceful degradation).
		console.error(
			`[pi-airpx] could not verify key (network); saving anyway: ${(err as Error).message}`,
		);
		return;
	} finally {
		clearTimeout(timer);
	}
	if (res.status === 401 || res.status === 403) {
		throw new KeyRejectedError(`key rejected (HTTP ${res.status})`);
	}
	if (!res.ok) {
		// Server-side hiccup (5xx etc.): don't block login either.
		console.error(`[pi-airpx] could not verify key (HTTP ${res.status}); saving anyway`);
	}
}

// ── pi-native catalog refresh ────────────────────────────────────────────────

/** Minimal shape of pi's refreshModels context (only what we consume). */
interface RefreshContext {
	signal?: AbortSignal;
	stored?: { models?: unknown } | null;
	publish?: (update: { persist?: unknown }) => void;
}

/**
 * Pi-driven catalog refresh.
 *
 * Contract notes that matter:
 *   * The returned list REPLACES the registered models, so every failure path
 *     must return a non-empty list when one is known. Returning [] on a network
 *     blip would empty the model picker.
 *   * `context.stored` is pi's own persisted snapshot for this provider — the
 *     same store the built-in providers use. We prefer it over our legacy cache
 *     file so there is one source of truth going forward.
 *   * Persisting is generation-checked by pi via `context.publish`, which is why
 *     we hand the snapshot back instead of writing a file ourselves.
 */
async function refreshCatalog(
	key: string | undefined,
	context: RefreshContext,
): Promise<PiModel[]> {
	try {
		const fresh = await fetchCatalog(key, context.signal);
		if (fresh.length === 0) throw new Error("empty catalog");
		syncPickerScope(lastFetchedRows.length > 0 ? lastFetchedRows : readCachedRows());
		try {
			context.publish?.({
				persist: {
					models: fresh,
					checkedAt: Date.now(),
					...(lastEtag ? { etag: lastEtag } : {}),
				},
			});
		} catch {
			/* persisting is best-effort; the live list is what matters */
		}
		return fresh;
	} catch (err) {
		// Keep whatever we already had: pi's snapshot first, then our own cache.
		const stored = Array.isArray(context.stored?.models)
			? (context.stored?.models as PiModel[])
			: [];
		if (stored.length > 0) return stored;
		const cached = attemptRead();
		if (cached.length > 0) return cached;
		console.error(`[pi-airpx] refresh failed and no catalog is known: ${String(err)}`);
		return [];
	}
}

/** Legacy cache read, kept as the last resort during the transition. */
function attemptRead(): PiModel[] {
	try {
		return parseCachedModels(readFileSync(CACHE_PATH, "utf8")) ?? [];
	} catch {
		return [];
	}
}

// ── picker scope sync (server-decided, Copilot-style) ────────────────────────
//
// The proxy owns the "offer this model by default" decision (picker_enabled /
// policy.state on /v1/models). pi resolves its Ctrl+P scope from `enabledModels`
// at session start and exposes it read-only, so the only way to honour server
// policy is to keep that setting in sync on disk.
//
// Rules, chosen so this can never surprise the user:
//   * ONLY `airpx/...` entries are ever touched. Anything else in the list is
//     preserved verbatim — this is a user setting we are borrowing, not owning.
//   * A missing `enabledModels` is left missing: absent means "everything is in
//     scope", and writing a list would silently NARROW the user's picker.
//   * Opt-out via LLM_PROXY_NO_PICKER_SYNC=1.
//   * Best-effort: any error is swallowed, because a settings file is never worth
//     breaking a session over.
// Takes effect on the next start, since pi reads the scope once per session.

const SETTINGS_PATH = join(homedir(), ".pi", "agent", "settings.json");

/** Raw rows from the most recent successful fetch; empty when only cache was used. */
let lastFetchedRows: ProxyModel[] = [];

/** Etag from the most recent successful fetch, persisted alongside the cache. */
let lastEtag = "";

/** Etag persisted by the last successful fetch. Best-effort, never throws. */
function readCachedEtag(): string {
	try {
		const parsed = JSON.parse(readFileSync(CACHE_PATH, "utf8")) as { etag?: unknown };
		return typeof parsed.etag === "string" ? parsed.etag : "";
	} catch {
		return "";
	}
}

/** Raw rows persisted by the last successful fetch. Best-effort, never throws. */
function readCachedRows(): ProxyModel[] {
	try {
		const parsed = JSON.parse(readFileSync(CACHE_PATH, "utf8")) as {
			rows?: unknown;
		};
		return Array.isArray(parsed.rows) ? (parsed.rows as ProxyModel[]) : [];
	} catch {
		return [];
	}
}

function syncPickerScope(rows: ProxyModel[]): void {
	if (process.env.LLM_PROXY_NO_PICKER_SYNC === "1") return;
	try {
		const raw = readFileSync(SETTINGS_PATH, "utf8");
		const settings = JSON.parse(raw) as { enabledModels?: unknown };
		const current = settings.enabledModels;
		// Not configured ⇒ every model is in scope. Do not narrow it.
		if (!Array.isArray(current)) return;
		const foreign = current.filter(
			(e) => typeof e === "string" && !e.startsWith("airpx/"),
		) as string[];
		const wanted = pickerPatterns(rows);
		if (wanted.length === 0) return; // catalog unavailable — keep what we have
		const next = [...foreign, ...wanted];
		const prev = current.filter((e) => typeof e === "string") as string[];
		if (prev.length === next.length && prev.every((v, i) => v === next[i])) return;
		settings.enabledModels = next;
		writeFileSync(SETTINGS_PATH, `${JSON.stringify(settings, null, 2)}\n`);
	} catch {
		/* best-effort only */
	}
}

// ── provider registration ────────────────────────────────────────────────────

/** Register (or replace) the provider with a given (possibly empty) model list. */
function register(pi: ExtensionAPI, key: string | undefined, models: PiModel[]): void {
	pi.registerProvider(PROVIDER, {
		name: "airpx (LLM proxy)",
		baseUrl: baseUrl(),
		api: "openai-completions" as Api,
		// apiKey is used when logging in via env/auth-file api_key; oauth.getApiKey
		// takes over when the credential is an oauth entry from /login.
		...(key ? { apiKey: key } : {}),
		models: models as unknown as Model<Api>[] as never,
		// Pi's own catalog-refresh hook. Registering it means airpx participates
		// in `pi update --models` and in any runtime refresh exactly like the
		// built-in providers, whose catalogs live in models-store.json with an
		// etag apiece — instead of this plugin owning a private cache file and a
		// private refresh moment.
		//
		// Returned models REPLACE the ones passed above, so a refresh that fails
		// must return the current list rather than an empty one: an empty return
		// would wipe the picker (the failure mode that emptied opencode's Copilot
		// provider). See refreshCatalog().
		refreshModels: (async (context: RefreshContext) =>
			await refreshCatalog(key, context)) as never,
		oauth: {
			name: "airpx (paste sk-proxy key)",
			async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
				const entered = await callbacks.onPrompt({
					message: "Paste your airpx API key (sk-proxy-...):",
				});
				const k = entered.trim();
				if (!k) throw new Error("No key entered");
				// Validate the key with an explicit KEYED probe BEFORE saving —
				// otherwise a bogus key would be saved silently (the default
				// catalog fetch is anonymous and never surfaces auth errors).
				try {
					await validateKey(k);
				} catch (err) {
					if (err instanceof KeyRejectedError) {
						// Surface a clear error and DO NOT save credentials.
						throw new Error(
							"airpx key rejected — check your sk-proxy key and try again.",
						);
					}
					throw err;
				}
				// Key is valid: register the DEFAULT catalog (anon main list unless
				// AIRPX_ALL_MODELS=1). We validated WITH the key, but the registered
				// model set stays the anon main list by default — we do NOT register
				// the full keyed catalog just because validation used the key. Live
				// re-register applies immediately (no /reload). Never throws.
				await loadAndRegister(pi, k);
				return {
					refresh: k,
					access: k,
					expires: Date.now() + 100 * 365 * 24 * 3600 * 1000,
				};
			},
			async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
				return credentials; // sk-proxy keys don't expire
			},
			getApiKey(credentials: OAuthCredentials): string {
				return credentials.access;
			},
		},
	} as never);
}

/**
 * Shared "fetch default catalog → replace cache → register" flow used by BOTH
 * the startup factory and `/login`.
 *
 * Cache contract:
 *   - SUCCESS ⇒ the fresh response REPLACES the cache (last-good snapshot);
 *     register exactly the fresh list.
 *   - FAILURE with a cache present ⇒ register the last-good cache so pi boots.
 *   - FAILURE with no cache ⇒ register an empty list (pi still boots; /login
 *     stays available).
 *
 * NEVER throws — the factory must not throw and login must not fail on a
 * transient catalog error after the key was already validated.
 */
async function loadAndRegister(pi: ExtensionAPI, key: string | undefined): Promise<void> {
	const { models } = await resolveCatalog(catalogDeps(key));
	// Register regardless — even an empty list keeps /login airpx working.
	// Registration is NOT filtered by picker policy: a deprecated model stays
	// addressable, it is merely not offered by default.
	register(pi, key, models);
	// Prefer this run's fresh rows; fall back to the cached ones so a start that
	// could not reach the proxy still applies the last-known server policy
	// instead of silently skipping the sync.
	syncPickerScope(lastFetchedRows.length > 0 ? lastFetchedRows : readCachedRows());

	// First-run onboarding hint: with no key configured the models are loaded but
	// pi keeps them UNAVAILABLE in /model until auth is set (documented pi
	// behavior). Tell the user exactly what to do instead of showing an empty
	// picker with no explanation.
	if (!key) {
		console.error(
			`[pi-airpx] no API key configured — run "/login airpx" and paste your ` +
				`sk-proxy key to enable ${models.length || "the"} model(s). ` +
				`Get a key at https://airpx.cc`,
		);
	}
}

// ── async factory (must NEVER throw) ─────────────────────────────────────────

export default async function (pi: ExtensionAPI): Promise<void> {
	await loadAndRegister(pi, resolveKey());
}