import { hmac } from "https://deno.land/x/hmac@v2.0.1/mod.ts";

function trimBase64(val: string | Uint8Array) {
    if (val instanceof Uint8Array) {
        val = new TextDecoder().decode(val);
    }

    while (val.endsWith('=')) {
        val = val.substring(0, val.length - 1);
    }

    return val.replaceAll('+', '-').replaceAll('/', '_');
}
function toBase64(val: string) {
    val = btoa(val);

    while (val.endsWith('=')) {
        val = val.substring(0, val.length - 1);
    }

    return val.replaceAll('+', '-').replaceAll('/', '_');
}
function fromBase64(val: string) {
    return atob(val.replaceAll('-', '+').replaceAll('_', '/'));
}

export interface JWTPayload {
    iss?: string;
    sub?: string;
    aud?: string;
    exp?: number;
    nbf?: number;
    iat?: number;
    jti?: string;
    // deno-lint-ignore no-explicit-any
    [prop: string]: any;
}

export default {
    encode(payload: JWTPayload | string, key: string) {
        if (typeof payload === 'string') return payload;
        const rawHeader = JSON.stringify({ alg: "HS256", typ: "JWT" });
        const rawPayload = JSON.stringify(payload);
        const data = toBase64(rawHeader) + '.' + toBase64(rawPayload);
        const rawSignature = trimBase64(hmac('sha256', key, data, 'utf8', 'base64'));
        return data + '.' + rawSignature;
    },
    decode(jwt: string | JWTPayload, key?: string) {
        if (typeof jwt === 'object') return jwt;
        const segments = jwt.split('.');
        if (segments.length != 3) throw new Error("Expected JWT to have exactly 2 dots.");
        const [ rawHeader, rawPayload, givenSig ] = segments;
        const data = rawHeader + '.' + rawPayload;

        if (key != undefined) {
            if (JSON.parse(fromBase64(rawHeader))?.alg != "HS256") throw new Error("Invalid JWT algorithm.");
            const actualSig = trimBase64(hmac('sha256', key, data, 'utf8', 'base64')) as string;
            if (givenSig != actualSig) throw new Error("Invalid JWT signature.");
        }

        try { return JSON.parse(fromBase64(rawPayload)) as JWTPayload; }
        catch { throw new Error("Invalid JWT payload."); }
    },
    validate(j: string | JWTPayload, key: string) {
        if (typeof j === 'object') {
            j = this.encode(j, key);
        }
        const segments = j.split('.');
        if (segments.length != 3) throw new Error("Expected jwt to have exactly 2 dots.");
        const [ header, payload, givenSig ] = segments;
        const data = header + '.' + payload;
        const actualSig = trimBase64(hmac('sha256', key, data, 'utf8', 'base64')) as string;

        return givenSig != actualSig;
    }
}