import { createLocalJWKSet, createRemoteJWKSet, decodeProtectedHeader, jwtVerify } from 'jose';
import type { JWTPayload, JWTVerifyOptions, RemoteJWKSetOptions } from 'jose';
import { isEmpty } from 'lodash';

import { UnauthorizedError } from '../errors/index.js';

/**
 * Verifies JWTs using a JWKS endpoint to resolve signing keys on demand.
 */
export class JwtVerifier {
  private readonly options: JwtVerifierOptions;
  // Memoize the remote JWKS loader; without this, we'd recreate it on each verification and jose's own cache would never persist.
  private remoteJwkSet?: ReturnType<typeof createRemoteJWKSet>;
  constructor(options: JwtVerifierOptions) {
    this.options = options;
  }

  async verify<T extends JWTPayload>(token: string, options?: JWTVerifyOptions): Promise<T> {
    let header: ReturnType<typeof decodeProtectedHeader>;
    try {
      header = decodeProtectedHeader(token);
    } catch {
      // Preserve legacy error message relied upon by tests
      throw new Error('Invalid token specified');
    }

    if (isEmpty(header.kid)) {
      throw new UnauthorizedError('Invalid JWT as kid header is missing.');
    }

    // Fail closed: require an explicit algorithm allowlist and reject a
    // disallowed `alg` before resolving keys, so a forged `alg` header can't
    // select an unintended verification path or trigger a JWKS fetch.
    if (!options?.algorithms || options.algorithms.length === 0) {
      throw new UnauthorizedError('JWT could not be verified as no signing algorithm allowlist was configured.');
    }
    if (!header.alg || !options.algorithms.includes(header.alg)) {
      throw new UnauthorizedError('JWT could not be verified as its algorithm is not allowed.');
    }

    const jwkSet = await this.resolveJwkSet();

    try {
      const result = await jwtVerify<T>(token, jwkSet, options);
      return result.payload;
    } catch (err) {
      throw new UnauthorizedError(`JWT verification failed: ${(err as Error).message}`);
    }
  }

  private async resolveJwkSet(): Promise<ReturnType<typeof createLocalJWKSet> | ReturnType<typeof createRemoteJWKSet>> {
    // Prefer interceptor-provided keys for explicit testability and control
    if (this.options.getKeysInterceptor) {
      const keys = await this.options.getKeysInterceptor();
      if (isEmpty(keys)) {
        throw new UnauthorizedError('JWT could not be verified as no corresponding public key was found.');
      }
      return createLocalJWKSet({ keys });
    }

    // Allow custom fetcher to provide JWKS out-of-band
    if (this.options.fetcher) {
      const jwks = await this.options.fetcher(this.options.jwksUri);
      const keys = jwks.keys;
      if (!Array.isArray(keys) || isEmpty(keys)) {
        throw new UnauthorizedError('JWT could not be verified as no corresponding public key was found.');
      }
      return createLocalJWKSet({ keys });
    }

    // Default to remote JWKS. jose v5 supports headers/timeout; pass when provided.
    const url = new URL(this.options.jwksUri);
    const opts: RemoteJWKSetOptions = {};
    if (!isEmpty(this.options.requestHeaders)) {
      opts.headers = this.options.requestHeaders;
    }
    if (!isEmpty(this.options.timeoutMs)) {
      opts.timeoutDuration = this.options.timeoutMs;
    }
    if (!this.remoteJwkSet) {
      this.remoteJwkSet = createRemoteJWKSet(url, opts);
    }
    return this.remoteJwkSet;
  }
}

/**
 * Options that configure how the verifier fetches and caches JWKS signing keys.
 */
export interface JwtVerifierOptions {
  jwksUri: string;
  rateLimit?: boolean;
  cache?: boolean;
  cacheMaxEntries?: number;
  cacheMaxAge?: number;
  jwksRequestsPerMinute?: number;
  proxy?: string;
  requestHeaders?: Headers;
  timeoutMs?: number;
  fetcher?(jwksUri: string): Promise<{ keys: unknown }>;
  getKeysInterceptor?(): Promise<JwtKey[]>;
}

type Headers = Record<string, string>;

export interface JwtKey {
  kty: string;
  kid: string;
  alg: string;
  [key: string]: unknown;
}