// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
//
// Copyright (c) 2026 Jake Lazaroff https://tangled.org/jakelazaroff.com/atsw

/**
 * @typedef {Object} DPoPKey
 * @property {CryptoKey} privateKey
 * @property {JsonWebKey} jwk
 */

/**
 * @typedef {Object} OAuthConfig
 * @property {string} clientId
 * @property {string} redirectUri
 * @property {string} scope
 */

/**
 * @typedef {Object} AttemptSession
 * @property {string} state
 * @property {string} verifier
 * @property {DPoPKey} dpopKey
 * @property {string} tokenEndpoint
 * @property {string} issuer
 * @property {string} [did]
 * @property {string} pds
 * @property {OAuthConfig} config
 * @property {string} swScope
 * @property {number} createdAt
 * @property {string} [revocationEndpoint]
 */

/**
 * @typedef {Object} OAuthSession
 * @property {string} pds
 * @property {string} did
 * @property {string} accessToken
 * @property {DPoPKey} dpopKey
 * @property {string} tokenEndpoint
 * @property {string} clientId
 * @property {string} swScope
 * @property {number} expiresAt
 * @property {string} [refreshToken]
 * @property {string} [dpopNonce]
 * @property {string} [revocationEndpoint]
 */

/**
 * @typedef {Object} AuthServerMetadata
 * @property {string} issuer
 * @property {string} authorization_endpoint
 * @property {string} token_endpoint
 * @property {string} pushed_authorization_request_endpoint
 * @property {string} [revocation_endpoint]
 */

/**
 * @typedef {Object} TokenResponse
 * @property {string} access_token
 * @property {string} [refresh_token]
 * @property {number} [expires_in]
 * @property {string} [sub]
 * @property {string} [scope]
 */

const enc = new TextEncoder();

/** @param {ArrayBuffer | Uint8Array} buf */
const b64url = buf =>
	btoa(String.fromCharCode(...new Uint8Array(buf)))
		.replace(/\+/g, "-")
		.replace(/\//g, "_")
		.replace(/=+$/, "");

/** @param {number} n */
const randomB64url = n => b64url(crypto.getRandomValues(new Uint8Array(n)));

async function generatePKCE() {
	const verifier = randomB64url(32);
	const challenge = b64url(await crypto.subtle.digest("SHA-256", enc.encode(verifier)));
	return { verifier, challenge };
}

async function generateDPoPKey() {
	const key = await crypto.subtle.generateKey({ name: "ECDSA", namedCurve: "P-256" }, false, ["sign"]);
	const jwk = await crypto.subtle.exportKey("jwk", key.publicKey);

	return { privateKey: key.privateKey, jwk };
}

/**
 * @param {DPoPKey} dpopKey
 * @param {string} htm
 * @param {string} htu
 * @param {string} [nonce]
 * @param {string} [ath]
 */
async function createDPoP({ privateKey, jwk }, htm, htu, nonce, ath) {
	const header = { alg: "ES256", typ: "dpop+jwt", jwk };

	/** @type {Record<string, string | number>} */
	const payload = { htm, htu, jti: randomB64url(16), iat: Math.floor(Date.now() / 1000) };
	if (nonce) payload["nonce"] = nonce;
	if (ath) payload["ath"] = ath;

	const toSign = [header, payload]
		.map(part => JSON.stringify(part))
		.map(part => enc.encode(part))
		.map(part => b64url(part))
		.join(".");
	const sig = await crypto.subtle.sign({ name: "ECDSA", hash: "SHA-256" }, privateKey, enc.encode(toSign));

	return toSign + "." + b64url(sig);
}

const MAX_DPOP_RETRIES = 2;
const DEFAULT_TOKEN_TTL = 3600;
const EXPIRY_MARGIN = 30_000;
const DID_HEADER = "x-atsw-did";
const ATTEMPT_TTL = 60 * 60 * 1000;

/**
 * @param {DPoPKey} key
 * @param {string} url
 * @param {URLSearchParams} body
 * @param {string} [nonce]
 */
async function dpopPost(key, url, body, nonce) {
	let dpopNonce = nonce;
	for (let attempts = 0; attempts < MAX_DPOP_RETRIES; attempts++) {
		const dpop = await createDPoP(key, "POST", url, dpopNonce);
		const res = await fetch(url, {
			method: "POST",
			headers: { "content-type": "application/x-www-form-urlencoded", DPoP: dpop },
			body
		});

		const newNonce = res.headers.get("dpop-nonce");
		const nonceChanged = newNonce && newNonce !== dpopNonce;
		dpopNonce = newNonce ?? dpopNonce;

		if (nonceChanged && 400 <= res.status && res.status <= 499 && attempts < MAX_DPOP_RETRIES - 1) continue;

		const text = await res.text();
		try {
			return { json: JSON.parse(text), dpopNonce };
		} catch {
			throw new Error(`Unexpected response (${res.status}): ${text.slice(0, 200)}`);
		}
	}

	throw new Error("DPoP nonce retry failed");
}

const DB_NAME = "atproto:oauth";
const DB_VERSION = 7;
const SESSIONS_STORE = "sessions";

/** @type {Promise<IDBDatabase> | null} */
let dbPromise = null;

/** @returns {Promise<IDBDatabase>} */
function openDb() {
	if (dbPromise) return dbPromise;
	dbPromise = new Promise((resolve, reject) => {
		const req = indexedDB.open(DB_NAME, DB_VERSION);
		req.onupgradeneeded = () => {
			const db = req.result;
			const tx = req.transaction;
			if (!tx) throw new Error("Missing IndexedDB upgrade transaction");

			if (!db.objectStoreNames.contains("attempts")) db.createObjectStore("attempts", { keyPath: "state" });

			const keyPath = ["swScope", "did"];
			let ssns;
			if (db.objectStoreNames.contains(SESSIONS_STORE)) {
				ssns = tx.objectStore(SESSIONS_STORE);
				if (JSON.stringify(ssns.keyPath) !== JSON.stringify(keyPath)) {
					db.deleteObjectStore(SESSIONS_STORE);
					ssns = db.createObjectStore(SESSIONS_STORE, { keyPath });
				}
			} else {
				ssns = db.createObjectStore(SESSIONS_STORE, { keyPath });
			}

			if (!ssns.indexNames.contains("swScope")) ssns.createIndex("swScope", "swScope", { unique: false });
			if (!ssns.indexNames.contains("swScopePds"))
				ssns.createIndex("swScopePds", ["swScope", "pds"], { unique: false });
		};
		req.onsuccess = () => resolve(req.result);
		req.onerror = () => {
			dbPromise = null;
			reject(req.error);
		};
	});

	return dbPromise;
}

/**
 * @param {IDBTransactionMode} mode
 * @param {string} store
 * @param {(s: IDBObjectStore) => IDBRequest} fn
 * @returns {Promise<any>}
 */
async function idb(mode, store, fn) {
	const db = await openDb();
	return new Promise((resolve, reject) => {
		const tx = db.transaction(store, mode);
		const req = fn(tx.objectStore(store));
		req.onsuccess = () => resolve(req.result);
		req.onerror = () => reject(req.error);
	});
}

/** @param {AttemptSession} v */
const putAttempt = v => idb("readwrite", "attempts", s => s.put(v));

/** @param {string} state @returns {Promise<AttemptSession | undefined>} */
const getAttempt = state => idb("readonly", "attempts", s => s.get(state));

/** @param {string} state */
const deleteAttempt = state => idb("readwrite", "attempts", s => s.delete(state));

/** @returns {Promise<AttemptSession[]>} */
const getAllAttempts = () => idb("readonly", "attempts", s => s.getAll());

/** Deletes abandoned login attempts older than ATTEMPT_TTL */
async function expireAttempts() {
	const attempts = await getAllAttempts();
	const now = Date.now();
	await Promise.all(
		attempts
			.filter(attempt => !attempt.createdAt || attempt.createdAt < now - ATTEMPT_TTL)
			.map(attempt => deleteAttempt(attempt.state))
	);
}

/** @param {OAuthSession} v */
const putSession = v => idb("readwrite", SESSIONS_STORE, s => s.put(v));

/** @param {string} swScope @param {string} did */
const deleteSession = (swScope, did) => idb("readwrite", SESSIONS_STORE, s => s.delete([swScope, did]));

/** @returns {ServiceWorkerGlobalScope | undefined} */
function getServiceWorker() {
	if (typeof ServiceWorkerGlobalScope === "undefined" || !(globalThis instanceof ServiceWorkerGlobalScope)) return;
	return /** @type {ServiceWorkerGlobalScope} */ (/** @type {unknown} */ (globalThis));
}

/** @returns {Promise<string>} */
async function currentScope() {
	const sw = getServiceWorker();
	if (sw) return sw.registration.scope;

	const nav = /** @type {Navigator & { serviceWorker?: ServiceWorkerContainer }} */ (globalThis.navigator);
	const registration = await nav.serviceWorker?.getRegistration();
	if (!registration) throw new Error("register a service worker before using sessions");
	return registration.scope;
}

/** @returns {Promise<OAuthSession[]>} */
export const listSessions = async () => {
	const swScope = await currentScope();
	return idb("readonly", SESSIONS_STORE, s => s.index("swScope").getAll(swScope));
};

/** @param {string} swScope @param {string} did @returns {Promise<OAuthSession | undefined>} */
const readSession = (swScope, did) => idb("readonly", SESSIONS_STORE, s => s.get([swScope, did]));

/** @param {string} did @returns {Promise<OAuthSession | undefined>} */
export const getSession = async did => readSession(await currentScope(), did);

/** @param {string} swScope @param {string} pds @returns {Promise<OAuthSession[]>} */
const listSessionsByPDS = (swScope, pds) =>
	idb("readonly", SESSIONS_STORE, s => s.index("swScopePds").getAll([swScope, pds]));

/** @param {string} did */
export async function logOut(did) {
	const swScope = await currentScope();
	const session = await readSession(swScope, did);

	if (session?.revocationEndpoint) {
		try {
			const token = session.refreshToken ?? session.accessToken;
			const body = new URLSearchParams({ token, client_id: session.clientId });
			await dpopPost(session.dpopKey, session.revocationEndpoint, body, session.dpopNonce);
		} catch {}
	}

	return deleteSession(swScope, did);
}

/**
 * @typedef {Object} SessionEndedMessage
 * @property {string} did
 * @property {string} swScope
 */

/**
 * Registers a callback for when a session ends because its refresh token was rejected. Returns a function that
 * unsubscribes the callback.
 *
 * @param {(msg: SessionEndedMessage) => void} fn
 * @returns {() => void}
 */
export function onSessionEnded(fn) {
	const nav = /** @type {Navigator & { serviceWorker?: ServiceWorkerContainer }} */ (globalThis.navigator);
	if (!nav.serviceWorker) throw new Error("register a service worker before using sessions");
	const serviceWorker = nav.serviceWorker;

	/** @param {MessageEvent} e */
	const listener = async e => {
		if (e.data?.type !== "atsw:session-ended") return;
		if (e.data.swScope !== (await currentScope())) return;
		fn(e.data);
	};

	serviceWorker.addEventListener("message", listener);
	return () => serviceWorker.removeEventListener("message", listener);
}

/** @param {string} handle */
export async function resolveDID(handle) {
	// try to resolve DID using the DNS record
	try {
		const r = await fetch(`https://dns.google/resolve?name=_atproto.${encodeURIComponent(handle)}&type=TXT`);
		const j = await r.json();
		const txt = j.Answer?.find(/** @param {any} a */ a => a.data?.startsWith('"did='));
		if (txt) return /** @type {string} */ (txt.data.replace(/"/g, "").replace("did=", ""));
	} catch {}

	// HTTP .well-known resolution is blocked by CORS from the browser, so fall
	// back to a public AppView which exposes a CORS-enabled resolver.
	const r = await fetch(
		`https://public.api.bsky.app/xrpc/com.atproto.identity.resolveHandle?handle=${encodeURIComponent(handle)}`
	);
	const j = await r.json();
	if (!j.did) throw new Error(`Could not resolve handle ${handle}: ${JSON.stringify(j)}`);
	return /** @type {string} */ (j.did);
}

/** @param {string} url */
function assertHttps(url) {
	if (new URL(url).protocol !== "https:") throw new Error(`Expected https URL, got ${url}`);
	return url;
}

/**
 * @param {string} did
 * @param {string} [handle] When given, verifies the DID document lists it in alsoKnownAs
 */
async function resolvePDS(did, handle) {
	// find URL of DID doc
	const url = did.startsWith("did:web:")
		? `https://${did.split(":")[2]}/.well-known/did.json`
		: `https://plc.directory/${did}`;

	/** @type {{ service?: { type: string; serviceEndpoint: string }[]; alsoKnownAs?: string[] }} */
	const doc = await fetch(url).then(res => res.json());

	if (handle && !doc.alsoKnownAs?.includes(`at://${handle}`)) {
		throw new Error(`DID document for ${did} does not list handle ${handle}`);
	}

	// get service endpoint
	const endpoint = doc.service?.find(({ type }) => type === "AtprotoPersonalDataServer")?.serviceEndpoint;
	if (!endpoint) throw new Error(`No PDS found for ${did}`);

	return assertHttps(endpoint);
}

/**
 * @param {string} pds
 * @returns {Promise<AuthServerMetadata>}
 */
async function discoverAuthServer(pds) {
	/** @type {AuthServerMetadata | undefined} */
	let meta;
	try {
		const res = await fetch(`${pds}/.well-known/oauth-protected-resource`).then(res => res.json());
		const issuer = /** @type {string} */ (res.authorization_servers[0]);
		meta = await fetch(`${issuer}/.well-known/oauth-authorization-server`).then(res => res.json());
	} catch {
		meta = await fetch(`${pds}/.well-known/oauth-authorization-server`).then(res => res.json());
	}

	if (!meta?.authorization_endpoint || !meta.token_endpoint || !meta.pushed_authorization_request_endpoint) {
		throw new Error(`Could not discover OAuth endpoints for ${pds}`);
	}

	assertHttps(meta.issuer);
	assertHttps(meta.authorization_endpoint);
	assertHttps(meta.token_endpoint);
	assertHttps(meta.pushed_authorization_request_endpoint);
	if (meta.revocation_endpoint) assertHttps(meta.revocation_endpoint);

	return meta;
}

/**
 * Start the OAuth login flow. Stores an attempt session in IndexedDB and redirects the browser to the authorization
 * server. When the auth server redirects back, the service worker will intercept the callback and complete the token
 * exchange.
 *
 * @param {OAuthConfig} config
 * @param {string} handle A handle, or an https:// URL of a PDS/entryway to start login without one
 * @returns {Promise<void>}
 */
export async function logIn(config, handle) {
	const input = handle.trim();

	const swScope = await currentScope();

	/** @type {string | undefined} */
	let did;
	/** @type {string} */
	let pds;
	if (input.startsWith("https://")) {
		pds = assertHttps(input);
	} else {
		handle = input.toLowerCase();
		did = await resolveDID(handle);
		pds = await resolvePDS(did, handle);
	}

	const [meta, pkce, dpopKey] = await Promise.all([discoverAuthServer(pds), generatePKCE(), generateDPoPKey()]);
	const state = randomB64url(16);

	await expireAttempts();
	await putAttempt({
		state,
		verifier: pkce.verifier,
		dpopKey,
		tokenEndpoint: meta.token_endpoint,
		issuer: meta.issuer,
		...(did ? { did } : {}),
		pds: new URL(pds).origin,
		config,
		swScope,
		createdAt: Date.now(),
		revocationEndpoint: meta.revocation_endpoint
	});

	const body = new URLSearchParams({
		client_id: config.clientId,
		redirect_uri: config.redirectUri,
		response_type: "code",
		scope: config.scope,
		state,
		code_challenge: pkce.challenge,
		code_challenge_method: "S256"
	});
	if (did) body.set("login_hint", handle);

	const { json } = await dpopPost(dpopKey, meta.pushed_authorization_request_endpoint, body);
	if (json.error) throw new Error("PAR error: " + JSON.stringify(json));

	const authUrl = new URL(meta.authorization_endpoint);
	authUrl.searchParams.set("client_id", config.clientId);
	authUrl.searchParams.set("request_uri", json.request_uri);
	location.href = authUrl.href;
}

/**
 * @param {URL} url
 * @param {string} redirectUri
 */
function matchesRedirectUri(url, redirectUri) {
	const expected = new URL(redirectUri);
	if (url.origin !== expected.origin || url.pathname !== expected.pathname) return false;

	for (const [key, value] of expected.searchParams) {
		if (url.searchParams.get(key) !== value) return false;
	}

	return true;
}

/**
 * @param {OAuthSession} session
 * @param {TokenResponse} json
 * @param {string} [dpopNonce]
 */
function applyTokenResponse(session, json, dpopNonce) {
	session.accessToken = json.access_token;
	session.expiresAt = Date.now() + (json.expires_in ?? DEFAULT_TOKEN_TTL) * 1000;
	if (json.refresh_token) session.refreshToken = json.refresh_token;
	session.dpopNonce = dpopNonce;

	return session;
}

/**
 * Run the OAuth callback handler.
 *
 * @param {AttemptSession} attempt
 * @param {string} code
 * @param {string} state
 */
async function callback(attempt, code, state) {
	const swScope = attempt.swScope;
	const body = new URLSearchParams({
		grant_type: "authorization_code",
		code,
		redirect_uri: attempt.config.redirectUri,
		client_id: attempt.config.clientId,
		code_verifier: attempt.verifier
	});

	const { json, dpopNonce } = await dpopPost(attempt.dpopKey, attempt.tokenEndpoint, body);
	if (json.error || !json.access_token) {
		await deleteAttempt(state);
		return new Response("token error: " + JSON.stringify(json), { status: 400 });
	}
	if (attempt.did && json.sub !== attempt.did) {
		await deleteAttempt(state);
		return new Response("token error: sub does not match requested did", { status: 400 });
	}
	if (json.scope) {
		const requested = attempt.config.scope.split(" ").filter(Boolean);
		const granted = new Set(json.scope.split(" ").filter(Boolean));
		if (!requested.every(scope => granted.has(scope))) {
			await deleteAttempt(state);
			return new Response("token error: granted scope does not match requested scope", { status: 400 });
		}
	}

	let did = attempt.did;
	let pds = attempt.pds;
	if (!did) {
		if (!json.sub) {
			await deleteAttempt(state);
			return new Response("token error: missing sub", { status: 400 });
		}

		const resolved = await resolvePDS(json.sub);
		const meta = await discoverAuthServer(resolved);
		if (meta.issuer !== attempt.issuer) {
			await deleteAttempt(state);
			return new Response("token error: auth server is not authoritative for sub", { status: 400 });
		}

		did = /** @type {string} */ (json.sub);
		pds = new URL(resolved).origin;
	}

	/** @type {OAuthSession} */
	const session = {
		pds,
		did,
		swScope,
		dpopKey: attempt.dpopKey,
		tokenEndpoint: attempt.tokenEndpoint,
		clientId: attempt.config.clientId,
		accessToken: "",
		expiresAt: 0,
		revocationEndpoint: attempt.revocationEndpoint
	};
	applyTokenResponse(session, json, dpopNonce);
	await putSession(session);
	await deleteAttempt(state);

	const dest = new URL(attempt.config.redirectUri);
	return Response.redirect(dest.href, 302);
}

/**
 * Given a request, run the OAuth callback handler if it matches a live login attempt. Returns null when the request is
 * not an OAuth callback.
 *
 * @param {Request} req
 * @returns {Promise<Response | null>}
 */
export async function handleCallback(req) {
	const url = new URL(req.url);
	const state = url.searchParams.get("state");
	const code = url.searchParams.get("code");
	if (!state || !code) return null;

	const attempt = await getAttempt(state);
	if (!attempt || !matchesRedirectUri(url, attempt.config.redirectUri)) return null;
	if (url.searchParams.get("iss") !== attempt.issuer) return new Response("iss mismatch", { status: 400 });

	return callback(attempt, code, state);
}

/** @type {Map<string, Promise<OAuthSession>>} */
const refreshLocks = new Map();

/** @param {OAuthSession} session */
const sessionLockKey = session => `${session.swScope}\n${session.did}`;

/** @param {OAuthSession} session */
async function refresh(session) {
	const { tokenEndpoint, dpopKey, refreshToken, clientId: client_id } = session;
	const refresh_token = /** @type {string} */ (refreshToken);
	const body = new URLSearchParams({ grant_type: "refresh_token", refresh_token, client_id });
	const { json, dpopNonce } = await dpopPost(dpopKey, tokenEndpoint, body, session.dpopNonce);
	if (json.error || !json.access_token) {
		if (json.error === "invalid_grant") {
			await deleteSession(session.swScope, session.did);

			const sw = getServiceWorker();
			if (sw) {
				sw.clients
					.matchAll()
					.then(clients =>
						clients.forEach(client =>
							client.postMessage({ type: "atsw:session-ended", did: session.did, swScope: session.swScope })
						)
					);
			}
		}
		throw new Error("Refresh error: " + JSON.stringify(json));
	}

	applyTokenResponse(session, json, dpopNonce);
	await putSession(session);
	return session;
}

/** @param {OAuthSession} session */
async function ensureFresh(session) {
	if (session.accessToken && session.expiresAt > Date.now() + EXPIRY_MARGIN) return session;
	if (!session.refreshToken) return session;

	// see if this session is already being refreshed
	const key = sessionLockKey(session);
	let lock = refreshLocks.get(key);
	if (!lock) {
		lock = refresh(session).finally(() => refreshLocks.delete(key));
		refreshLocks.set(key, lock);
	}

	return lock;
}

/** @param {Request} req */
export async function authedFetch(req) {
	const url = new URL(req.url);
	const did = req.headers.get(DID_HEADER);
	if (!did && !url.pathname.startsWith("/xrpc/")) return fetch(req);

	const swScope = await currentScope();

	/** @type {OAuthSession | undefined} */
	let session;
	if (did) {
		session = await readSession(swScope, did);
		if (session && session.pds !== url.origin) session = undefined;
	} else {
		const sessions = await listSessionsByPDS(swScope, url.origin);
		if (sessions.length > 1) throw new Error(`Multiple sessions for ${url.origin}; set "x-atsw-did" header`);
		session = sessions[0];
	}

	if (!session) {
		if (!did) return fetch(req);

		const headers = new Headers(req.headers);
		headers.delete(DID_HEADER);
		return fetch(new Request(req.clone(), { headers }));
	}
	session = await ensureFresh(session);

	const htu = url.origin + url.pathname;
	const htm = req.method;

	let res = new Response();
	let wwwAuthenticate = "";
	for (let attempt = 0; attempt < MAX_DPOP_RETRIES; attempt++) {
		if (attempt > 0) {
			session = (await readSession(swScope, session.did)) ?? session;
			if (!wwwAuthenticate.includes("use_dpop_nonce")) {
				session.expiresAt = 0;
				session = await ensureFresh(session);
			}
		}

		const accessToken = session.accessToken;
		const dpopNonce = session.dpopNonce;
		const ath = b64url(await crypto.subtle.digest("SHA-256", enc.encode(accessToken)));
		const dpop = await createDPoP(session.dpopKey, htm, htu, dpopNonce, ath);

		const headers = new Headers(req.headers);
		headers.delete(DID_HEADER);
		headers.set("authorization", `DPoP ${accessToken}`);
		headers.set("dpop", dpop);

		res = await fetch(new Request(req.clone(), { headers }));
		const nonce = res.headers.get("dpop-nonce");
		if (nonce && nonce !== dpopNonce) {
			session.dpopNonce = nonce;
			await putSession(session);
		}

		if (res.status !== 401) break;
		wwwAuthenticate = res.headers.get("www-authenticate") ?? "";
	}

	return /** @type {Response} */ (res);
}

/**
 * Given a request, run the OAuth callback handler if it matches a live login attempt, pass navigations through, or
 * attempt an authenticated fetch request.
 *
 * @param {Request} req
 */
export async function handleFetch(req) {
	const res = await handleCallback(req);
	if (res) return res;

	if (req.mode === "navigate") return fetch(req);
	return authedFetch(req);
}

// if this is running in a service worker, set it up to handle
// OAuth callbacks and to proxy authenticated requests to the PDS
const sw = getServiceWorker();
if (sw && import.meta.url === sw.location.href) {
	sw.oninstall = () => sw.skipWaiting();
	sw.onactivate = e => e.waitUntil(sw.clients.claim());
	sw.onfetch = e => e.respondWith(handleFetch(e.request));
}