import { env } from "cloudflare:workers";
import { z } from "zod";
import { emailAllowed, isBlockedOutbound } from "./naming";
import { isForgeHostname, sitesDomain } from "./config";
import type { SiteMeta } from "./config";

type OidcConfig = NonNullable<SiteMeta["oidc"]>;

const DISCOVERY_CACHE_SECONDS = 3600;
const CLOCK_SKEW_SECONDS = 60;

const discoverySchema = z.object({
  issuer: z.string(),
  authorization_endpoint: z.string(),
  token_endpoint: z.string(),
  jwks_uri: z.string(),
});
type Discovery = z.infer<typeof discoverySchema>;

const jwkSchema = z.object({ kid: z.string().optional(), kty: z.string(), alg: z.string().optional() }).passthrough();
const claimsSchema = z.object({
  iss: z.string(),
  aud: z.union([z.string(), z.array(z.string())]),
  exp: z.number(),
  nonce: z.string().optional(),
  email: z.string().optional(),
  email_verified: z.union([z.boolean(), z.string()]).optional(),
});

/** Only public HTTPS endpoints; a site owner must not be able to point the dispatcher at internal hosts. */
export function assertPublicHttps(raw: string): URL {
  const url = new URL(raw);
  const blocked = [new URL(env.PUBLIC_URL).hostname, sitesDomain()].filter((host): host is string => Boolean(host));
  if (url.protocol !== "https:" || isBlockedOutbound(url, blocked) || isForgeHostname(url.hostname)) {
    throw new Error(`OIDC endpoint ${url.hostname} must be a public https URL`);
  }
  return url;
}

async function getJson(url: string, init?: RequestInit): Promise<unknown> {
  const response = await fetch(assertPublicHttps(url), { ...init, redirect: "error" });
  if (!response.ok) throw new Error(`OIDC request to ${new URL(url).hostname} failed (${response.status})`);
  return response.json();
}

async function discover(issuer: string): Promise<Discovery> {
  const key = `sites:oidc:${issuer}`;
  const cached = await env.KV.get<Discovery>(key, { type: "json", cacheTtl: DISCOVERY_CACHE_SECONDS });
  if (cached) return cached;
  const document = discoverySchema.parse(
    await getJson(`${issuer.replace(/\/$/, "")}/.well-known/openid-configuration`),
  );
  if (document.issuer.replace(/\/$/, "") !== issuer.replace(/\/$/, "")) throw new Error("OIDC issuer mismatch");
  await env.KV.put(key, JSON.stringify(document), { expirationTtl: DISCOVERY_CACHE_SECONDS });
  return document;
}

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

function fromBase64url(value: string): Uint8Array<ArrayBuffer> {
  const binary = atob(value.replace(/-/g, "+").replace(/_/g, "/"));
  return Uint8Array.from(binary, (char) => char.charCodeAt(0));
}

function random(): string {
  return base64url(crypto.getRandomValues(new Uint8Array(32)));
}

export async function startOidc(
  config: OidcConfig,
  redirectUri: string,
): Promise<{ location: string; state: string; nonce: string; verifier: string }> {
  const discovery = await discover(config.issuer);
  const state = random();
  const nonce = random();
  const verifier = random();
  const challenge = base64url(
    new Uint8Array(await crypto.subtle.digest("SHA-256", new TextEncoder().encode(verifier))),
  );
  const url = assertPublicHttps(discovery.authorization_endpoint);
  url.searchParams.set("response_type", "code");
  url.searchParams.set("client_id", config.clientId);
  url.searchParams.set("redirect_uri", redirectUri);
  url.searchParams.set("scope", "openid email");
  url.searchParams.set("state", state);
  url.searchParams.set("nonce", nonce);
  url.searchParams.set("code_challenge", challenge);
  url.searchParams.set("code_challenge_method", "S256");
  return { location: url.toString(), state, nonce, verifier };
}

async function importJwk(jwk: z.infer<typeof jwkSchema>, alg: string): Promise<CryptoKey> {
  if (alg === "RS256" && jwk.kty === "RSA") {
    return crypto.subtle.importKey("jwk", jwk as JsonWebKey, { name: "RSASSA-PKCS1-v1_5", hash: "SHA-256" }, false, [
      "verify",
    ]);
  }
  if (alg === "ES256" && jwk.kty === "EC") {
    return crypto.subtle.importKey("jwk", jwk as JsonWebKey, { name: "ECDSA", namedCurve: "P-256" }, false, ["verify"]);
  }
  throw new Error(`Unsupported id_token algorithm ${alg}`);
}

/** Verifies an id_token's signature against the issuer's JWKS and returns its claims. */
async function verifyIdToken(token: string, discovery: Discovery): Promise<z.infer<typeof claimsSchema>> {
  const [headerPart, payloadPart, signaturePart] = token.split(".");
  if (!headerPart || !payloadPart || !signaturePart) throw new Error("Malformed id_token");
  const header = z
    .object({ alg: z.string(), kid: z.string().optional() })
    .parse(JSON.parse(new TextDecoder().decode(fromBase64url(headerPart))));
  const jwks = z.object({ keys: z.array(jwkSchema) }).parse(await getJson(discovery.jwks_uri));
  const jwk = jwks.keys.find((key) => (header.kid ? key.kid === header.kid : true));
  if (!jwk) throw new Error("Signing key not found");
  if (jwk.alg && jwk.alg !== header.alg) throw new Error("id_token algorithm does not match its key");
  const key = await importJwk(jwk, header.alg);
  const algorithm = header.alg === "ES256" ? { name: "ECDSA", hash: "SHA-256" } : { name: "RSASSA-PKCS1-v1_5" };
  const valid = await crypto.subtle.verify(
    algorithm,
    key,
    fromBase64url(signaturePart),
    new TextEncoder().encode(`${headerPart}.${payloadPart}`),
  );
  if (!valid) throw new Error("Invalid id_token signature");
  return claimsSchema.parse(JSON.parse(new TextDecoder().decode(fromBase64url(payloadPart))));
}

/** Exchanges the code (with PKCE) and returns the visitor's verified, allow-listed email, or null. */
export async function finishOidc(
  config: OidcConfig,
  clientSecret: string,
  input: { code: string; nonce: string; verifier: string; redirectUri: string },
): Promise<string | null> {
  const discovery = await discover(config.issuer);
  const tokens = z.object({ id_token: z.string() }).parse(
    await getJson(discovery.token_endpoint, {
      method: "POST",
      headers: { "content-type": "application/x-www-form-urlencoded", accept: "application/json" },
      body: new URLSearchParams({
        grant_type: "authorization_code",
        code: input.code,
        redirect_uri: input.redirectUri,
        client_id: config.clientId,
        client_secret: clientSecret,
        code_verifier: input.verifier,
      }),
    }),
  );
  const claims = await verifyIdToken(tokens.id_token, discovery);
  const audience = Array.isArray(claims.aud) ? claims.aud : [claims.aud];
  const now = Math.floor(Date.now() / 1000);
  if (claims.iss !== discovery.issuer || !audience.includes(config.clientId)) return null;
  if (claims.exp + CLOCK_SKEW_SECONDS < now || claims.nonce !== input.nonce) return null;
  const verified = claims.email_verified === true || claims.email_verified === "true";
  if (!claims.email || !verified) return null;
  return emailAllowed(claims.email, config.allowed) ? claims.email.toLowerCase() : null;
}