#include "key_derivation.h"

#include <stddef.h>
#include <stdint.h>
#include <string.h>

#include "esp_random.h"

#include "mbedtls/md.h"

#include "secp256k1.h"
#include "secp256k1_extrakeys.h"
#include "secp256k1_schnorrsig.h"

#define BIP32_HARDENED_FLAG 0x80000000u

typedef struct {
    uint8_t priv[32];
    uint8_t chain[32];
    uint8_t pub[33];
} hd_key_t;

static int hmac_sha512(const uint8_t *key,
                       size_t key_len,
                       const uint8_t *data,
                       size_t data_len,
                       uint8_t out[64])
{
    const mbedtls_md_info_t *md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA512);
    if (md_info == NULL) {
        return -1;
    }

    if (mbedtls_md_hmac(md_info, key, key_len, data, data_len, out) != 0) {
        return -1;
    }

    return 0;
}

static secp256k1_context *create_context(void)
{
    uint8_t rand32[32] = {0};
    secp256k1_context *ctx = secp256k1_context_create(SECP256K1_CONTEXT_SIGN | SECP256K1_CONTEXT_VERIFY);

    if (ctx == NULL) {
        return NULL;
    }

    esp_fill_random(rand32, sizeof(rand32));
    if (!secp256k1_context_randomize(ctx, rand32)) {
        memset(rand32, 0, sizeof(rand32));
        secp256k1_context_destroy(ctx);
        return NULL;
    }

    memset(rand32, 0, sizeof(rand32));
    return ctx;
}

static int compute_compressed_pubkey(const secp256k1_context *ctx, const uint8_t priv[32], uint8_t pub33[33])
{
    secp256k1_pubkey pubkey;
    size_t out_len = 33;

    if (!secp256k1_ec_pubkey_create(ctx, &pubkey, priv)) {
        return -1;
    }

    if (!secp256k1_ec_pubkey_serialize(ctx, pub33, &out_len, &pubkey, SECP256K1_EC_COMPRESSED)) {
        return -1;
    }

    if (out_len != 33) {
        return -1;
    }

    return 0;
}

static int bip32_master_from_seed(const secp256k1_context *ctx, const uint8_t seed[64], hd_key_t *master)
{
    static const uint8_t kBitcoinSeed[] = "Bitcoin seed";
    uint8_t i64[64] = {0};

    if (hmac_sha512(kBitcoinSeed, sizeof(kBitcoinSeed) - 1, seed, 64, i64) != 0) {
        return -1;
    }

    memcpy(master->priv, i64, 32);
    memcpy(master->chain, i64 + 32, 32);
    memset(i64, 0, sizeof(i64));

    if (!secp256k1_ec_seckey_verify(ctx, master->priv)) {
        return -1;
    }

    return compute_compressed_pubkey(ctx, master->priv, master->pub);
}

static int bip32_ckd_priv(const secp256k1_context *ctx,
                          const hd_key_t *parent,
                          uint32_t index,
                          hd_key_t *child)
{
    uint8_t data[37] = {0};
    uint8_t i64[64] = {0};

    if (index & BIP32_HARDENED_FLAG) {
        data[0] = 0x00;
        memcpy(data + 1, parent->priv, 32);
    } else {
        memcpy(data, parent->pub, 33);
    }

    data[33] = (uint8_t)((index >> 24) & 0xFF);
    data[34] = (uint8_t)((index >> 16) & 0xFF);
    data[35] = (uint8_t)((index >> 8) & 0xFF);
    data[36] = (uint8_t)(index & 0xFF);

    if (hmac_sha512(parent->chain, 32, data, sizeof(data), i64) != 0) {
        return -1;
    }

    memcpy(child->priv, parent->priv, 32);
    memcpy(child->chain, i64 + 32, 32);

    if (!secp256k1_ec_seckey_tweak_add(ctx, child->priv, i64)) {
        memset(i64, 0, sizeof(i64));
        return -1;
    }

    memset(i64, 0, sizeof(i64));

    if (!secp256k1_ec_seckey_verify(ctx, child->priv)) {
        return -1;
    }

    return compute_compressed_pubkey(ctx, child->priv, child->pub);
}

static int derive_nostr_key_internal(const uint8_t seed[64], uint32_t nostr_index, uint8_t privkey[32], uint8_t pubkey[32])
{
    const uint32_t path[5] = {
        44u | BIP32_HARDENED_FLAG,
        1237u | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
        0u,
        nostr_index,
    };

    secp256k1_context *ctx = NULL;
    hd_key_t node;

    if (seed == NULL || privkey == NULL || pubkey == NULL) {
        return -1;
    }

    memset(&node, 0, sizeof(node));

    ctx = create_context();
    if (ctx == NULL) {
        return -1;
    }

    if (bip32_master_from_seed(ctx, seed, &node) != 0) {
        secp256k1_context_destroy(ctx);
        return -1;
    }

    for (size_t i = 0; i < sizeof(path) / sizeof(path[0]); ++i) {
        hd_key_t next;
        memset(&next, 0, sizeof(next));

        if (bip32_ckd_priv(ctx, &node, path[i], &next) != 0) {
            memset(&node, 0, sizeof(node));
            secp256k1_context_destroy(ctx);
            return -1;
        }

        memset(&node, 0, sizeof(node));
        node = next;
    }

    memcpy(privkey, node.priv, 32);

    {
        secp256k1_keypair kp;
        secp256k1_xonly_pubkey xonly;

        if (!secp256k1_keypair_create(ctx, &kp, node.priv)) {
            memset(&node, 0, sizeof(node));
            secp256k1_context_destroy(ctx);
            return -1;
        }

        if (!secp256k1_keypair_xonly_pub(ctx, &xonly, NULL, &kp)) {
            memset(&node, 0, sizeof(node));
            secp256k1_context_destroy(ctx);
            return -1;
        }

        if (!secp256k1_xonly_pubkey_serialize(ctx, pubkey, &xonly)) {
            memset(&node, 0, sizeof(node));
            secp256k1_context_destroy(ctx);
            return -1;
        }
    }

    memset(&node, 0, sizeof(node));
    secp256k1_context_destroy(ctx);
    return 0;
}

int derive_nostr_key(const uint8_t seed[64], uint8_t privkey[32], uint8_t pubkey[32])
{
    return derive_nostr_key_internal(seed, 0u, privkey, pubkey);
}

int derive_nostr_key_index(const uint8_t seed[64], uint32_t nostr_index, uint8_t privkey[32], uint8_t pubkey[32])
{
    return derive_nostr_key_internal(seed, nostr_index, privkey, pubkey);
}

int schnorr_sign32(const uint8_t privkey[32], const uint8_t msg32[32], uint8_t sig64[64])
{
    secp256k1_context *ctx = NULL;
    secp256k1_keypair kp;
    uint8_t aux[32] = {0};

    if (privkey == NULL || msg32 == NULL || sig64 == NULL) {
        return -1;
    }

    ctx = create_context();
    if (ctx == NULL) {
        return -1;
    }

    if (!secp256k1_keypair_create(ctx, &kp, privkey)) {
        secp256k1_context_destroy(ctx);
        return -1;
    }

    esp_fill_random(aux, sizeof(aux));
    if (!secp256k1_schnorrsig_sign32(ctx, sig64, msg32, &kp, aux)) {
        memset(aux, 0, sizeof(aux));
        secp256k1_context_destroy(ctx);
        return -1;
    }

    memset(aux, 0, sizeof(aux));
    secp256k1_context_destroy(ctx);
    return 0;
}