// =============================================================================
// Std::Crypto::Jwt — JSON Web Token (RFC 7519) encode/decode
// Supports HS256, HS384, HS512, RS256, RS384, RS512, ES256, ES384, EdDSA
// =============================================================================
module Std::Crypto::Jwt {

    import Std::Mem::{Alloc, Free};
    import Std::String::{String_Len, String_Eq, String_StartsWith, String_Concat};
    import Std::Crypto::Base64::{Base64URL_Encode, Base64URL_Decode};
    import Std::Crypto::Hash::{Hash_Sha256Raw, Hash_Sha384Raw, Hash_Sha512Raw};
    import Std::Crypto::Hmac::{Hmac_Sha256Raw, Hmac_Sha384Raw, Hmac_Sha512Raw};
    import Std::Crypto::Rsa::{Rsa_SignSha256, Rsa_SignSha384, Rsa_SignSha512,
        Rsa_VerifySha256, Rsa_VerifySha384, Rsa_VerifySha512};
    import Std::Crypto::Ecdsa::{Ecdsa_SignP256, Ecdsa_SignP384, Ecdsa_VerifyP256, Ecdsa_VerifyP384};
    import Std::Crypto::Ed25519::{Ed25519_Sign, Ed25519_Verify};

    extern func bux_base64_encode(data: String, len: int) -> String;
    extern func bux_str_split_count(s: String, delim: String) -> uint;
    extern func bux_str_split_part(s: String, delim: String, index: uint) -> String;

    // --- JWT Algorithm enum ---
    enum JwtAlg {
        HS256,
        HS384,
        HS512,
        RS256,
        RS384,
        RS512,
        ES256,
        ES384,
        EdDSA,
    }

    // --- Header ---

    // Jwt_MakeHeader: build the JWT header JSON string for the given algorithm
    func Jwt_MakeHeader(alg: JwtAlg) -> String {
        if alg.tag == JwtAlg_HS256 { return "{\"alg\":\"HS256\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_HS384 { return "{\"alg\":\"HS384\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_HS512 { return "{\"alg\":\"HS512\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_RS256 { return "{\"alg\":\"RS256\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_RS384 { return "{\"alg\":\"RS384\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_RS512 { return "{\"alg\":\"RS512\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_ES256 { return "{\"alg\":\"ES256\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_ES384 { return "{\"alg\":\"ES384\",\"typ\":\"JWT\"}"; }
        if alg.tag == JwtAlg_EdDSA { return "{\"alg\":\"EdDSA\",\"typ\":\"JWT\"}"; }
        return "{\"alg\":\"none\",\"typ\":\"JWT\"}";
    }

    // --- Signing ---

    // Sign the JWT signing input with the given algorithm
    func Jwt_Sign(alg: JwtAlg, signingInput: String, key: String) -> String {
        // --- HMAC algorithms ---
        if alg.tag == JwtAlg_HS256 {
            let buf: *void = Alloc(32);
            Hmac_Sha256Raw(key, signingInput, buf);
            let result: String = bux_base64_encode(buf as String, 32);
            Free(buf);
            return result;
        }
        if alg.tag == JwtAlg_HS384 {
            let buf: *void = Alloc(48);
            Hmac_Sha384Raw(key, signingInput, buf);
            let result: String = bux_base64_encode(buf as String, 48);
            Free(buf);
            return result;
        }
        if alg.tag == JwtAlg_HS512 {
            let buf: *void = Alloc(64);
            Hmac_Sha512Raw(key, signingInput, buf);
            let result: String = bux_base64_encode(buf as String, 64);
            Free(buf);
            return result;
        }

        // --- RSA algorithms (key is PEM private key) ---
        if alg.tag == JwtAlg_RS256 {
            let raw: String = Rsa_SignSha256(key, signingInput);
            return bux_base64_encode(raw, String_Len(raw) as int);
        }
        if alg.tag == JwtAlg_RS384 {
            let raw: String = Rsa_SignSha384(key, signingInput);
            return bux_base64_encode(raw, String_Len(raw) as int);
        }
        if alg.tag == JwtAlg_RS512 {
            let raw: String = Rsa_SignSha512(key, signingInput);
            return bux_base64_encode(raw, String_Len(raw) as int);
        }

        // --- ECDSA algorithms (key is PEM private key) ---
        if alg.tag == JwtAlg_ES256 {
            let raw: String = Ecdsa_SignP256(key, signingInput);
            return bux_base64_encode(raw, String_Len(raw) as int);
        }
        if alg.tag == JwtAlg_ES384 {
            let raw: String = Ecdsa_SignP384(key, signingInput);
            return bux_base64_encode(raw, String_Len(raw) as int);
        }

        // --- EdDSA (key is 32-byte raw private key) ---
        if alg.tag == JwtAlg_EdDSA {
            return Ed25519_Sign(key, signingInput);
        }

        return "";
    }

    // --- Verify ---

    // Verify a JWT signature
    func Jwt_Verify(alg: JwtAlg, signingInput: String, signatureB64: String, key: String) -> bool {
        // --- HMAC algorithms ---
        if alg.tag == JwtAlg_HS256 {
            let expectBuf: *void = Alloc(32);
            Hmac_Sha256Raw(key, signingInput, expectBuf);
            let expectB64: String = bux_base64_encode(expectBuf as String, 32);
            Free(expectBuf);
            return String_Eq(expectB64, signatureB64);
        }
        if alg.tag == JwtAlg_HS384 {
            let expectBuf: *void = Alloc(48);
            Hmac_Sha384Raw(key, signingInput, expectBuf);
            let expectB64: String = bux_base64_encode(expectBuf as String, 48);
            Free(expectBuf);
            return String_Eq(expectB64, signatureB64);
        }
        if alg.tag == JwtAlg_HS512 {
            let expectBuf: *void = Alloc(64);
            Hmac_Sha512Raw(key, signingInput, expectBuf);
            let expectB64: String = bux_base64_encode(expectBuf as String, 64);
            Free(expectBuf);
            return String_Eq(expectB64, signatureB64);
        }

        // --- RSA algorithms ---
        if alg.tag == JwtAlg_RS256 { return Rsa_VerifySha256(key, signingInput, signatureB64); }
        if alg.tag == JwtAlg_RS384 { return Rsa_VerifySha384(key, signingInput, signatureB64); }
        if alg.tag == JwtAlg_RS512 { return Rsa_VerifySha512(key, signingInput, signatureB64); }

        // --- ECDSA algorithms ---
        if alg.tag == JwtAlg_ES256 { return Ecdsa_VerifyP256(key, signingInput, signatureB64); }
        if alg.tag == JwtAlg_ES384 { return Ecdsa_VerifyP384(key, signingInput, signatureB64); }

        // --- EdDSA ---
        if alg.tag == JwtAlg_EdDSA { return Ed25519_Verify(key, signatureB64, signingInput); }

        return false;
    }

    // --- Encode ---

    // Jwt_Encode: create a signed JWT
    //   headerJson — JSON header string (use Jwt_MakeHeader or custom)
    //   payloadJson — JSON payload/claims string
    //   alg — signing algorithm
    //   key — signing key (HMAC secret, RSA PEM, ECDSA PEM, or Ed25519 raw privkey)
    // Returns the complete "header.payload.signature" JWT string
    func Jwt_Encode(headerJson: String, payloadJson: String, alg: JwtAlg, key: String) -> String {
        let headerB64: String = Base64URL_Encode(headerJson);
        let payloadB64: String = Base64URL_Encode(payloadJson);
        let signingInput: String = String_Concat(headerB64, ".");
        let signingInputFull: String = String_Concat(signingInput, payloadB64);

        let sigB64: String = Jwt_Sign(alg, signingInputFull, key);

        let part1: String = String_Concat(signingInputFull, ".");
        return String_Concat(part1, sigB64);
    }

    // --- Decode ---

    // Jwt_Decode: decode and verify a JWT.
    //   token — the full "header.payload.signature" string
    //   alg — expected algorithm
    //   key — verification key
    //   headerOut — receives decoded header JSON
    //   payloadOut — receives decoded payload JSON
    // Returns true if signature is valid.
    func Jwt_Decode(token: String, alg: JwtAlg, key: String,
    headerOut: *String, payloadOut: *String) -> bool {
        // Split by "."
        let partCount: uint = bux_str_split_count(token, ".");
        if partCount != 3 { return false; }

        let headerB64: String = bux_str_split_part(token, ".", 0);
        let payloadB64: String = bux_str_split_part(token, ".", 1);
        let sigB64: String = bux_str_split_part(token, ".", 2);

        // Build signing input
        let input: String = String_Concat(headerB64, ".");
        let signingInput: String = String_Concat(input, payloadB64);

        // Verify signature
        if !Jwt_Verify(alg, signingInput, sigB64, key) {
            return false;
        }

        // Decode
        headerOut[0] = Base64URL_Decode(headerB64);
        payloadOut[0] = Base64URL_Decode(payloadB64);
        return true;
    }

    // --- Convenience: Encode with standard header ---

    func Jwt_EncodeHS256(payloadJson: String, secret: String) -> String {
        let header: String = Jwt_MakeHeader(JwtAlg { tag: JwtAlg_HS256 });
        return Jwt_Encode(header, payloadJson, JwtAlg { tag: JwtAlg_HS256 }, secret);
    }

    func Jwt_EncodeHS384(payloadJson: String, secret: String) -> String {
        let header: String = Jwt_MakeHeader(JwtAlg { tag: JwtAlg_HS384 });
        return Jwt_Encode(header, payloadJson, JwtAlg { tag: JwtAlg_HS384 }, secret);
    }

    func Jwt_EncodeHS512(payloadJson: String, secret: String) -> String {
        let header: String = Jwt_MakeHeader(JwtAlg { tag: JwtAlg_HS512 });
        return Jwt_Encode(header, payloadJson, JwtAlg { tag: JwtAlg_HS512 }, secret);
    }

    func Jwt_EncodeRS256(payloadJson: String, pemPrivateKey: String) -> String {
        let header: String = Jwt_MakeHeader(JwtAlg { tag: JwtAlg_RS256 });
        return Jwt_Encode(header, payloadJson, JwtAlg { tag: JwtAlg_RS256 }, pemPrivateKey);
    }

    func Jwt_EncodeES256(payloadJson: String, pemPrivateKey: String) -> String {
        let header: String = Jwt_MakeHeader(JwtAlg { tag: JwtAlg_ES256 });
        return Jwt_Encode(header, payloadJson, JwtAlg { tag: JwtAlg_ES256 }, pemPrivateKey);
    }

    func Jwt_EncodeEdDSA(payloadJson: String, rawPrivKey: String) -> String {
        let header: String = Jwt_MakeHeader(JwtAlg { tag: JwtAlg_EdDSA });
        return Jwt_Encode(header, payloadJson, JwtAlg { tag: JwtAlg_EdDSA }, rawPrivKey);
    }
}
