import {
  CanActivate,
  ExecutionContext,
  Injectable,
  Logger,
  UnauthorizedException,
} from '@nestjs/common'
import { ConfigService } from '@nestjs/config'
import { createRemoteJWKSet, jwtVerify } from 'jose'
import type { AuthentikJwtPayload } from './jwt-payload.interface.js'

@Injectable()
export class JwtAuthGuard implements CanActivate {
  private readonly logger = new Logger(JwtAuthGuard.name)
  private jwks: ReturnType<typeof createRemoteJWKSet> | null = null
  private readonly issuer: string
  private readonly audience: string
  private readonly jwksUri: string

  constructor(config: ConfigService) {
    this.issuer = config.getOrThrow<string>('AUTHENTIK_ISSUER')
    this.audience = config.getOrThrow<string>('AUTHENTIK_AUDIENCE')
    this.jwksUri = config.getOrThrow<string>('AUTHENTIK_JWKS_URI')
  }

  // Lazy init so misconfigured env doesn't crash the module bootstrap; the first
  // request will surface a useful 401 instead of a startup failure.
  private getJwks() {
    if (!this.jwks) {
      this.jwks = createRemoteJWKSet(new URL(this.jwksUri), {
        cacheMaxAge: 10 * 60 * 1000,
        cooldownDuration: 30 * 1000,
      })
    }
    return this.jwks
  }

  async canActivate(context: ExecutionContext): Promise<boolean> {
    const req = context.switchToHttp().getRequest<{
      headers: Record<string, string | string[]>
      user?: AuthentikJwtPayload
    }>()
    const authHeader = req.headers['authorization']
    const headerValue = Array.isArray(authHeader) ? authHeader[0] : authHeader

    if (!headerValue?.startsWith('Bearer ')) {
      throw new UnauthorizedException('Missing or malformed Authorization header')
    }

    const token = headerValue.slice('Bearer '.length).trim()

    try {
      const { payload } = await jwtVerify(token, this.getJwks(), {
        issuer: this.issuer,
        audience: this.audience,
      })
      req.user = payload as unknown as AuthentikJwtPayload
      return true
    } catch (err) {
      this.logger.warn(`JWT verification failed: ${(err as Error).message}`)
      throw new UnauthorizedException('Invalid access token')
    }
  }
}