import { createRemoteJWKSet, jwtVerify, type JWTPayload, type JWTVerifyGetKey } from "jose";

type KeyInput = Parameters<typeof jwtVerify>[1];

export interface VerifierConfig {
  issuer: string;
  /** Token is accepted if its `aud` includes at least one of these. */
  audiences: string[];
  /** Token must carry this Zitadel project role (granted in Zitadel, asserted into the access token). */
  requiredRole: string;
  /** Test seam: resolve the verification key directly instead of fetching JWKS. */
  keyResolver?: (token: string) => Promise<KeyInput>;
}

/**
 * Zitadel asserts project roles into JWT access tokens as
 * `urn:zitadel:iam:org:project:roles` (and a `…:project:<id>:roles` variant),
 * each an object keyed by role key.
 */
export function rolesFromClaims(payload: JWTPayload): Set<string> {
  const roles = new Set<string>();
  for (const [claim, value] of Object.entries(payload)) {
    if (!/^urn:zitadel:iam:org:project:(\d+:)?roles$/.test(claim)) continue;
    if (value && typeof value === "object" && !Array.isArray(value)) {
      for (const role of Object.keys(value)) roles.add(role);
    }
  }
  return roles;
}

/**
 * Build a Zitadel access-token verifier. Verifies signature + issuer + expiry via JWKS,
 * then enforces the audience and the required project role. Throws on any failure.
 */
export function makeVerifier(cfg: VerifierConfig) {
  const jwks: JWTVerifyGetKey | null = cfg.keyResolver
    ? null
    : createRemoteJWKSet(new URL(`${cfg.issuer}/oauth/v2/keys`));

  return async function verify(token: string): Promise<JWTPayload> {
    const key: KeyInput = cfg.keyResolver ? await cfg.keyResolver(token) : (jwks as JWTVerifyGetKey);
    const { payload } = await jwtVerify(token, key, { issuer: cfg.issuer });

    const aud = Array.isArray(payload.aud) ? payload.aud : payload.aud ? [payload.aud] : [];
    if (!cfg.audiences.some((a) => aud.includes(a))) {
      throw new Error("token audience not accepted");
    }
    if (!rolesFromClaims(payload).has(cfg.requiredRole)) {
      throw new Error(`required role "${cfg.requiredRole}" not granted`);
    }
    return payload;
  };
}

function splitCsv(v: unknown): string[] {
  return String(v ?? "")
    .split(",")
    .map((s) => s.trim())
    .filter(Boolean);
}

let _cached: ReturnType<typeof makeVerifier> | null = null;

/** Process-wide verifier built from environment (read at runtime, not baked at build). */
export function getVerifier() {
  if (!_cached) {
    _cached = makeVerifier({
      issuer: process.env.ZITADEL_ISSUER || "https://auth.kuns.dev",
      audiences: splitCsv(process.env.ZITADEL_AUDIENCE),
      requiredRole: process.env.REQUIRED_ROLE || "user",
    });
  }
  return _cached;
}