// Copyright (c) 2019-2024 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module ecdsa

import hash

// NID constants
//
// NIST P-256 is referred to as secp256r1 and prime256v1, defined as #define NID_X9_62_prime256v1 415
// Different names, but they are all the same.
// https://www.rfc-editor.org/rfc/rfc4492.html#appendix-A
const nid_prime256v1 = C.NID_X9_62_prime256v1

// NIST P-384, ie, secp384r1 curve, defined as #define NID_secp384r1 715
const nid_secp384r1 = C.NID_secp384r1

// NIST P-521, ie, secp521r1 curve, defined as #define NID_secp521r1 716
const nid_secp521r1 = C.NID_secp521r1

// Bitcoin curve, defined as #define NID_secp256k1 714
const nid_secp256k1 = C.NID_secp256k1

// #define NID_X9_62_id_ecPublicKey   408
const nid_ec_publickey = C.NID_X9_62_id_ecPublicKey
// C.EVP_PKEY_EC = NID_X9_62_id_ecPublicKey
const nid_evp_pkey_ec = C.EVP_PKEY_EC
// we only support this
const openssl_ec_named_curve = C.OPENSSL_EC_NAMED_CURVE

// https://docs.openssl.org/3.0/man3/EVP_PKEY_fromdata/#selections
const evp_pkey_keypair = C.EVP_PKEY_KEYPAIR

// POINT_CONVERSION FORAMT
const point_conversion_uncompressed = 4

// Nid is an enumeration of the supported curves
pub enum Nid {
    prime256v1 = C.NID_X9_62_prime256v1
    secp384r1  = C.NID_secp384r1
    secp521r1  = C.NID_secp521r1
    secp256k1  = C.NID_secp256k1
}

// we need this group (cruve) name representation to pass them into needed routines
fn (nid Nid) str() string {
    match nid {
        // TODO: maybe better relies on info from underlying C defined constants,
        // ie, #define SN_X9_62_prime256v1     "prime256v1" etc
        .prime256v1 { return 'prime256v1' }
        .secp384r1 { return 'secp384r1' }
        .secp521r1 { return 'secp521r1' }
        .secp256k1 { return 'secp256k1' }
    }
}

// CurveOptions represents configuration options to drive keypair generation.
@[params]
pub struct CurveOptions {
pub mut:
    // default to NIST P-256 curve
    nid Nid = .prime256v1
    // by default, allow arbitrary size of seed bytes as key.
    // Set it to `true` when you need fixed size, using the curve key size.
    // Its main purposes is to support the `.new_key_from_seed` call.
    fixed_size bool
}

// HashConfig is an enumeration of the possible options for key signing (verifying).
pub enum HashConfig {
    with_recommended_hash
    with_no_hash
    with_custom_hash
}

// SignerOpts represents configuration options to drive signing and verifying process.
@[params]
pub struct SignerOpts {
pub mut:
    // default to .with_recommended_hash
    hash_config HashConfig = .with_recommended_hash
    // make sense when HashConfig != with_recommended_hash
    allow_smaller_size bool
    allow_custom_hash  bool
    // set to non-nil if allow_custom_hash was true
    custom_hash &hash.Hash = unsafe { nil }
}

// KeyFlag is an enumeration of possible options to support flexible of PrivateKey key size.
enum KeyFlag {
    // flexible flag to allow flexible-size of seed bytes
    flexible
    // fixed flag for using underlying curve key size
    fixed
}

// generate_key generates a new key pair. If opt was not provided, its default to prime256v1 curve.
// If you want another curve, use `pubkey, pivkey := ecdsa.generate_key(nid: .secp384r1)!` instead.
pub fn generate_key(opt CurveOptions) !(PublicKey, PrivateKey) {
    // This can be simplified to just more simpler one
    pv := PrivateKey.new(opt)!
    pb := pv.public_key()!

    return pb, pv
}

// new_key_from_seed creates a new private key from the seed bytes. If opt was not provided,
// its default to prime256v1 curve.
//
// Notes on the seed:
//
// You should make sure, the seed bytes come from a cryptographically secure random generator,
// likes the `crypto.rand` or other trusted sources.
// Internally, the seed size's would be checked to not exceed the key size of underlying curve,
// ie, 32 bytes length for p-256 and secp256k1, 48 bytes length for p-384 and 66 bytes length for p-521.
// Its recommended to use seed with bytes length matching with underlying curve key size.
pub fn new_key_from_seed(seed []u8, opt CurveOptions) !PrivateKey {
    // Early exit check
    if seed.len == 0 {
        return error('Seed with null-length was not allowed')
    }
    evpkey := evpkey_from_seed(seed, opt) or { return err }
    num_bits := C.EVP_PKEY_get_bits(evpkey)
    key_size := (num_bits + 7) / 8
    if seed.len > key_size {
        C.EVP_PKEY_free(evpkey)
        return error('Seed length exceeds key size')
    }
    // Check if its using fixed key size or flexible one
    if opt.fixed_size {
        if seed.len != key_size {
            C.EVP_PKEY_free(evpkey)
            return error('seed size doesnt match with curve key size')
        }
    }
    mut pvkey := PrivateKey{
        evpkey: evpkey
    }

    if opt.fixed_size {
        // using fixed one
        pvkey.ks_flag = .fixed
        pvkey.ks_size = key_size
    } else {
        pvkey.ks_size = seed.len
    }

    return pvkey
}

// PrivateKey represents ECDSA private key. Actually its a key pair,
// contains private key and public key parts.
pub struct PrivateKey {
    // The new high level of keypair opaque
    evpkey &C.EVP_PKEY
mut:
    // ks_flag with .flexible value allowing
    // flexible-size seed bytes as key.
    // When it is `.fixed`, it will use the underlying key size.
    ks_flag KeyFlag = .flexible
    // ks_size stores size of the seed bytes when ks_flag was .flexible.
    // You should set it to a non zero value
    ks_size int
}

// PrivateKey.new creates a new key pair. By default, it would create a prime256v1 based key.
// Dont forget to call `.free()` after finish with your key.
pub fn PrivateKey.new(opt CurveOptions) !PrivateKey {
    evpkey := C.EVP_PKEY_new()
    pctx := C.EVP_PKEY_CTX_new_id(nid_evp_pkey_ec, 0)
    if pctx == 0 {
        C.EVP_PKEY_free(evpkey)
        C.EVP_PKEY_CTX_free(pctx)
        return error('C.EVP_PKEY_CTX_new_id failed')
    }
    nt := C.EVP_PKEY_keygen_init(pctx)
    if nt <= 0 {
        C.EVP_PKEY_free(evpkey)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_keygen_init failed')
    }
    // set the group (curve)
    cn := C.EVP_PKEY_CTX_set_ec_paramgen_curve_nid(pctx, int(opt.nid))
    if cn <= 0 {
        C.EVP_PKEY_free(evpkey)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_CTX_set_ec_paramgen_curve_nid')
    }
    // explicitly only allowing named curve, likely its the default on 3.0.
    pn := C.EVP_PKEY_CTX_set_ec_param_enc(pctx, openssl_ec_named_curve)
    if pn <= 0 {
        C.EVP_PKEY_free(evpkey)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_CTX_set_ec_param_enc failed')
    }
    // generates keypair
    nr := C.EVP_PKEY_keygen(pctx, &evpkey)
    if nr <= 0 {
        C.EVP_PKEY_free(evpkey)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_keygen failed')
    }

    // Cleans up the context
    C.EVP_PKEY_CTX_free(pctx)
    // when using default this function, its using underlying curve key size
    // and discarded opt.fixed_size flag when its not set.
    priv_key := PrivateKey{
        evpkey:  evpkey
        ks_flag: .fixed
    }
    return priv_key
}

// sign performs signing the message with the options. By default options,
// it will perform hashing before signing the message.
pub fn (pv PrivateKey) sign(message []u8, opt SignerOpts) ![]u8 {
    digest := calc_digest_with_evpkey(pv.evpkey, message, opt)!
    return sign_digest(pv.evpkey, digest)!
}

// sign_with_options signs message with the options. It will be deprecated,
// Use `PrivateKey.sign()` instead.
@[deprecated: 'use PrivateKey.sign() instead']
pub fn (pv PrivateKey) sign_with_options(message []u8, opt SignerOpts) ![]u8 {
    return pv.sign(message, opt)
}

// bytes represent private key as bytes.
pub fn (pv PrivateKey) bytes() ![]u8 {
    mut bn := &C.BIGNUM(unsafe { nil })
    // retrieves a BIGNUM value associated with a 'priv' key name
    n := C.EVP_PKEY_get_bn_param(pv.evpkey, c'priv', &bn)
    if n <= 0 {
        C.BN_free(bn)
        return error('EVP_PKEY_get_bn_param failed')
    }
    num_bytes := (C.BN_num_bits(bn) + 7) / 8
    // Get the buffer size to store the seed.
    size := if pv.ks_flag == .flexible {
        // should be non zero
        pv.ks_size
    } else {
        num_bytes
    }
    mut buf := []u8{len: int(size)}
    res := C.BN_bn2binpad(bn, buf.data, i32(size))
    if res <= 0 {
        C.BN_free(bn)
        return error('Failed to convert BIGNUM to bytes')
    }
    C.BN_free(bn)
    return buf
}

// seed gets the seed (private key bytes). It will be deprecated.
// Use `PrivateKey.bytes()` instead.
@[deprecated: 'use PrivateKey.bytes() instead']
pub fn (pv PrivateKey) seed() ![]u8 {
    return pv.bytes()
}

// public_key gets the PublicKey from private key.
pub fn (pv PrivateKey) public_key() !PublicKey {
    // Using duplicate key and removes (clears out) priv key
    pbkey := C.EVP_PKEY_dup(pv.evpkey)
    bn := C.BN_new()
    n := C.EVP_PKEY_set_bn_param(pbkey, c'priv', bn)
    assert n == 1
    // cleansup
    C.BN_free(bn)
    return PublicKey{
        evpkey: pbkey
    }
}

// equal compares two private keys was equal.
pub fn (priv_key PrivateKey) equal(other PrivateKey) bool {
    eq := C.EVP_PKEY_eq(voidptr(priv_key.evpkey), voidptr(other.evpkey))
    return eq == 1
}

// free clears out allocated memory for PrivateKey. Dont use PrivateKey after calling `.free()`
pub fn (pv &PrivateKey) free() {
    C.EVP_PKEY_free(pv.evpkey)
}

// PublicKey represents ECDSA public key for verifying message.
pub struct PublicKey {
    // The new high level of keypair opaque
    evpkey &C.EVP_PKEY
}

// verify verifies a message with the signature are valid with public key provided .
// You should provide it with the same SignerOpts used with the `.sign()` call.
// or verify would fail (false).
pub fn (pb PublicKey) verify(message []u8, sig []u8, opt SignerOpts) !bool {
    digest := calc_digest_with_evpkey(pb.evpkey, message, opt)!
    return verify_signature(pb.evpkey, sig, digest)
}

// equal compares two public keys was equal.
pub fn (pub_key PublicKey) equal(other PublicKey) bool {
    eq := C.EVP_PKEY_eq(voidptr(pub_key.evpkey), voidptr(other.evpkey))
    return eq == 1
}

// free clears out allocated memory for PublicKey. Dont use PublicKey after calling `.free()`
pub fn (pb &PublicKey) free() {
    C.EVP_PKEY_free(pb.evpkey)
}

// Helpers
//
// calc_digest_with_evpkey get the digest of the messages under the EVP_PKEY and options
fn calc_digest_with_evpkey(key &C.EVP_PKEY, message []u8, opt SignerOpts) ![]u8 {
    if message.len == 0 {
        return error('null-length messages')
    }
    match opt.hash_config {
        .with_no_hash, .with_recommended_hash {
            md := default_digest(key)!
            return calc_digest_with_md(message, md)!
        }
        .with_custom_hash {
            mut cfg := opt
            if !cfg.allow_custom_hash {
                return error('custom hash was not allowed, set it into true')
            }
            if cfg.custom_hash == unsafe { nil } {
                return error('Custom hasher was not defined')
            }
            key_size := evp_key_size(key)!
            if key_size > cfg.custom_hash.size() {
                if !cfg.allow_smaller_size {
                    return error('Hash into smaller size than current key size was not allowed')
                }
            }
            // we need to reset the custom hash before writes message
            cfg.custom_hash.reset()
            _ := cfg.custom_hash.write(message)!
            digest := cfg.custom_hash.sum([]u8{})

            return digest
        }
    }

    return error('Not should be here')
}

// sign_digest signs the digest with the key. Under the hood, EVP_PKEY_sign() does not
// hash the data to be signed, and therefore is normally used to sign digests.
fn sign_digest(key &C.EVP_PKEY, digest []u8) ![]u8 {
    ctx := C.EVP_PKEY_CTX_new(key, 0)
    if ctx == 0 {
        C.EVP_PKEY_CTX_free(ctx)
        return error('EVP_PKEY_CTX_new failed')
    }
    sin := C.EVP_PKEY_sign_init(ctx)
    if sin != 1 {
        C.EVP_PKEY_CTX_free(ctx)
        return error('EVP_PKEY_sign_init failed')
    }
    // siglen was used to store the size of the signature output. When EVP_PKEY_sign
    // was called with NULL signature buffer, siglen will tell maximum size of signature.
    siglen := usize(C.EVP_PKEY_size(key))
    sig := []u8{len: int(siglen)}

    // calls directly with sign
    do := C.EVP_PKEY_sign(ctx, sig.data, &siglen, digest.data, digest.len)
    if do <= 0 {
        unsafe { sig.free() }
        C.EVP_PKEY_CTX_free(ctx)
        return error('EVP_PKEY_sign fails to sign message')
    }
    // siglen now contains actual length of the signature buffer.
    signed := sig[..int(siglen)].clone()

    // Cleans up
    unsafe { sig.free() }
    C.EVP_PKEY_CTX_free(ctx)

    return signed
}

// verify_signature verifies the signature for the digest under the provided key.
fn verify_signature(key &C.EVP_PKEY, sig []u8, digest []u8) bool {
    ctx := C.EVP_PKEY_CTX_new(key, 0)
    if ctx == 0 {
        C.EVP_PKEY_CTX_free(ctx)
        return false
    }
    vinit := C.EVP_PKEY_verify_init(ctx)
    if vinit != 1 {
        C.EVP_PKEY_CTX_free(ctx)
        return false
    }
    res := C.EVP_PKEY_verify(ctx, sig.data, sig.len, digest.data, digest.len)
    if res <= 0 {
        C.EVP_PKEY_CTX_free(ctx)
        return false
    }
    C.EVP_PKEY_CTX_free(ctx)
    return res == 1
}

// calc_digest_with_md get the digest of the msg using md digest algorithm
fn calc_digest_with_md(msg []u8, md &C.EVP_MD) ![]u8 {
    ctx := C.EVP_MD_CTX_new()
    if ctx == 0 {
        C.EVP_MD_CTX_free(ctx)
        return error('EVP_MD_CTX_new failed')
    }
    nt := C.EVP_DigestInit(ctx, md)
    assert nt == 1
    upd := C.EVP_DigestUpdate(ctx, msg.data, msg.len)
    assert upd == 1

    size := u32(C.EVP_MD_get_size(md))
    out := []u8{len: int(size)}

    fin := C.EVP_DigestFinal(ctx, out.data, &size)
    assert fin == 1

    digest := out[..size].clone()
    // cleans up
    unsafe { out.free() }
    C.EVP_MD_CTX_free(ctx)

    return digest
}

// default_digest gets the default digest (hash) algorithm for this key.
fn default_digest(key &C.EVP_PKEY) !&C.EVP_MD {
    // get bits size of this key
    bits_size := C.EVP_PKEY_get_bits(key)
    if bits_size <= 0 {
        return error(' this size isnt available.')
    }
    // based on this bits_size, choose appropriate digest algorithm
    match true {
        bits_size <= 256 {
            return voidptr(C.EVP_sha256())
        }
        bits_size > 256 && bits_size <= 384 {
            return voidptr(C.EVP_sha384())
        }
        bits_size > 384 {
            return voidptr(C.EVP_sha512())
        }
        else {
            return error('Unsupported bits size')
        }
    }

    return error('should not here')
}

// Build EVP_PKEY from raw seed of bytes and options.
fn evpkey_from_seed(seed []u8, opt CurveOptions) !&C.EVP_PKEY {
    // This routine mostly comes from the official docs with adds some checking at
    // https://docs.openssl.org/3.0/man3/EVP_PKEY_fromdata/#creating-an-ecc-keypair-using-raw-key-data
    //
    // convert the seed bytes to BIGNUM.
    bn := C.BN_bin2bn(seed.data, seed.len, 0)
    if bn == 0 {
        C.BN_free(bn)
        return error('BN_bin2bn failed from seed')
    }
    // build the group (curve) from the options.
    group := C.EC_GROUP_new_by_curve_name(int(opt.nid))
    if group == 0 {
        C.EC_GROUP_free(group)
        C.BN_free(bn)
        return error('EC_GROUP_new_by_curve_name failed')
    }
    // Build EC_POINT from this BIGNUM and gets bytes represantion of this point
    // in uncompressed format.
    point := ec_point_mult(group, bn) or { return err }
    pub_bytes := point_2_buf(group, point, point_conversion_uncompressed)!

    // Lets build params builder
    param_bld := C.OSSL_PARAM_BLD_new()
    assert param_bld != 0

    // push the group, private and public key bytes infos into the builder
    n := C.OSSL_PARAM_BLD_push_utf8_string(param_bld, c'group', voidptr(opt.nid.str().str), 0)
    m := C.OSSL_PARAM_BLD_push_BN(param_bld, c'priv', bn)
    o := C.OSSL_PARAM_BLD_push_octet_string(param_bld, c'pub', pub_bytes.data, usize(pub_bytes.len))
    if n <= 0 || m <= 0 || o <= 0 {
        C.EC_POINT_free(point)
        C.BN_free(bn)
        C.EC_GROUP_free(group)
        C.OSSL_PARAM_BLD_free(param_bld)
        return error('OSSL_PARAM_BLD_push FAILED')
    }
    // Setup the new key
    mut pkey := C.EVP_PKEY_new()
    assert pkey != 0

    // build parameter, initialize and build the key from params
    params := C.OSSL_PARAM_BLD_to_param(param_bld)
    pctx := C.EVP_PKEY_CTX_new_id(nid_evp_pkey_ec, 0)
    if params == 0 || pctx == 0 {
        C.EC_POINT_free(point)
        C.BN_free(bn)
        C.EC_GROUP_free(group)
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_free(pkey)
        if pctx == 0 {
            C.EVP_PKEY_CTX_free(pctx)
        }
        return error('EVP_PKEY_CTX_new or OSSL_PARAM_BLD_to_param failed')
    }
    // initialize key and build the key from builded params context.
    p := C.EVP_PKEY_fromdata_init(pctx)
    q := C.EVP_PKEY_fromdata(pctx, &pkey, evp_pkey_keypair, params)
    if p <= 0 || q <= 0 {
        C.EC_POINT_free(point)
        C.BN_free(bn)
        C.EC_GROUP_free(group)
        C.OSSL_PARAM_BLD_free(param_bld)
        C.OSSL_PARAM_free(params)
        C.EVP_PKEY_free(pkey)
        C.EVP_PKEY_CTX_free(pctx)
        return error('EVP_PKEY_fromdata failed')
    }
    // After this step, we have build the key in pkey
    // TODO: right way to check the builded key

    // Cleans up
    C.EC_POINT_free(point)
    C.BN_free(bn)
    C.EC_GROUP_free(group)
    C.OSSL_PARAM_BLD_free(param_bld)
    C.OSSL_PARAM_free(params)
    C.EVP_PKEY_CTX_free(pctx)

    return pkey
}

// ec_point_mult performs point multiplications, point = bn * generator
fn ec_point_mult(group &C.EC_GROUP, bn &C.BIGNUM) !&C.EC_POINT {
    // Create a new EC_POINT object for the public key
    point := C.EC_POINT_new(group)
    // Create a new BN_CTX object for efficient BIGNUM operations
    ctx := C.BN_CTX_new()
    if ctx == 0 {
        C.EC_POINT_free(point)
        C.BN_CTX_free(ctx)
        return error('Failed to create BN_CTX')
    }

    // Perform the point multiplication to compute the public key: point = bn * G
    res := C.EC_POINT_mul(group, point, bn, 0, 0, ctx)
    if res != 1 {
        C.EC_POINT_free(point)
        C.BN_CTX_free(ctx)
        return error('Failed to compute public key')
    }
    C.BN_CTX_free(ctx)
    return point
}

// maximum key size we supported was 64 bytes.
const default_point_bufsize = 160 // 2 * 64 + 1 + extra

// point_2_buf gets bytes representation of the EC_POINT
fn point_2_buf(group &C.EC_GROUP, point &C.EC_POINT, fmt int) ![]u8 {
    ctx := C.BN_CTX_new()
    mut pbuf := &u8(unsafe { nil })
    // Notes from the docs:
    // EC_POINT_point2buf() allocates a buffer of suitable length and writes an EC_POINT to it in octet format.
    // The allocated buffer is written to *pbuf and its length is returned.
    // The caller must free up the allocated buffer with a call to OPENSSL_free().
    // Since the allocated buffer value is written to *pbuf the pbuf parameter MUST NOT be NULL.
    // So, we explicitly call `.OPENSSL_free` on the allocated buffer.
    n := C.EC_POINT_point2buf(group, point, fmt, &pbuf, ctx)
    if n <= 0 {
        C.BN_CTX_free(ctx)
        C.OPENSSL_free(voidptr(pbuf))
        return error('Get null length of buf')
    }
    // Gets the copy of the result with the correct length
    result := unsafe { pbuf.vbytes(int(n)).clone() }

    C.OPENSSL_free(voidptr(pbuf))
    C.BN_CTX_free(ctx)

    return result
}