import {Result} from 'rust'
import * as jose from 'jose'
import logger from 'log'

const log = logger('crypto:jwt')

export type JWTalgorithm = "HS256" | "HS512" | "ES256" | "ES512" | "EdDSA"
export type Key = jose.KeyLike | Uint8Array
/** KeyPair for asymmetric algorithms, [PrivKey, PubKey] */
export type KeyPair = [Key, Key]

export class JWTcontext {
    public constructor(
        private readonly key: Key | KeyPair,
        private readonly alg: JWTalgorithm,
    ) {}

    public static async gen_key(alg: JWTalgorithm, extractable: boolean = false) : Promise<Key | KeyPair> {
        log.debug(`Generate new ${alg} key`)
        log.trace('Key extractable :', extractable ? 'yes' : 'no')

        switch (alg) {
            case "HS256":
            case "HS512":
                return jose.generateSecret(alg, { extractable })
            case "EdDSA":
            case "ES256":
            case "ES512":
            {
                const key = await jose.generateKeyPair(alg)
                return [key.privateKey, key.publicKey]
            }
        }
    }

    public async sign<T>(message: T, set_issued: boolean = false, exp?: number | string | Date, audience?: string | string[], issuer?: string): Promise<string> {
        log.debug('sign JWT')
        log.trace('Config :', {
            set_issued,
            exp,
            issuer,
        })

        let jwt = new jose.SignJWT({message}).setProtectedHeader({ alg: this.alg })

        if (set_issued) jwt = jwt.setIssuedAt()
        if (issuer !== undefined) jwt = jwt.setIssuer(issuer)
        if (audience !== undefined) jwt = jwt.setAudience(audience)
        if (exp !== undefined) jwt = jwt.setExpirationTime(exp)

        const key = this.get_key(true)
        return await jwt.sign(key)
    }

    public async verify<T>(jwt: string, audience?: string | string[], issuer?: string | string[]): Promise<Result<T>> {
        log.debug('Verify JWT')
        log.trace('Issuers :', issuer)

        const key = this.get_key(false)

        try {
            let payload = await jose.jwtVerify(jwt, key)
            return Result.ok(payload.payload.message as T)
        } catch(e) {
            log.warn('JWT verification failed')
            log.debug(`Error : ${e}`)
        }

        return Result.error([])
    }

    private get_key(sign: boolean) : Key {
        switch (this.alg) {
            case "HS256":
            case "HS512":
                return this.key as Key
            case "ES256":
            case "ES512":
            case "EdDSA":
                return (this.key as KeyPair)[sign ? 0 : 1]
        }
    }
}