/**
 * SuperGrok / xAI OAuth 2.0 + PKCE (browser loopback + device code) and token refresh.
 */

import { createServer, type Server } from "node:http";
import type {
	OAuthCredentials,
	OAuthLoginCallbacks,
} from "@earendil-works/pi-ai";
import {
	AUTHORIZE_URL,
	CLIENT_ID,
	DEVICE_CODE_GRANT,
	DEVICE_CODE_URL,
	FORM_HEADERS,
	OAUTH_HOST,
	OAUTH_PORT,
	OAUTH_REDIRECT_PATH,
	REDIRECT_URI,
	REFRESH_SKEW_MS,
	SCOPE,
	TOKEN_URL,
} from "./constants.ts";

// -----------------------------------------------------------------------------
// PKCE helpers
// -----------------------------------------------------------------------------

function base64Url(bytes: Uint8Array): string {
	return btoa(String.fromCharCode(...bytes))
		.replace(/\+/g, "-")
		.replace(/\//g, "_")
		.replace(/=+$/, "");
}

async function generatePKCE(): Promise<{ verifier: string; challenge: string }> {
	const verifier = base64Url(crypto.getRandomValues(new Uint8Array(48)));
	const hash = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(verifier));
	return { verifier, challenge: base64Url(new Uint8Array(hash)) };
}

function randomState(): string {
	return base64Url(crypto.getRandomValues(new Uint8Array(32)));
}

// -----------------------------------------------------------------------------
// Token shapes
// -----------------------------------------------------------------------------

interface TokenResponse {
	access_token: string;
	refresh_token?: string;
	id_token?: string;
	token_type?: string;
	expires_in?: number;
	scope?: string;
}

/** Cap untrusted error bodies before they land in logs / thrown Errors. */
function truncateDetail(detail: string, max = 200): string {
	const cleaned = detail.replace(/\s+/g, " ").trim();
	if (cleaned.length <= max) return cleaned;
	return `${cleaned.slice(0, max)}…`;
}

function toCredentials(tokens: TokenResponse, fallbackRefresh = ""): OAuthCredentials {
	if (!tokens.access_token || typeof tokens.access_token !== "string") {
		throw new Error("xAI token response missing access_token");
	}
	const expiresIn =
		typeof tokens.expires_in === "number" && Number.isFinite(tokens.expires_in) && tokens.expires_in > 0
			? tokens.expires_in
			: 3600;
	const refresh =
		typeof tokens.refresh_token === "string" && tokens.refresh_token
			? tokens.refresh_token
			: fallbackRefresh;
	return {
		access: tokens.access_token,
		refresh,
		expires: Date.now() + expiresIn * 1000 - REFRESH_SKEW_MS,
	};
}

function tierError(detail: string): string {
	const trimmed = truncateDetail(detail);
	return (
		"xAI returned 403. OAuth API access may be restricted to certain SuperGrok tiers. " +
		"If login succeeds in the browser but requests 403, use an XAI_API_KEY provider instead." +
		(trimmed ? ` (${trimmed})` : "")
	);
}

// -----------------------------------------------------------------------------
// Authorization-code (browser / loopback) flow
// -----------------------------------------------------------------------------

function buildAuthorizeUrl(challenge: string, state: string, nonce: string): string {
	const params = new URLSearchParams({
		response_type: "code",
		client_id: CLIENT_ID,
		redirect_uri: REDIRECT_URI,
		scope: SCOPE,
		code_challenge: challenge,
		code_challenge_method: "S256",
		state,
		nonce,
		// `plan=generic` opts the consent screen into xAI's generic OAuth plan
		// tier; without it accounts.x.ai rejects loopback OAuth for this client.
		plan: "generic",
		referrer: "pi",
	});
	return `${AUTHORIZE_URL}?${params.toString()}`;
}

async function exchangeCode(code: string, verifier: string): Promise<TokenResponse> {
	const res = await fetch(TOKEN_URL, {
		method: "POST",
		headers: FORM_HEADERS,
		body: new URLSearchParams({
			grant_type: "authorization_code",
			code,
			redirect_uri: REDIRECT_URI,
			client_id: CLIENT_ID,
			code_verifier: verifier,
		}).toString(),
	});
	if (!res.ok) {
		const detail = truncateDetail(await res.text().catch(() => ""));
		if (res.status === 403) throw new Error(tierError(detail));
		throw new Error(`xAI token exchange failed (${res.status})${detail ? `: ${detail}` : ""}`);
	}
	return (await res.json()) as TokenResponse;
}

const SUCCESS_HTML = `<!doctype html><html><head><meta charset="utf-8"><title>pi · xAI login</title>
<style>body{font-family:system-ui,sans-serif;display:flex;justify-content:center;align-items:center;height:100vh;margin:0;background:#131010;color:#f1ecec}.c{text-align:center;padding:2rem}h1{margin-bottom:1rem}p{color:#b7b1b1}</style>
</head><body><div class="c"><h1>Login successful</h1><p>You can close this window and return to pi.</p></div>
<script>setTimeout(()=>window.close(),1500)</script></body></html>`;

function errorHtml(msg: string): string {
	const safe = msg.replace(/[&<>"']/g, (c) =>
		({ "&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;", "'": "&#39;" })[c] ?? c,
	);
	return `<!doctype html><html><head><meta charset="utf-8"><title>pi · xAI login failed</title>
<style>body{font-family:system-ui,sans-serif;display:flex;justify-content:center;align-items:center;height:100vh;margin:0;background:#131010;color:#f1ecec}.c{text-align:center;padding:2rem}h1{color:#fc533a;margin-bottom:1rem}code{color:#ff917b;font-family:monospace;display:block;margin-top:1rem;padding:1rem;background:#3c140d;border-radius:.5rem}</style>
</head><body><div class="c"><h1>Login failed</h1><code>${safe}</code></div></body></html>`;
}

/** Run a one-shot loopback server that resolves with the authorization code. */
function waitForLoopbackCode(state: string): { promise: Promise<string>; close: () => void } {
	let server: Server | undefined;
	let timer: ReturnType<typeof setTimeout> | undefined;
	/** After the first terminal /callback response, ignore further hits. */
	let settled = false;

	const close = () => {
		if (timer) {
			clearTimeout(timer);
			timer = undefined;
		}
		server?.close();
		server = undefined;
	};

	const promise = new Promise<string>((resolve, reject) => {
		const settle = (fn: () => void) => {
			if (settled) return;
			settled = true;
			if (timer) {
				clearTimeout(timer);
				timer = undefined;
			}
			fn();
			// Drop the listener as soon as we have a terminal outcome so a second
			// request (browser prefetch / double-hit) cannot race the code exchange.
			queueMicrotask(() => close());
		};

		server = createServer((req, res) => {
			if (settled) {
				res.writeHead(410, { "Content-Type": "text/plain" });
				res.end("Already handled");
				return;
			}

			const url = new URL(req.url || "/", `http://${OAUTH_HOST}:${OAUTH_PORT}`);
			if (url.pathname !== OAUTH_REDIRECT_PATH) {
				res.writeHead(404);
				res.end("Not found");
				return;
			}

			const err = url.searchParams.get("error");
			const errDesc = url.searchParams.get("error_description");
			const code = url.searchParams.get("code");
			const gotState = url.searchParams.get("state");

			if (err) {
				const m = truncateDetail(errDesc || err, 400);
				res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
				res.end(errorHtml(m));
				settle(() => reject(new Error(m)));
				return;
			}
			if (!code) {
				res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
				res.end(errorHtml("Missing authorization code"));
				settle(() => reject(new Error("Missing authorization code")));
				return;
			}
			if (gotState !== state) {
				res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
				res.end(errorHtml("State mismatch (possible CSRF)"));
				settle(() => reject(new Error("OAuth state mismatch")));
				return;
			}
			res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
			res.end(SUCCESS_HTML);
			settle(() => resolve(code));
		});

		server.once("error", (e) => {
			settle(() => reject(e));
		});
		server.listen(OAUTH_PORT, OAUTH_HOST);

		timer = setTimeout(
			() => settle(() => reject(new Error("OAuth callback timed out after 5 minutes"))),
			5 * 60 * 1000,
		);
	});

	return { promise, close };
}

async function loginBrowser(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
	const { verifier, challenge } = await generatePKCE();
	const state = randomState();
	const nonce = randomState();

	let listener: { promise: Promise<string>; close: () => void };
	try {
		listener = waitForLoopbackCode(state);
	} catch (e) {
		throw new Error(
			`Could not bind ${REDIRECT_URI} (${(e as Error).message}). ` +
				"Close anything using that port, or use the device-code login instead.",
		);
	}

	try {
		callbacks.onAuth({
			url: buildAuthorizeUrl(challenge, state, nonce),
			instructions:
				"Approve in the browser. xAI may show a page titled \"Grok Build\" with a " +
				"code to copy \u2014 ignore it and just close the tab; pi captures the login " +
				"automatically on 127.0.0.1:56121.",
		});
		const code = await listener.promise;
		const tokens = await exchangeCode(code, verifier);
		return toCredentials(tokens);
	} finally {
		listener.close();
	}
}

// -----------------------------------------------------------------------------
// Device-code flow (headless)
// -----------------------------------------------------------------------------

interface DeviceCodeResponse {
	device_code: string;
	user_code: string;
	verification_uri: string;
	verification_uri_complete?: string;
	expires_in?: number;
	interval?: number;
}

async function loginDevice(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
	const res = await fetch(DEVICE_CODE_URL, {
		method: "POST",
		headers: FORM_HEADERS,
		body: new URLSearchParams({ client_id: CLIENT_ID, scope: SCOPE }).toString(),
	});
	if (!res.ok) {
		const detail = truncateDetail(await res.text().catch(() => ""));
		throw new Error(`xAI device code request failed (${res.status})${detail ? `: ${detail}` : ""}`);
	}
	const device = (await res.json()) as DeviceCodeResponse;
	if (!device.device_code || !device.user_code || !device.verification_uri) {
		throw new Error("xAI device code response was malformed");
	}

	const intervalSeconds = Math.max(device.interval ?? 5, 1);
	const expiresInSeconds = device.expires_in ?? 300;
	callbacks.onDeviceCode({
		userCode: device.user_code,
		verificationUri: device.verification_uri_complete ?? device.verification_uri,
		intervalSeconds,
		expiresInSeconds,
	});

	const deadline = Date.now() + expiresInSeconds * 1000;
	let intervalMs = intervalSeconds * 1000;
	const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));

	while (Date.now() < deadline) {
		await sleep(intervalMs);
		const poll = await fetch(TOKEN_URL, {
			method: "POST",
			headers: FORM_HEADERS,
			body: new URLSearchParams({
				grant_type: DEVICE_CODE_GRANT,
				client_id: CLIENT_ID,
				device_code: device.device_code,
			}).toString(),
		});
		if (poll.ok) return toCredentials((await poll.json()) as TokenResponse);

		const body = (await poll.json().catch(() => ({}))) as {
			error?: string;
			error_description?: string;
		};
		if (body.error === "authorization_pending") continue;
		if (body.error === "slow_down") {
			intervalMs += 5000;
			continue;
		}
		if (body.error === "access_denied" || body.error === "authorization_denied") {
			throw new Error("xAI device authorization was denied");
		}
		if (body.error === "expired_token") {
			throw new Error("xAI device code expired — run /login supergrok again");
		}
		const detail = truncateDetail(body.error_description ?? body.error ?? "");
		throw new Error(
			`xAI device token exchange failed (${poll.status})${detail ? `: ${detail}` : ""}`,
		);
	}
	throw new Error("xAI device authorization timed out");
}

// -----------------------------------------------------------------------------
// Public API
// -----------------------------------------------------------------------------

export async function login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
	const method = await callbacks.onSelect({
		message: "How do you want to log in to SuperGrok / xAI?",
		options: [
			{ id: "browser", label: "Browser (this machine) — opens accounts.x.ai" },
			{ id: "device", label: "Device code (headless / SSH / container)" },
		],
	});
	if (!method) throw new Error("Login cancelled");
	return method === "device" ? loginDevice(callbacks) : loginBrowser(callbacks);
}

export async function refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
	if (!credentials.refresh) {
		throw new Error("No refresh token stored — run /login supergrok again");
	}
	const res = await fetch(TOKEN_URL, {
		method: "POST",
		headers: FORM_HEADERS,
		body: new URLSearchParams({
			grant_type: "refresh_token",
			refresh_token: credentials.refresh,
			client_id: CLIENT_ID,
		}).toString(),
	});
	if (!res.ok) {
		const detail = truncateDetail(await res.text().catch(() => ""));
		if (res.status === 403) throw new Error(tierError(detail));
		throw new Error(`xAI token refresh failed (${res.status})${detail ? `: ${detail}` : ""}`);
	}
	return toCredentials((await res.json()) as TokenResponse, credentials.refresh);
}