import {
  BadRequestException,
  Injectable,
  ServiceUnavailableException,
} from '@nestjs/common';
import { createHash, randomBytes } from 'crypto';

type JoseModule = typeof import('jose');
type RemoteJwkSet = ReturnType<JoseModule['createRemoteJWKSet']>;

export interface OidcProfile {
  subject: string;
  email: string;
  name?: string;
}

interface OidcDiscovery {
  authorization_endpoint: string;
  token_endpoint: string;
  jwks_uri: string;
  issuer: string;
}

interface PendingOidcState {
  codeVerifier: string;
  nonce: string;
  expiresAt: number;
}

interface TokenResponse {
  id_token?: string;
  error?: string;
  error_description?: string;
}

@Injectable()
export class OidcService {
  private readonly pendingStates = new Map<string, PendingOidcState>();
  private discovery?: OidcDiscovery;
  private jose?: Promise<JoseModule>;
  private jwks?: RemoteJwkSet;

  async createAuthorizationUrl(): Promise<string> {
    const config = this.getConfig();
    const discovery = await this.getDiscovery(config.discoveryUrl);
    const state = this.createOpaqueToken();
    const nonce = this.createOpaqueToken();
    const codeVerifier = this.createOpaqueToken();
    const codeChallenge = this.codeChallenge(codeVerifier);
    const authorizationUrl = new URL(discovery.authorization_endpoint);

    this.pendingStates.set(state, {
      codeVerifier,
      nonce,
      expiresAt: Date.now() + 10 * 60 * 1000,
    });
    this.deleteExpiredStates();

    authorizationUrl.searchParams.set('response_type', 'code');
    authorizationUrl.searchParams.set('client_id', config.clientId);
    authorizationUrl.searchParams.set('redirect_uri', config.callbackUrl);
    authorizationUrl.searchParams.set('scope', 'openid email profile');
    authorizationUrl.searchParams.set('state', state);
    authorizationUrl.searchParams.set('nonce', nonce);
    authorizationUrl.searchParams.set('code_challenge', codeChallenge);
    authorizationUrl.searchParams.set('code_challenge_method', 'S256');

    return authorizationUrl.toString();
  }

  async exchangeCallback(code?: string, state?: string): Promise<OidcProfile> {
    if (!code || !state) {
      throw new BadRequestException('OIDC code and state are required.');
    }

    const pendingState = this.pendingStates.get(state);
    this.pendingStates.delete(state);

    if (!pendingState || pendingState.expiresAt <= Date.now()) {
      throw new BadRequestException('OIDC state is invalid or expired.');
    }

    const config = this.getConfig();
    const discovery = await this.getDiscovery(config.discoveryUrl);
    const tokenResponse = await this.requestTokens(
      discovery,
      config,
      code,
      pendingState.codeVerifier,
    );

    if (!tokenResponse.id_token) {
      throw new ServiceUnavailableException(
        tokenResponse.error_description ??
          tokenResponse.error ??
          'OIDC token response did not include an ID token.',
      );
    }

    const [{ jwtVerify }, jwks] = await Promise.all([
      this.getJose(),
      this.getJwks(discovery.jwks_uri),
    ]);
    const { payload } = await jwtVerify(tokenResponse.id_token, jwks, {
      issuer: discovery.issuer,
      audience: config.clientId,
    });

    if (payload.nonce !== pendingState.nonce) {
      throw new BadRequestException('OIDC nonce is invalid.');
    }

    if (!payload.sub || typeof payload.sub !== 'string') {
      throw new BadRequestException('OIDC subject is missing.');
    }

    const email = typeof payload.email === 'string' ? payload.email : undefined;

    if (!email) {
      throw new BadRequestException('OIDC email claim is missing.');
    }

    return {
      subject: payload.sub,
      email,
      name: typeof payload.name === 'string' ? payload.name : undefined,
    };
  }

  private async requestTokens(
    discovery: OidcDiscovery,
    config: ReturnType<OidcService['getConfig']>,
    code: string,
    codeVerifier: string,
  ): Promise<TokenResponse> {
    const body = new URLSearchParams({
      grant_type: 'authorization_code',
      code,
      redirect_uri: config.callbackUrl,
      client_id: config.clientId,
      code_verifier: codeVerifier,
    });

    if (config.clientSecret) {
      body.set('client_secret', config.clientSecret);
    }

    const response = await fetch(discovery.token_endpoint, {
      method: 'POST',
      headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
      body,
    });
    const payload = (await response.json().catch(() => ({}))) as TokenResponse;

    if (!response.ok) {
      throw new ServiceUnavailableException(
        payload.error_description ??
          payload.error ??
          'OIDC token request failed.',
      );
    }

    return payload;
  }

  private async getDiscovery(discoveryUrl: string): Promise<OidcDiscovery> {
    if (this.discovery) {
      return this.discovery;
    }

    const response = await fetch(discoveryUrl);

    if (!response.ok) {
      throw new ServiceUnavailableException('OIDC discovery failed.');
    }

    this.discovery = (await response.json()) as OidcDiscovery;
    return this.discovery;
  }

  private async getJwks(jwksUri: string): Promise<RemoteJwkSet> {
    const { createRemoteJWKSet } = await this.getJose();

    this.jwks ??= createRemoteJWKSet(new URL(jwksUri));
    return this.jwks;
  }

  private getJose(): Promise<JoseModule> {
    this.jose ??= import('jose');
    return this.jose;
  }

  private getConfig() {
    const issuerUrl = process.env.OIDC_ISSUER_URL;
    const explicitDiscoveryUrl = process.env.OIDC_DISCOVERY_URL;
    const clientId = process.env.OIDC_CLIENT_ID;
    const callbackUrl = process.env.OIDC_CALLBACK_URL;

    if (!issuerUrl || !clientId || !callbackUrl) {
      throw new ServiceUnavailableException(
        'OIDC configuration is incomplete.',
      );
    }

    return {
      issuerUrl,
      discoveryUrl:
        explicitDiscoveryUrl ??
        `${issuerUrl.replace(/\/$/, '')}/.well-known/openid-configuration`,
      clientId,
      callbackUrl,
      clientSecret: process.env.OIDC_CLIENT_SECRET,
    };
  }

  private createOpaqueToken(): string {
    return randomBytes(32).toString('base64url');
  }

  private codeChallenge(codeVerifier: string): string {
    return createHash('sha256').update(codeVerifier).digest('base64url');
  }

  private deleteExpiredStates(): void {
    const now = Date.now();

    for (const [state, pendingState] of this.pendingStates.entries()) {
      if (pendingState.expiresAt <= now) {
        this.pendingStates.delete(state);
      }
    }
  }
}