/**
 * Pi extension: xAI / Grok OAuth via shared Grok Build auth (no API key).
 *
 * Model after systray-ai `pi_auth_sync.py` (Codex → Pi):
 * - Source of truth = CLI auth file (`~/.grok/auth.json`, same OIDC client as `grok login`)
 * - Login DERIVES a copy into Pi's auth.json — no token endpoint, no second OAuth session
 * - refreshToken RE-DERIVES from the Grok file (picks up CLI-side refreshes)
 * - NEVER independently rotates refresh tokens (xAI refresh is single-use / revokes prior)
 *
 * Install:
 *   pi install /absolute/path/to/packages/pi-xai-oauth
 *   /login xai-oauth   # or: node import-from-grok.mjs
 */

import type { OAuthCredentials, OAuthLoginCallbacks } from "@earendil-works/pi-ai";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { homedir } from "node:os";
import { join } from "node:path";
import { readFile } from "node:fs/promises";

const PROVIDER_ID = "xai-oauth";
const PROVIDER_NAME = "xAI (Grok Build auth)";
const BASE_URL = "https://api.x.ai/v1";
/** Skew only for "is this still usable" checks — stored expires comes from JWT `exp`. */
const EXPIRY_SKEW_MS = 30_000;

type GrokAuthEntry = {
	key?: string;
	access_token?: string;
	refresh_token?: string;
	expires_at?: string;
	oidc_client_id?: string;
	auth_mode?: string;
	user_id?: string;
	principal_id?: string;
	team_id?: string;
};

type Derived = OAuthCredentials & { accountId?: string };

// Only models user wants — verified against api.x.ai chat with Grok OAuth.
const MODELS = [
	{
		id: "grok-4.5",
		name: "Grok 4.5",
		reasoning: true,
		input: ["text", "image"] as ("text" | "image")[],
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
		contextWindow: 256000,
		maxTokens: 65536,
	},
	{
		id: "grok-composer-2.5-fast",
		name: "Composer 2.5",
		reasoning: false,
		input: ["text", "image"] as ("text" | "image")[],
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
		contextWindow: 256000,
		maxTokens: 65536,
	},
];

function grokAuthPath(): string {
	return process.env.GROK_AUTH_FILE || join(homedir(), ".grok", "auth.json");
}

function decodeJwtPayload(access: string): Record<string, unknown> {
	const parts = access.split(".");
	if (parts.length < 2) throw new Error("access token is not a JWT");
	const segment = parts[1];
	const padding = "=".repeat((4 - (segment.length % 4)) % 4);
	const json = Buffer.from(segment + padding, "base64url").toString("utf8");
	const claims = JSON.parse(json) as Record<string, unknown>;
	if (!claims || typeof claims !== "object") throw new Error("invalid JWT payload");
	return claims;
}

/**
 * Derive Pi oauth credential from Grok Build auth.json — COPY only, no network.
 * Matches systray-ai PiAuthSync._derive_credential discipline for Codex.
 */
async function deriveFromGrokAuth(): Promise<Derived> {
	const path = grokAuthPath();
	let raw: string;
	try {
		raw = await readFile(path, "utf8");
	} catch {
		throw new Error(
			`Grok auth not found at ${path}. Run \`grok login\` once, then /login ${PROVIDER_ID} again.`,
		);
	}

	let payload: Record<string, GrokAuthEntry>;
	try {
		payload = JSON.parse(raw) as Record<string, GrokAuthEntry>;
	} catch {
		throw new Error(`Invalid JSON in ${path}`);
	}

	let chosen: GrokAuthEntry | null = null;
	for (const [slot, entry] of Object.entries(payload)) {
		if (!entry || typeof entry !== "object") continue;
		const access = (entry.key || entry.access_token || "").trim();
		const refresh = (entry.refresh_token || "").trim();
		if (!access || !refresh) continue;
		if (!slot.includes("auth.x.ai") && !entry.oidc_client_id) continue;
		chosen = entry;
		break;
	}
	if (!chosen) {
		throw new Error(
			`No usable Grok OAuth entry in ${path} (need access + refresh). Run \`grok login\`.`,
		);
	}

	const access = (chosen.key || chosen.access_token || "").trim();
	const refresh = (chosen.refresh_token || "").trim();
	if (!access || !refresh) {
		throw new Error(`Empty tokens in ${path}`);
	}

	const claims = decodeJwtPayload(access);
	const expSec = claims.exp;
	if (typeof expSec !== "number" || !Number.isFinite(expSec)) {
		throw new Error("Grok access token missing numeric JWT exp — refuse to guess lifetime");
	}

	const accountId =
		(typeof chosen.principal_id === "string" && chosen.principal_id) ||
		(typeof claims.principal_id === "string" && claims.principal_id) ||
		(typeof claims.sub === "string" && claims.sub) ||
		undefined;

	const out: Derived = {
		access,
		refresh,
		// Pi stores expires as epoch milliseconds (see systray-ai pi_auth_sync).
		expires: Math.trunc(expSec * 1000),
	};
	if (accountId) out.accountId = accountId;
	return out;
}

async function loginFromGrok(_callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
	// No onSelect / device code: single path = derive shared Grok session.
	// Device code would mint a second session and can revoke the only login.
	return deriveFromGrokAuth();
}

/**
 * Re-copy from Grok auth file. Does NOT call auth.x.ai.
 * If Grok CLI refreshed the file, Pi picks it up. If access is already expired,
 * fail closed — do not rotate refresh here (would kill Grok Build + Pi together).
 */
async function refreshFromGrokFile(_credentials: OAuthCredentials): Promise<OAuthCredentials> {
	const derived = await deriveFromGrokAuth();
	if (derived.expires <= Date.now() + EXPIRY_SKEW_MS) {
		throw new Error(
			"Grok access token expired in ~/.grok/auth.json. " +
				"Refresh by using the Grok CLI (or `grok login` if needed), " +
				"then /login xai-oauth again. " +
				"This extension never hits the token endpoint — dual refresh revokes the shared session.",
		);
	}
	return derived;
}

export default function (pi: ExtensionAPI) {
	pi.registerProvider(PROVIDER_ID, {
		name: PROVIDER_NAME,
		baseUrl: BASE_URL,
		api: "openai-completions",
		authHeader: true,
		compat: {
			supportsStore: false,
			supportsDeveloperRole: false,
			supportsReasoningEffort: false,
		},
		models: MODELS,
		oauth: {
			name: PROVIDER_NAME,
			login: loginFromGrok,
			refreshToken: refreshFromGrokFile,
			getApiKey(credentials: OAuthCredentials) {
				return credentials.access;
			},
		},
	});
}