import {
  importPKCS8,
  importSPKI,
  JWTHeaderParameters,
  JWTPayload,
  jwtVerify,
  SignJWT,
  errors,
} from 'jose'
import { jwt as jwtConfig } from '~/config'
import { type UserInfo } from '~/modules/auth/schema'

export const JWTInvalid = errors.JWTInvalid
export const JWTExpired = errors.JWTExpired

type Alg = 'HS256' | 'RS256' | 'ES256'

interface SignOptions {
  issuer?: string
  audience?: string | string[]
  expiresIn?: string | number
  kid?: string
  alg?: Alg // optional override
}

interface VerifyOptions {
  issuer?: string | string[]
  audience?: string | string[]
  clockTolerance?: string | number
  alg?: Alg // optional override
}

export type JwtClaims = JWTPayload & {
  user: UserInfo
}

let cachedAlg: Alg | null = null
let cachedSignerKey: CryptoKey | null = null
let cachedVerifierKey: CryptoKey | null = null

function getAlg(override?: Alg): Alg {
  const envAlg = jwtConfig.algorithm.toUpperCase()
  const alg = (override || envAlg) as Alg
  if (!['HS256', 'RS256', 'ES256'].includes(alg)) {
    throw new Error(`Unsupported JWT_ALG: ${alg}`)
  }
  return alg
}

async function importHmacKey(secret: string): Promise<CryptoKey> {
  const raw = new TextEncoder().encode(secret)
  return crypto.subtle.importKey(
    'raw',
    raw,
    { name: 'HMAC', hash: 'SHA-256' },
    false,
    ['sign', 'verify'],
  )
}

async function loadKeys(alg: Alg): Promise<{ signKey: CryptoKey, verifyKey: CryptoKey }> {
  if (cachedAlg === alg && cachedSignerKey && cachedVerifierKey) {
    return { signKey: cachedSignerKey, verifyKey: cachedVerifierKey }
  }

  if (alg === 'HS256') {
    const key = await importHmacKey(jwtConfig.key.secret)
    cachedAlg = alg
    cachedSignerKey = key
    cachedVerifierKey = key
    return { signKey: key, verifyKey: key }
  }

  const priv = await importPKCS8(jwtConfig.key.private, alg)
  const pub = await importSPKI(jwtConfig.key.public, alg)
  cachedAlg = alg
  cachedSignerKey = priv
  cachedVerifierKey = pub
  return { signKey: priv, verifyKey: pub }
}

export async function sign(
  claims: JwtClaims,
  opts: SignOptions = {},
): Promise<string> {
  const alg = getAlg(opts.alg)
  const { signKey } = await loadKeys(alg)

  const header: JWTHeaderParameters = { alg, ...(opts.kid ? { kid: opts.kid } : {}) }

  const builder = new SignJWT(claims).setProtectedHeader(header).setIssuedAt()

  if (opts.issuer) builder.setIssuer(opts.issuer)
  if (opts.audience) builder.setAudience(opts.audience)
  builder.setExpirationTime(opts.expiresIn ?? jwtConfig.expires.access)

  return builder.sign(signKey)
}

export async function verify<T extends JWTPayload = JwtClaims>(
  token: string,
  opts: VerifyOptions = {},
): Promise<T> {
  const alg = getAlg(opts.alg)
  const { verifyKey } = await loadKeys(alg)

  const { payload } = await jwtVerify(token, verifyKey, {
    algorithms: [alg],
    issuer: opts.issuer,
    audience: opts.audience,
    clockTolerance: opts.clockTolerance ?? '5s',
  })

  return payload as T
}