// =============================================================================
// jwt-pitbul — JWT CLI Tool (encode, decode, verify, keygen)
// Rewritten with modern Bux: algebraic enums, methods, string interpolation,
// for-in loops, generic arrays, and structured error handling.
// =============================================================================
module Main {

    import Std::Io::{PrintLine, Print, PrintInt};
    import Std::String::{
        String_Len, String_Eq,
        String_SplitCount, String_SplitPart,
        StringBuilder, StringBuilder_New, StringBuilder_Append,
        StringBuilder_AppendInt, StringBuilder_Build, StringBuilder_Free
    };
    import Std::Crypto::Jwt::{
        JwtAlg,
        Jwt_MakeHeader,
        Jwt_Encode,
        Jwt_Decode,
        Jwt_EncodeHS256, Jwt_EncodeHS384, Jwt_EncodeHS512,
        Jwt_EncodeRS256, Jwt_EncodeES256, Jwt_EncodeEdDSA
    };
    import Std::Crypto::Base64::{Base64URL_Decode, Base64_Encode};
    import Std::Crypto::Ed25519::{Ed25519_Keypair};

    extern func bux_argc() -> int;
    extern func bux_argv(index: int) -> String;
    extern func bux_alloc(size: uint) -> *void;
    extern func bux_read_file(path: String) -> String;
    extern func bux_file_exists(path: String) -> int;
    extern func bux_base64_encode(data: String, len: int) -> String;

    // =============================================================================
    // Constants
    // =============================================================================
    const AppName: String = "jwt-pitbul";
    const Version: String = "0.2.0";

    // =============================================================================
    // Algebraic enums for optionals and results
    // =============================================================================
    enum AlgOption {
        Some(JwtAlg),
        None,
    }

    func AlgOption_MakeSome(value: JwtAlg) -> AlgOption {
        let o: AlgOption = AlgOption { tag: AlgOption_Some };
        o.data.Some_0 = value;
        return o;
    }

    func AlgOption_MakeNone() -> AlgOption {
        return AlgOption { tag: AlgOption_None };
    }

    enum CmdOption {
        Some(Cmd),
        None,
    }

    func CmdOption_MakeSome(value: Cmd) -> CmdOption {
        let o: CmdOption = CmdOption { tag: CmdOption_Some };
        o.data.Some_0 = value;
        return o;
    }

    func CmdOption_MakeNone() -> CmdOption {
        return CmdOption { tag: CmdOption_None };
    }

    enum KeyTypeOption {
        Some(KeyType),
        None,
    }

    func KeyTypeOption_MakeSome(value: KeyType) -> KeyTypeOption {
        let o: KeyTypeOption = KeyTypeOption { tag: KeyTypeOption_Some };
        o.data.Some_0 = value;
        return o;
    }

    func KeyTypeOption_MakeNone() -> KeyTypeOption {
        return KeyTypeOption { tag: KeyTypeOption_None };
    }

    enum KeyResult {
        Ok(String),
        Err(String),
    }

    func KeyResult_MakeOk(value: String) -> KeyResult {
        let r: KeyResult = KeyResult { tag: KeyResult_Ok };
        r.data.Ok_0 = value;
        return r;
    }

    func KeyResult_MakeErr(msg: String) -> KeyResult {
        let r: KeyResult = KeyResult { tag: KeyResult_Err };
        r.data.Err_0 = msg;
        return r;
    }

    enum ExitResult {
        Ok(int),
        Err(String),
    }

    func ExitResult_MakeOk(value: int) -> ExitResult {
        let r: ExitResult = ExitResult { tag: ExitResult_Ok };
        r.data.Ok_0 = value;
        return r;
    }

    func ExitResult_MakeErr(msg: String) -> ExitResult {
        let r: ExitResult = ExitResult { tag: ExitResult_Err };
        r.data.Err_0 = msg;
        return r;
    }

    // =============================================================================
    // Commands
    // =============================================================================
    enum Cmd {
        Sign,
        Verify,
        Decode,
        Keygen,
        Help,
    }

    func ParseCmd(name: String) -> CmdOption {
        if String_Eq(name, "sign")   { return CmdOption_MakeSome(Cmd { tag: Cmd_Sign }); }
        if String_Eq(name, "verify") { return CmdOption_MakeSome(Cmd { tag: Cmd_Verify }); }
        if String_Eq(name, "decode") { return CmdOption_MakeSome(Cmd { tag: Cmd_Decode }); }
        if String_Eq(name, "keygen") { return CmdOption_MakeSome(Cmd { tag: Cmd_Keygen }); }
        if String_Eq(name, "help") || String_Eq(name, "--help") || String_Eq(name, "-h") {
            return CmdOption_MakeSome(Cmd { tag: Cmd_Help });
        }
        return CmdOption_MakeNone();
    }

    // =============================================================================
    // Key type generation
    // =============================================================================
    enum KeyType {
        Rsa,
        Ecdsa,
        Ed25519,
    }

    func ParseKeyType(name: String) -> KeyTypeOption {
        if String_Eq(name, "rsa")     { return KeyTypeOption_MakeSome(KeyType { tag: KeyType_Rsa }); }
        if String_Eq(name, "ecdsa")   { return KeyTypeOption_MakeSome(KeyType { tag: KeyType_Ecdsa }); }
        if String_Eq(name, "ed25519") { return KeyTypeOption_MakeSome(KeyType { tag: KeyType_Ed25519 }); }
        return KeyTypeOption_MakeNone();
    }

    // =============================================================================
    // Extend JwtAlg with parsing / introspection methods
    // =============================================================================
    extend JwtAlg {
        func Name(self: JwtAlg) -> String {
            match self {
                JwtAlg::HS256 => "HS256",
                JwtAlg::HS384 => "HS384",
                JwtAlg::HS512 => "HS512",
                JwtAlg::RS256 => "RS256",
                JwtAlg::RS384 => "RS384",
                JwtAlg::RS512 => "RS512",
                JwtAlg::ES256 => "ES256",
                JwtAlg::ES384 => "ES384",
                JwtAlg::EdDSA => "EdDSA",
                _ => "HS256",
            }
        }

        func IsHmac(self: JwtAlg) -> bool {
            return self.tag == JwtAlg_HS256 || self.tag == JwtAlg_HS384 || self.tag == JwtAlg_HS512;
        }

        func NeedsPemFile(self: JwtAlg) -> bool {
            return self.tag == JwtAlg_RS256 || self.tag == JwtAlg_RS384 || self.tag == JwtAlg_RS512 ||
            self.tag == JwtAlg_ES256 || self.tag == JwtAlg_ES384;
        }

    }

    func ParseAlg(name: String) -> AlgOption {
        if String_Eq(name, "HS256") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_HS256 }); }
        if String_Eq(name, "HS384") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_HS384 }); }
        if String_Eq(name, "HS512") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_HS512 }); }
        if String_Eq(name, "RS256") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_RS256 }); }
        if String_Eq(name, "RS384") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_RS384 }); }
        if String_Eq(name, "RS512") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_RS512 }); }
        if String_Eq(name, "ES256") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_ES256 }); }
        if String_Eq(name, "ES384") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_ES384 }); }
        if String_Eq(name, "EdDSA") { return AlgOption_MakeSome(JwtAlg { tag: JwtAlg_EdDSA }); }
        return AlgOption_MakeNone();
    }

    // =============================================================================
    // Help / Usage
    // =============================================================================
    func PrintUsage() {
        PrintLine("╔══════════════════════════════════════════════════════╗");
        PrintLine(f"║  {AppName} — JWT CLI Tool v{Version}                  ║");
        PrintLine("║  Sign, verify, decode JSON Web Tokens               ║");
        PrintLine("╚══════════════════════════════════════════════════════╝");
        PrintLine("");
        PrintLine("Usage:");
        PrintLine(f"  {AppName} sign <alg> <key> <claims>");
        PrintLine(f"  {AppName} verify <token> <alg> <key>");
        PrintLine(f"  {AppName} decode <token>");
        PrintLine(f"  {AppName} keygen <type>");
        PrintLine("");
        PrintLine("Commands:");
        PrintLine("  sign     Create a signed JWT from claims JSON");
        PrintLine("  verify   Verify a JWT signature and print payload");
        PrintLine("  decode   Decode a JWT without verification");
        PrintLine("  keygen   Generate cryptographic keys");
        PrintLine("");
        PrintLine("Algorithms:");
        PrintLine("  HS256, HS384, HS512     — HMAC (symmetric)");
        PrintLine("  RS256, RS384, RS512     — RSA PKCS#1 v1.5");
        PrintLine("  ES256, ES384            — ECDSA P-256 / P-384");
        PrintLine("  EdDSA                   — Ed25519");
        PrintLine("");
        PrintLine("Key formats:");
        PrintLine("  HMAC:  raw secret string");
        PrintLine("  RSA:   path to PEM private/public key file");
        PrintLine("  ECDSA: path to PEM private/public key file");
        PrintLine("  EdDSA: base64-encoded 32-byte raw key");
        PrintLine("");
        PrintLine("Key generation:");
        PrintLine(f"  {AppName} keygen rsa       # RSA 2048-bit (PEM)");
        PrintLine(f"  {AppName} keygen ecdsa     # ECDSA P-256 (PEM)");
        PrintLine(f"  {AppName} keygen ed25519   # Ed25519 (raw base64)");
        PrintLine("");
        PrintLine("Examples:");
        Print(f"  {AppName} sign HS256 'my-secret' ");
        PrintLine("'{\"sub\":\"123\"}'");
        PrintLine(f"  {AppName} verify eyJh... HS256 'my-secret'");
        PrintLine(f"  {AppName} decode eyJh...");
        PrintLine(f"  {AppName} keygen ed25519");
    }

    // =============================================================================
    // Key resolution — for RSA/ECDSA, read PEM file; for HMAC/EdDSA, pass through
    // =============================================================================
    func ResolveKey(alg: JwtAlg, keyArg: String) -> KeyResult {
        if !alg.NeedsPemFile() {
            return KeyResult_MakeOk(keyArg);
        }

        if bux_file_exists(keyArg) == 0 {
            return KeyResult_MakeErr(f"ERROR: PEM file not found: {keyArg}");
        }

        let pem: String = bux_read_file(keyArg);
        if String_Len(pem) == 0 {
            return KeyResult_MakeErr(f"ERROR: could not read PEM file: {keyArg}");
        }
        return KeyResult_MakeOk(pem);
    }

    // =============================================================================
    // Command: sign
    // =============================================================================
    func CmdSign(algName: String, keyArg: String, claimsJson: String) -> ExitResult {
        let algOpt: AlgOption = ParseAlg(algName);
        if algOpt.tag != AlgOption_Some {
            return ExitResult_MakeErr(f"ERROR: unknown algorithm '{algName}'");
        }
        let alg: JwtAlg = algOpt.data.Some_0;

        let keyRes: KeyResult = ResolveKey(alg, keyArg);
        if keyRes.tag != KeyResult_Ok {
            return ExitResult_MakeErr(keyRes.data.Err_0);
        }
        let key: String = keyRes.data.Ok_0;

        let header: String = Jwt_MakeHeader(alg);
        let token: String = Jwt_Encode(header, claimsJson, alg, key);

        if String_Len(token) == 0 {
            return ExitResult_MakeErr("ERROR: signing failed");
        }

        PrintLine(token);
        return ExitResult_MakeOk(0);
    }

    // =============================================================================
    // Command: verify
    // =============================================================================
    func CmdVerify(token: String, algName: String, keyArg: String) -> ExitResult {
        let algOpt: AlgOption = ParseAlg(algName);
        if algOpt.tag != AlgOption_Some {
            return ExitResult_MakeErr(f"ERROR: unknown algorithm '{algName}'");
        }
        let alg: JwtAlg = algOpt.data.Some_0;

        let keyRes: KeyResult = ResolveKey(alg, keyArg);
        if keyRes.tag != KeyResult_Ok {
            return ExitResult_MakeErr(keyRes.data.Err_0);
        }
        let key: String = keyRes.data.Ok_0;

        var header: String = "";
        var payload: String = "";

        if !Jwt_Decode(token, alg, key, &header, &payload) {
            return ExitResult_MakeErr("✗ Signature INVALID (or malformed token)");
        }

        PrintLine("✓ Signature valid");
        PrintLine("");
        PrintLine("Header:");
        PrintLine(header);
        PrintLine("");
        PrintLine("Payload:");
        PrintLine(payload);
        return ExitResult_MakeOk(0);
    }

    // =============================================================================
    // Command: decode (no verification)
    // =============================================================================
    func CmdDecode(token: String) -> ExitResult {
        let partCount: uint = String_SplitCount(token, ".");
        if partCount != 3 {
            return ExitResult_MakeErr("ERROR: not a valid JWT (expected 3 parts)");
        }

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

        let headerJson: String = Base64URL_Decode(headerB64);
        let payloadJson: String = Base64URL_Decode(payloadB64);

        PrintLine("Decoded (no verification):");
        PrintLine("");
        PrintLine("Header:");
        PrintLine(headerJson);
        PrintLine("");
        PrintLine("Payload:");
        PrintLine(payloadJson);
        PrintLine("");
        Print("Signature (base64url): ");
        PrintLine(sigB64);

        return ExitResult_MakeOk(0);
    }

    // =============================================================================
    // Command: keygen
    // =============================================================================
    func CmdKeygen(keyType: String) -> ExitResult {
        let ktOpt: KeyTypeOption = ParseKeyType(keyType);
        if ktOpt.tag != KeyTypeOption_Some {
            return ExitResult_MakeErr(f"ERROR: unknown key type '{keyType}'. Use: rsa, ecdsa, ed25519");
        }
        let kt: KeyType = ktOpt.data.Some_0;

        if kt.tag == KeyType_Ed25519 {
            let pubBuf: *void = bux_alloc(32);
            let priv: *void = bux_alloc(32);
            if !Ed25519_Keypair(pubBuf, priv) {
                return ExitResult_MakeErr("ERROR: Ed25519 key generation failed (OpenSSL 1.1.1+ required)");
            }
            let pubB64: String = bux_base64_encode(pubBuf as String, 32);
            let privB64: String = bux_base64_encode(priv as String, 32);
            PrintLine("Ed25519 keypair (base64):");
            PrintLine(f"  Public:  {pubB64}");
            PrintLine(f"  Private: {privB64}");
            return ExitResult_MakeOk(0);
        }
        if kt.tag == KeyType_Rsa {
            PrintLine("RSA key generation requires OpenSSL CLI:");
            PrintLine("  openssl genpkey -algorithm RSA -out private.pem -pkeyopt rsa_keygen_bits:2048");
            PrintLine("  openssl rsa -in private.pem -pubout -out public.pem");
            PrintLine("");
            PrintLine("(Key generation from within Bux requires PEM write support — coming soon)");
            return ExitResult_MakeOk(0);
        }
        if kt.tag == KeyType_Ecdsa {
            PrintLine("ECDSA key generation requires OpenSSL CLI:");
            PrintLine("  openssl ecparam -genkey -name prime256v1 -noout -out ec_private.pem");
            PrintLine("  openssl ec -in ec_private.pem -pubout -out ec_public.pem");
            PrintLine("");
            PrintLine("(Key generation from within Bux requires PEM write support — coming soon)");
            return ExitResult_MakeOk(0);
        }
        return ExitResult_MakeErr(f"ERROR: unknown key type '{keyType}'");
    }

    // =============================================================================
    // Dispatch a parsed command to its handler
    // =============================================================================
    func RunCommand(cmd: Cmd, args: Array<String>) -> ExitResult {
        if cmd.tag == Cmd_Help {
            PrintUsage();
            return ExitResult_MakeOk(0);
        }
        if cmd.tag == Cmd_Sign {
            if args.len < 5 {
                return ExitResult_MakeErr("ERROR: 'sign' requires: <alg> <key> <claims_json>");
            }
            return CmdSign(args[2], args[3], args[4]);
        }
        if cmd.tag == Cmd_Verify {
            if args.len < 5 {
                return ExitResult_MakeErr("ERROR: 'verify' requires: <token> <alg> <key>");
            }
            return CmdVerify(args[2], args[3], args[4]);
        }
        if cmd.tag == Cmd_Decode {
            if args.len < 3 {
                return ExitResult_MakeErr("ERROR: 'decode' requires: <token>");
            }
            return CmdDecode(args[2]);
        }
        if cmd.tag == Cmd_Keygen {
            if args.len < 3 {
                return ExitResult_MakeErr("ERROR: 'keygen' requires: <type>");
            }
            return CmdKeygen(args[2]);
        }
        return ExitResult_MakeErr("ERROR: unhandled command");
    }

    // =============================================================================
    // Main Entry Point
    // =============================================================================
    func Main() -> int {
        let argc: int = bux_argc();

        // Collect CLI arguments into a generic Array<String>
        let args: Array<String> = Array_New<String>(argc as uint);
        for i in 0..argc {
            Array_Push<String>(&args, bux_argv(i));
        }

        if args.len < 2 {
            PrintUsage();
            return 0;
        }

        let command: String = args[1];
        let cmdOpt: CmdOption = ParseCmd(command);
        if cmdOpt.tag != CmdOption_Some {
            PrintLine(f"ERROR: unknown command '{command}'");
            PrintLine(f"Run '{AppName} help' for usage.");
            return 1;
        }
        let cmd: Cmd = cmdOpt.data.Some_0;

        let result: ExitResult = RunCommand(cmd, args);
        if result.tag == ExitResult_Err {
            PrintLine(result.data.Err_0);
            return 1;
        }
        return result.data.Ok_0;
    }

} // module Main
