/**
 * Google Antigravity provider extension.
 *
 * Provides Google Cloud Code Assist (Antigravity) login, live model discovery,
 * and streaming. The wire protocol, OAuth flow, header spoofing, and request
 * envelope mirror oh-my-pi's google-antigravity implementation.
 *
 * Usage:
 *   prime-agent -e ./packages/coding-agent/examples/extensions/google-antigravity
 *   # Then /login google-antigravity
 *
 * Or drop the directory into ~/.prime/agent/extensions/ for auto-discovery.
 */
import { readFileSync, writeFileSync } from "node:fs";
import { homedir } from "node:os";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import type { Api, AssistantMessageEventStream, Context, Model } from "@earendil-works/pi-ai";
import type { ExtensionAPI, ProviderModelConfig } from "@earendil-works/pi-coding-agent";
import { fetchAntigravityModels } from "./discovery.js";
import type { AntigravitySimpleStreamOptions } from "./helpers.js";
import { STATIC_PROVIDER_MODELS } from "./models.js";
import { antigravityOAuth } from "./oauth.js";
import { streamSimpleAntigravity } from "./stream.js";
import { ANTIGRAVITY_DAILY_ENDPOINT } from "./wire.js";

export const PROVIDER_ID = "google-antigravity";
export const PROVIDER_API = "google-antigravity" as Api;

/**
 * Mirror prime's getAgentDir() without a runtime import of
 * @earendil-works/pi-coding-agent (its dist barrel pulls pi-ai/mcp through
 * jiti and fails to load inside an extension).
 */
function getAgentDir(): string {
	const envDir = process.env.PRIME_AGENT_CODING_AGENT_DIR ?? process.env.PI_CODING_AGENT_DIR;
	if (envDir) {
		if (envDir === "~") return homedir();
		if (envDir.startsWith("~/")) return join(homedir(), envDir.slice(2));
		return envDir;
	}
	return join(homedir(), ".prime", "agent");
}

interface StoredCredential {
	type?: string;
	access?: string;
}

/**
 * auth.json stores a bare credential per provider, or an array once more than
 * one account is logged in (multi-account credential pooling).
 */
function readExistingAccessToken(): string | undefined {
	try {
		const authPath = join(getAgentDir(), "auth.json");
		const auth = JSON.parse(readFileSync(authPath, "utf-8")) as Record<string, StoredCredential | StoredCredential[]>;
		const stored = auth[PROVIDER_ID];
		const credentials = Array.isArray(stored) ? stored : [stored];
		const entry = credentials.find(
			(credential) =>
				credential?.type === "oauth" && typeof credential.access === "string" && credential.access.length > 0,
		);
		if (entry) return entry.access;
	} catch {
		// No auth file or not logged in yet.
	}
	return undefined;
}

const DISCOVERY_CACHE_FILE = "models-cache.json";
const DISCOVERY_CACHE_TTL_MS = 6 * 60 * 60 * 1000;

function discoveryCachePath(): string {
	return join(dirname(fileURLToPath(import.meta.url)), DISCOVERY_CACHE_FILE);
}

/**
 * A cached empty list records a failed discovery: it keeps the next start off the
 * network for one TTL while leaving the static models in place.
 */
function readDiscoveryCache(): { models?: ProviderModelConfig[]; fresh: boolean } | undefined {
	try {
		const cached = JSON.parse(readFileSync(discoveryCachePath(), "utf-8")) as {
			fetchedAt?: unknown;
			models?: unknown;
		};
		if (typeof cached.fetchedAt !== "string" || !Array.isArray(cached.models)) {
			return undefined;
		}
		const fetchedAt = Date.parse(cached.fetchedAt);
		return {
			...(cached.models.length > 0 ? { models: cached.models as ProviderModelConfig[] } : {}),
			fresh: Number.isFinite(fetchedAt) && Date.now() - fetchedAt < DISCOVERY_CACHE_TTL_MS,
		};
	} catch {
		return undefined;
	}
}

function writeDiscoveryCache(models: ProviderModelConfig[]): void {
	try {
		writeFileSync(
			discoveryCachePath(),
			`${JSON.stringify({ fetchedAt: new Date().toISOString(), models }, null, "\t")}\n`,
		);
	} catch (error) {
		// The cache is an optimization; a read-only extension dir must not break loading.
		console.error(`google-antigravity: could not write ${DISCOVERY_CACHE_FILE}: ${String(error)}`);
	}
}

function registerAntigravityProvider(pi: ExtensionAPI, baseUrl: string, models: ProviderModelConfig[]): void {
	pi.registerProvider(PROVIDER_ID, {
		name: "Google Antigravity",
		baseUrl,
		api: PROVIDER_API,
		models,
		oauth: antigravityOAuth,
		streamSimple: (model: Model<Api>, context: Context, options?: AntigravitySimpleStreamOptions): AssistantMessageEventStream =>
			streamSimpleAntigravity(model, context, options),
	});
}

/**
 * Live discovery is a request to Cloud Code Assist, so it must never hold up
 * extension loading: both the client and every session worker load this module,
 * and awaiting it added ~0.8s to each of them. Models come from the disk cache
 * when it is fresh, otherwise from the static list, and a stale cache refreshes
 * in the background and re-registers (the registry upserts providers).
 */
export default async function (pi: ExtensionAPI) {
	const baseUrl = ANTIGRAVITY_DAILY_ENDPOINT;
	const cached = readDiscoveryCache();
	registerAntigravityProvider(pi, baseUrl, cached?.models ?? STATIC_PROVIDER_MODELS);

	const accessToken = readExistingAccessToken();
	if (!accessToken || cached?.fresh === true) {
		return;
	}
	void (async () => {
		try {
			const discovered = await fetchAntigravityModels(accessToken);
			writeDiscoveryCache(discovered ?? []);
			if (discovered && discovered.length > 0) {
				registerAntigravityProvider(pi, baseUrl, discovered);
			}
		} catch {
			// Keep the models already registered. A failed fetch stays uncached so the
			// next start retries instead of hiding a transient outage for a whole TTL.
		}
	})();
}