// Copyright (c) blackshirt. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module slhdsa

// Configuration options used in SLH-DSA key generation.
@[params]
pub struct KeyOpts {
pub mut:
    // An opaque represents the kind of SLH-DSA keys want to built.
    // See `enum Kind` for available options.
    kind Kind = .sha2_128s

    // flag, 0=random (default), 1= use seed bytes, 2 = use priv bytes, otherwise error
    flag int
    // when you set flag=1, builder will use seed bytes as a params,
    // so, make sure to pass seed bytes length != 0
    seed []u8
    // when flag=2, you should ensure private bytes length != 0
    priv []u8

    // This option below was not supported yet.
    //
    // Sets properties to be used when fetching algorithm implementations
    // used for SLH-DSA hashing operations.
    propq string
}

// Configurations parameters used for signing (and or verifying)
@[params]
pub struct SignerOpts {
pub mut:
    // optional context string up to 255 length, used in signing (verifying)
    context string

    // "message-encoding"
    // The default value of 1 uses 'Pure SLH-DSA Signature Generation'.
    // Setting it to 0 does not encode the message, which is used for testing,
    // but can also be used for 'Pre Hash SLH-DSA Signature Generation'.
    // If you set encoding to 0, you should provide the entropy.
    encoding int = 1

    // "test-entropy" used for testing to pass a optional random value.
    entropy []u8 // octet-string

    // "deterministic" integer option.
    // The default value of 0 generates a random value (using a DRBG) this is used when
    // processing the message. Setting this to 1 causes the private key seed to be used instead.
    // This value is ignored if "test-entropy" is set.
    deterministic int
}

// from_seed creates a new SLH-DSA PrivateKey from seed bytes and kind options.
// If the seed length was zero, it will create a key based on randomly generated seed.
// You should make sure, the seed bytes comes from trusted cryptographic secure source.
// The seed size was 3 times of `𝑛` parameter defined in the standard.
// The `𝑛` size maybe 16, 24 or 32 bytes length, depend on the chosen type.
fn PrivateKey.from_seed(seed []u8, kind Kind) !PrivateKey {
    // when seed bytes length was zero, build the key based on the random ones
    if seed.len == 0 {
        return PrivateKey.new(kind: kind)!
    }
    // Get the n size parameter set from the options
    nsize := kind.nsize()
    // The length of the seed bytes supplied must be 3 * nsize.
    if seed.len != 3 * nsize {
        return error('Unmatching seed length with kind supplied, need ${3 * nsize} bytes')
    }

    param_bld := C.OSSL_PARAM_BLD_new()
    assert param_bld != 0

    m := C.OSSL_PARAM_BLD_push_octet_string(param_bld, c'seed', seed.data, seed.len)
    if m <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        return error('OSSL_PARAM_BLD_push failed')
    }

    // build params
    params := C.OSSL_PARAM_BLD_to_param(param_bld)
    // create a desired context
    pctx := C.EVP_PKEY_CTX_new_from_name(0, kind.long_name(), 0)
    if params == 0 || pctx == 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_CTX_new or OSSL_PARAM_BLD_to_param failed')
    }
    pkey := C.EVP_PKEY_new()
    assert pkey != 0
    ke := C.EVP_PKEY_keygen_init(pctx)
    if ke <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        C.EVP_PKEY_free(pkey)
        return error('EVP_PKEY_keygen_init failed')
    }
    // Use EVP_PKEY_CTX_set_params() after calling EVP_PKEY_keygen_init().
    s := C.EVP_PKEY_CTX_set_params(pctx, params)
    if s != 1 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        C.EVP_PKEY_free(pkey)
        return error('EVP_PKEY_CTX_set_params failed')
    }

    ss := C.EVP_PKEY_keygen(pctx, &pkey)
    if ss <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        C.EVP_PKEY_free(pkey)
        return error('EVP_PKEY_keygen failed')
    }
    // TODO: right way to check the key
    pvkey := PrivateKey{
        key: pkey
    }
    // Cleans up
    C.OSSL_PARAM_BLD_free(param_bld)
    C.OSSL_PARAM_free(params)
    C.EVP_PKEY_CTX_free(pctx)

    return pvkey
}

fn PrivateKey.from_bytes(bytes []u8, kind Kind) !PrivateKey {
    // when bytes length was zero, build the key based on the random ones
    if bytes.len == 0 {
        return PrivateKey.new(kind: kind)!
    }
    // Get the n size parameter set from the options
    nsize := kind.nsize()
    // The private key has a size of 4 * n bytes.
    if bytes.len != 4 * nsize {
        return error('Unmatching private bytes length with kind supplied, need ${4 * nsize} bytes')
    }

    param_bld := C.OSSL_PARAM_BLD_new()
    assert param_bld != 0

    m := C.OSSL_PARAM_BLD_push_octet_string(param_bld, c'priv', bytes.data, bytes.len)
    if m <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        return error('OSSL_PARAM_BLD_push FAILED')
    }

    // build params
    params := C.OSSL_PARAM_BLD_to_param(param_bld)
    // create a desired context
    pctx := C.EVP_PKEY_CTX_new_from_name(0, kind.long_name(), 0)
    if params == 0 || pctx == 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_CTX_new or OSSL_PARAM_BLD_to_param failed')
    }

    pkey := C.EVP_PKEY_new()
    assert pkey != 0

    s := C.EVP_PKEY_fromdata_init(pctx)
    if s <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        C.EVP_PKEY_free(pkey)
        return error('EVP_PKEY_fromdata_init failed')
    }

    ss := C.EVP_PKEY_fromdata(pctx, &pkey, evp_pkey_keypair, params)
    if ss <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        C.EVP_PKEY_free(pkey)
        return error('EVP_PKEY_fromdata failed')
    }

    pvkey := PrivateKey{
        key: pkey
    }
    // Cleans up
    C.OSSL_PARAM_BLD_free(param_bld)
    C.OSSL_PARAM_free(params)
    C.EVP_PKEY_CTX_free(pctx)

    return pvkey
}

// from_bytes creates a new PublicKey with type of supported key and bytes array.
// If you dont provide the bytes, ie, supplied with zero-length bytes,
// it will be generated with random bytes for you.
pub fn PublicKey.from_bytes(bytes []u8, kind Kind) !PublicKey {
    // when bytes length was zero, build the key based on the random ones
    if bytes.len == 0 {
        pv := PrivateKey.new(kind: kind)!
        pbk := pv.public_key()!
        return pbk
    }
    // Get the n size parameter set from the options
    nsize := kind.nsize()
    // The public key has a size of 2 * n bytes.
    if bytes.len != 2 * nsize {
        return error('Unmatching public bytes length with kind supplied, need ${2 * nsize} bytes')
    }

    param_bld := C.OSSL_PARAM_BLD_new()
    assert param_bld != 0

    m := C.OSSL_PARAM_BLD_push_octet_string(param_bld, c'pub', bytes.data, bytes.len)
    if m <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        return error('OSSL_PARAM_BLD_push FAILED')
    }

    // build params
    params := C.OSSL_PARAM_BLD_to_param(param_bld)
    // create a desired context
    pctx := C.EVP_PKEY_CTX_new_from_name(0, kind.long_name(), 0)
    if params == 0 || pctx == 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_CTX_new or OSSL_PARAM_BLD_to_param failed')
    }

    pkey := C.EVP_PKEY_new()
    assert pkey != 0

    s := C.EVP_PKEY_fromdata_init(pctx)
    if s <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        C.EVP_PKEY_free(pkey)
        return error('EVP_PKEY_fromdata failed')
    }

    ss := C.EVP_PKEY_fromdata(pctx, &pkey, evp_pkey_public_key, params)
    if ss <= 0 {
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_CTX_free(pctx)
        C.EVP_PKEY_free(pkey)
        return error('EVP_PKEY_fromdata failed')
    }

    pbkey := PublicKey{
        key: pkey
    }
    // Cleans up
    C.OSSL_PARAM_BLD_free(param_bld)
    C.OSSL_PARAM_free(params)
    C.EVP_PKEY_CTX_free(pctx)

    return pbkey
}

// The enumeration of NID of SLHDSA parameters set. <br>
// See Table 2. SLH-DSA parameter sets of the Chapter 11. Parameter Sets<br>
// Each sets name indicates:
//
//     - the hash function family (SHA2 or SHAKE) that is used to instantiate the hash functions.
//    - the length in bits of the security parameter, in the 128, 192, and 256 respectives number.
//    - the mnemonic name indicates parameter to create relatively small signatures (`s`)
//      or to have relatively fast signature generation (`f`).
pub enum Kind {
    // SHA2-based family
    sha2_128s = C.NID_SLH_DSA_SHA2_128s
    sha2_128f = C.NID_SLH_DSA_SHA2_128f
    sha2_192s = C.NID_SLH_DSA_SHA2_192s
    sha2_192f = C.NID_SLH_DSA_SHA2_192f
    sha2_256s = C.NID_SLH_DSA_SHA2_256s
    sha2_256f = C.NID_SLH_DSA_SHA2_256f
    // SHAKE-based family
    shake_128s = C.NID_SLH_DSA_SHAKE_128s
    shake_128f = C.NID_SLH_DSA_SHAKE_128f
    shake_192s = C.NID_SLH_DSA_SHAKE_192s
    shake_192f = C.NID_SLH_DSA_SHAKE_192f
    shake_256s = C.NID_SLH_DSA_SHAKE_256s
    shake_256f = C.NID_SLH_DSA_SHAKE_256f
}

// nsize returns the size of underlying n parameter from current type.
@[inline]
fn (n Kind) nsize() int {
    match n {
        .sha2_128s, .sha2_128f, .shake_128s, .shake_128f { return 16 }
        .sha2_192s, .sha2_192f, .shake_192s, .shake_192f { return 24 }
        .sha2_256s, .sha2_256f, .shake_256s, .shake_256f { return 32 }
    }
}

fn (n Kind) str() string {
    match n {
        // vfmt off
        // SHA2-based family
        .sha2_128s { return "sha2_128s" }
        .sha2_128f { return "sha2_128f" }
        .sha2_192s { return "sha2_192s" }
        .sha2_192f { return "sha2_192f" }
        .sha2_256s { return "sha2_256s" }
        .sha2_256f { return "sha2_256f" }
        // SHAKE-based family
        .shake_128s { return "shake_128s" }
        .shake_128f { return "shake_128f" }
        .shake_192s { return "shake_192s" }
        .shake_192f { return "shake_192f" }
        .shake_256s { return "shake_256s" }
        .shake_256f { return "shake_256f" }
        // vfmt on
    }
}

// Kind long name as v string
@[inline]
fn (n Kind) ln_to_vstr() string {
    out := unsafe { n.long_name().vstring() }
    return out
}

// Kind long name as c-style string
@[inline]
fn (n Kind) long_name() &char {
    match n {
        // vfmt off
        // SHA2-based family
        .sha2_128s { return ln_slhdsa_sha2_128s }
        .sha2_128f { return ln_slhdsa_sha2_128f }
        .sha2_192s { return ln_slhdsa_sha2_192s }
        .sha2_192f { return ln_slhdsa_sha2_192f }
        .sha2_256s { return ln_slhdsa_sha2_256s }
        .sha2_256f { return ln_slhdsa_sha2_256f }
        // SHAKE-based family
        .shake_128s { return ln_slhdsa_shake_128s }
        .shake_128f { return ln_slhdsa_shake_128f }
        .shake_192s { return ln_slhdsa_shake_192s }
        .shake_192f { return ln_slhdsa_shake_192f }
        .shake_256s { return ln_slhdsa_shake_256s }
        .shake_256f { return ln_slhdsa_shake_256f }
        // vfmt on
    }
}

// Chapter 11. Parameters Set
struct ParamSet {
    // Algorithm name
    id  Kind
    n   int
    h   int
    d   int
    hp  int
    a   int
    k   int
    lgw int = 4
    m   int
    sc  int
    pkb int
    sig int
}

// Table 2. SLH-DSA parameter sets
const paramset = {
    //                              id     𝑛     ℎ     𝑑     ℎ′  𝑎     𝑘     𝑙𝑔𝑤 𝑚  sc pkb  sig
    'sha2_128s':  ParamSet{.sha2_128s, 16, 63, 7, 9, 12, 14, 4, 30, 1, 32, 7856}
    'sha2_128f':  ParamSet{.sha2_128f, 16, 66, 22, 3, 6, 33, 4, 34, 1, 32, 17088}
    'sha2_192s':  ParamSet{.sha2_192s, 24, 63, 7, 9, 14, 17, 4, 39, 3, 48, 16224}
    'sha2_192f':  ParamSet{.sha2_192f, 24, 66, 22, 3, 8, 33, 4, 42, 3, 48, 35664}
    'sha2_256s':  ParamSet{.sha2_256s, 32, 64, 8, 8, 14, 22, 4, 47, 5, 64, 29792}
    'sha2_256f':  ParamSet{.sha2_256f, 32, 68, 17, 4, 9, 35, 4, 49, 5, 64, 49856}
    // SHAKE family
    'shake_128s': ParamSet{.shake_128s, 16, 63, 7, 9, 12, 14, 4, 30, 1, 32, 7856}
    'shake_128f': ParamSet{.shake_128f, 16, 66, 22, 3, 6, 33, 4, 34, 1, 32, 17088}
    'shake_192s': ParamSet{.shake_192s, 24, 63, 7, 9, 14, 17, 4, 39, 3, 48, 16224}
    'shake_192f': ParamSet{.shake_192f, 24, 66, 22, 3, 8, 33, 4, 42, 3, 48, 35664}
    'shake_256s': ParamSet{.shake_256s, 32, 64, 8, 8, 14, 22, 4, 47, 5, 64, 29792}
    'shake_256f': ParamSet{.shake_256f, 32, 68, 17, 4, 9, 35, 4, 49, 5, 64, 49856}
}

fn ParamSet.from_kind(k Kind) ParamSet {
    return paramset[k.str()]
}

// Some helpers
//
fn key_algo_name(key &C.EVP_PKEY) &char {
    name := voidptr(C.EVP_PKEY_get0_type_name(key))
    assert name != 0
    return name
}

fn key_type_name(key &C.EVP_PKEY) &char {
    kn := voidptr(C.EVP_PKEY_get0_type_name(key))
    assert kn != 0
    return kn
}

fn key_description(key &C.EVP_PKEY) &char {
    kd := voidptr(C.EVP_PKEY_get0_description(key))
    assert kd != 0
    return kd
}