// Verifying atproto service auth on a Node server (Hono or any other; rust/host/src/serviceauth.rs
// is the Rust host's): an `Authorization: Bearer <JWT>` a caller's PDS minted with
// com.atproto.server.getServiceAuth, signed (ES256K or ES256, low S) by the caller's #atproto key.

import { createHash, createPublicKey, verify } from "node:crypto";

import { unbase58 } from "../encoding/index.js";

/** Longest a token may still be valid for, and the clock skew allowed on its iat. */
export const MAX_TTL_S = 3600;
export const SKEW_S = 60;
/** Tokens a replay cache remembers at once, and how long and how many DID documents are kept. */
const REPLAY_MAX = 10_000;
const DOC_TTL_MS = 300_000;
const DOC_MAX = 1000;

/**
 * The tokens verifyServiceAuth accepted, each until its exp: a token is good for one request (a
 * replay of it is refused). Keyed by iss + jti, or the token's sha256 when it has no jti.
 * Past `max` live tokens new ones are refused ("busy") rather than a live one dropped,
 * which would let it be replayed. ponytail: one cap for every issuer; cap per iss if one floods it.
 */
export function replayCache(max = REPLAY_MAX) {
	const seen = new Map(); // key -> exp
	return {
		seen,
		admit(key, exp, now) {
			if ((seen.get(key) ?? 0) > now) throw new Error("the token was already used");
			if (seen.size >= max) for (const [k, e] of seen) if (e <= now) seen.delete(k);
			if (seen.size >= max) throw new Error("busy: too many tokens in use; try again later");
			seen.set(key, exp);
		},
	};
}

/**
 * `resolve` with its documents kept for ttlMs, at most `max` of them (expired ones go first, then
 * the oldest). Failures are not kept.
 * ponytail: a key rotated within ttlMs fails the issuer's new tokens until the entry expires.
 */
export function cachedResolve(resolve, { ttlMs = DOC_TTL_MS, max = DOC_MAX, clock = Date.now } = {}) {
	const docs = new Map(); // did -> [at, doc], oldest first
	return async (did) => {
		const hit = docs.get(did);
		if (hit && clock() - hit[0] < ttlMs) return hit[1];
		const doc = await resolve(did);
		docs.delete(did);
		if (docs.size >= max) for (const [k, [at]] of docs) if (clock() - at >= ttlMs) docs.delete(k);
		if (docs.size >= max) docs.delete(docs.keys().next().value);
		docs.set(did, [clock(), doc]);
		return doc;
	};
}

// multicodec prefix (secp256k1-pub 0xe7 0x01, p256-pub 0x80 0x24) -> [alg, curve OID, group order]
const CURVES = {
	"231,1": ["ES256K", [0x2b, 0x81, 0x04, 0x00, 0x0a], 0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141n],
	"128,36": ["ES256", [0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07], 0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551n],
};

// The #atproto key of a DID document: (multicodec prefix, compressed key).
function atprotoKey(doc) {
	const multibase = doc?.verificationMethod?.find?.((m) => typeof m?.id === "string" && m.id.endsWith("#atproto"))?.publicKeyMultibase;
	if (typeof multibase !== "string") throw new Error("the issuer's DID document has no #atproto key");
	if (!multibase.startsWith("z")) throw new Error("publicKeyMultibase is not base58btc");
	const bytes = unbase58(multibase.slice(1));
	if (bytes.length < 3) throw new Error("publicKeyMultibase is too short");
	return [`${bytes[0]},${bytes[1]}`, bytes.subarray(2)];
}

const der = (tag, bytes) => [tag, bytes.length, ...bytes];

function checkSignature(alg, prefix, key, message, signature) {
	const curve = CURVES[prefix];
	if (!curve || curve[0] !== alg) throw new Error(`alg ${alg} does not match the issuer's key type`);
	const [, oid, n] = curve;
	if (signature.length !== 64) throw new Error("bad signature: not compact r||s");
	if (BigInt(`0x${Buffer.from(signature.subarray(32)).toString("hex")}`) > n / 2n) throw new Error("bad signature: high S");
	// SubjectPublicKeyInfo { { ecPublicKey, curve }, BIT STRING compressed point }
	const spki = der(0x30, [...der(0x30, [...der(0x06, [0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02, 0x01]), ...der(0x06, oid)]), ...der(0x03, [0, ...key])]);
	let publicKey;
	try {
		publicKey = createPublicKey({ key: Buffer.from(spki), format: "der", type: "spki" });
	} catch (e) {
		throw new Error(`bad signature: ${e.message}`);
	}
	if (!verify("sha256", message, { key: publicKey, dsaEncoding: "ieee-p1363" }, signature)) throw new Error("bad signature: signature error");
}

/**
 * did:plc via plc.directory, did:web via /.well-known/did.json (uncached; verifyServiceAuth's
 * default resolve is this through cachedResolve).
 * ponytail: unguarded fetch of the issuer's document; hardening is the open SSRF item (one hardened DID resolver).
 */
export async function resolveDid(did, fetch = globalThis.fetch) {
	const url = did.startsWith("did:web:") ? `https://${did.slice(8)}/.well-known/did.json` : did.startsWith("did:plc:") ? `https://plc.directory/${did}` : null;
	if (!url) throw new Error(`cannot resolve ${did}`);
	const response = await fetch(url);
	if (!response.ok) throw new Error(`resolving ${did}: ${response.status}`);
	return response.json();
}

const defaultResolve = cachedResolve((did) => resolveDid(did));
const defaultReplay = replayCache();

/**
 * The issuer DID of a valid `Authorization: Bearer <service-auth JWT>` for method `lxm`: aud one of
 * `audiences` (or that DID with a #fragment), iat no later than now + SKEW_S, exp in the future but
 * at most MAX_TTL_S away, iss in `issuers` (empty: anyone), signed by the issuer's #atproto key, and
 * not seen by `replay` before (a caller mints one token per request). Throws an Error with the
 * reason otherwise (answer 401). Defaults: DID documents cached (cachedResolve over resolveDid), one
 * replay cache for the process.
 * @param {{ audiences: string[], issuers?: string[] }} auth
 * @param {string | undefined | null} authorization  the Authorization header
 * @param {string} lxm
 * @param {{ now?: number, resolve?: (did: string) => Promise<object>, replay?: ReturnType<typeof replayCache> }} [options]  now: unix seconds
 */
export async function verifyServiceAuth({ audiences, issuers = [] }, authorization, lxm, { now = Math.floor(Date.now() / 1000), resolve = defaultResolve, replay = defaultReplay } = {}) {
	if (typeof authorization !== "string" || !authorization.startsWith("Bearer ")) throw new Error("Bearer service-auth token required");
	const parts = authorization.slice(7).trim().split(".");
	if (parts.length !== 3) throw new Error("the token is not a JWT");
	const [h, p, s] = parts;
	const part = (x) => {
		if (!/^[A-Za-z0-9_-]*$/.test(x)) throw new Error("the token is not base64url");
		try {
			return JSON.parse(Buffer.from(x, "base64url").toString());
		} catch {
			throw new Error("the token is not JSON");
		}
	};
	const [header, claims] = [part(h), part(p)];
	const claim = (k) => (typeof claims?.[k] === "string" ? claims[k] : "");
	const iss = claim("iss").split("#")[0];
	if (claim("lxm") !== lxm) throw new Error(`the token is for ${JSON.stringify(claim("lxm"))}, not ${lxm}`);
	const aud = claim("aud");
	if (!audiences.some((a) => a === aud || a === aud.split("#")[0])) throw new Error(`the token's aud ${JSON.stringify(aud)} is not this service`);
	const exp = Number.isSafeInteger(claims.exp) && claims.exp > 0 ? claims.exp : 0;
	if (exp <= now) throw new Error("the token has expired");
	if (exp > now + MAX_TTL_S) throw new Error(`the token's exp is more than ${MAX_TTL_S} s away`);
	if (!Number.isSafeInteger(claims.iat) || claims.iat < 0 || claims.iat > now + SKEW_S) throw new Error("the token's iat is missing or in the future");
	if (issuers.length && !issuers.includes(iss)) throw new Error(`issuer ${iss} is not allowed`);
	const [prefix, key] = atprotoKey(await resolve(iss));
	// Canonical only: Node decodes spare trailing bits away, so "AQID-g" and "AQID-h" would be one
	// signature under two replay keys.
	if (!/^[A-Za-z0-9_-]*$/.test(s) || Buffer.from(s, "base64url").toString("base64url") !== s) throw new Error("the signature is not base64url");
	checkSignature(typeof header?.alg === "string" ? header.alg : "", prefix, key, Buffer.from(`${h}.${p}`), new Uint8Array(Buffer.from(s, "base64url")));
	const jti = claim("jti");
	replay.admit(jti ? `${iss} ${jti}` : `sha256 ${createHash("sha256").update(`${h}.${p}.${s}`).digest("hex")}`, exp, now);
	return iss;
}