#include "key_derivation.h"
#include "pq_crypto_firmware.h"

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

#include "esp_random.h"
#include "esp_log.h"

#include "mbedtls/md.h"
#include "mbedtls/ecp.h"
#include "mbedtls/pk.h"
#include "psa/crypto.h"

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

static const char *KD_TAG = "key_derivation";

#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;
}

/* ====================================================================
 * Phase 7: ed25519, x25519, and post-quantum key derivation
 * ==================================================================== */

/* SLIP-0010 all-hardened derivation for ed25519/x25519.
 *
 * SLIP-0010 uses HMAC-SHA512 with a "ed25519 seed" or curve-specific key
 * for the master key, and all derivation steps are hardened (the parent
 * private key is prepended to the index data).
 *
 * For ed25519/x25519, the derived 512-bit HMAC output is split:
 *   - first 32 bytes = private key (the scalar)
 *   - last 32 bytes = chain code
 *
 * The private key IS the ed25519/x25519 secret — no tweak-add is needed
 * (unlike secp256k1 BIP-32 where the child priv = parent_priv + HMAC).
 */

/* SLIP-0010 master key from seed: HMAC-SHA512(key="ed25519 seed", data=seed) */
static int slip10_master_from_seed(const uint8_t seed[64],
                                   uint8_t priv[32], uint8_t chain[32]) {
    static const uint8_t kEd25519Seed[] = "ed25519 seed";
    uint8_t i64[64] = {0};

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

    memcpy(priv, i64, 32);
    memcpy(chain, i64 + 32, 32);
    memset(i64, 0, sizeof(i64));
    return 0;
}

/* SLIP-0010 hardened child derivation:
 * HMAC-SHA512(key=chain, data=0x00 || priv || index_be32) */
static int slip10_ckd_priv(const uint8_t parent_priv[32],
                           const uint8_t parent_chain[32],
                           uint32_t index,
                           uint8_t child_priv[32],
                           uint8_t child_chain[32]) {
    uint8_t data[37];
    uint8_t i64[64] = {0};

    /* Hardened derivation: 0x00 || priv || index (big-endian) */
    data[0] = 0x00;
    memcpy(data + 1, parent_priv, 32);
    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) {
        memset(data, 0, sizeof(data));
        return -1;
    }

    memcpy(child_priv, i64, 32);
    memcpy(child_chain, i64 + 32, 32);
    memset(data, 0, sizeof(data));
    memset(i64, 0, sizeof(i64));
    return 0;
}

/* Derive a 32-byte seed via SLIP-0010 all-hardened path.
 * path[] is an array of hardened indices (the caller sets the hardened flag).
 * Returns the final 32-byte private material in `out_seed`. */
static int slip10_derive_seed(const uint8_t seed[64],
                              const uint32_t *path, size_t path_len,
                              uint8_t out_seed[32]) {
    uint8_t priv[32], chain[32], next_priv[32], next_chain[32];
    size_t i;

    if (slip10_master_from_seed(seed, priv, chain) != 0) {
        return -1;
    }

    for (i = 0; i < path_len; i++) {
        if (slip10_ckd_priv(priv, chain, path[i],
                            next_priv, next_chain) != 0) {
            memset(priv, 0, sizeof(priv));
            memset(chain, 0, sizeof(chain));
            return -1;
        }
        memcpy(priv, next_priv, 32);
        memcpy(chain, next_chain, 32);
    }

    memcpy(out_seed, priv, 32);
    memset(priv, 0, sizeof(priv));
    memset(chain, 0, sizeof(chain));
    memset(next_priv, 0, sizeof(next_priv));
    memset(next_chain, 0, sizeof(next_chain));
    return 0;
}

/* --- ed25519 --- */

int derive_ed25519_key(const uint8_t seed[64], uint32_t index,
                       uint8_t privkey[32], uint8_t pubkey[32]) {
    /* m/44'/102001'/<index>'/0'/0' — all hardened (SLIP-0010) */
    const uint32_t path[5] = {
        44u | BIP32_HARDENED_FLAG,
        102001u | BIP32_HARDENED_FLAG,
        index | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
    };
    uint8_t derived_seed[32];

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

    if (slip10_derive_seed(seed, path, 5, derived_seed) != 0) {
        return -1;
    }

    /* The SLIP-0010 derived 32 bytes IS the ed25519 private key. */
    memcpy(privkey, derived_seed, 32);

    /* Derive the ed25519 public key via PSA crypto (IDF v5.x mbedtls has no
     * mbedtls_ed25519_make_public). Import the private key, export the pub. */
    psa_status_t status;
    psa_key_id_t key_id = 0;
    psa_key_attributes_t attrs = PSA_KEY_ATTRIBUTES_INIT;
    size_t pub_len = 0;

    psa_crypto_init();
    psa_set_key_type(&attrs, PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_TWISTED_EDWARDS));
    psa_set_key_bits(&attrs, 255);
    psa_set_key_usage_flags(&attrs, PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_MESSAGE);
    psa_set_key_algorithm(&attrs, PSA_ALG_PURE_EDDSA);

    status = psa_import_key(&attrs, privkey, 32, &key_id);
    if (status != PSA_SUCCESS) {
        ESP_LOGE(KD_TAG, "ed25519 psa_import failed: %d", (int)status);
        memset(derived_seed, 0, sizeof(derived_seed));
        memset(privkey, 0, 32);
        return -1;
    }
    status = psa_export_public_key(key_id, pubkey, 32, &pub_len);
    psa_destroy_key(key_id);
    if (status != PSA_SUCCESS || pub_len != 32) {
        ESP_LOGE(KD_TAG, "ed25519 psa_export_public failed: %d", (int)status);
        memset(derived_seed, 0, sizeof(derived_seed));
        memset(privkey, 0, 32);
        return -1;
    }

    memset(derived_seed, 0, sizeof(derived_seed));
    return 0;
}

/* ed25519 sign via PSA (PureEdDSA — signs the raw message, not pre-hashed). */
int ed25519_sign32(const uint8_t privkey[32], const uint8_t msg32[32],
                   uint8_t sig64[64]) {
    return ed25519_sign_msg(privkey, msg32, 32, sig64);
}

int ed25519_sign_msg(const uint8_t privkey[32], const uint8_t *msg, size_t msg_len,
                     uint8_t sig64[64]) {
    psa_status_t status;
    psa_key_id_t key_id = 0;
    psa_key_attributes_t attrs = PSA_KEY_ATTRIBUTES_INIT;
    size_t sig_len = 0;

    psa_crypto_init();
    psa_set_key_type(&attrs, PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_TWISTED_EDWARDS));
    psa_set_key_bits(&attrs, 255);
    psa_set_key_usage_flags(&attrs, PSA_KEY_USAGE_SIGN_MESSAGE);
    psa_set_key_algorithm(&attrs, PSA_ALG_PURE_EDDSA);

    status = psa_import_key(&attrs, privkey, 32, &key_id);
    if (status != PSA_SUCCESS) {
        ESP_LOGE(KD_TAG, "ed25519 sign psa_import failed: %d", (int)status);
        return -1;
    }
    status = psa_sign_message(key_id, PSA_ALG_PURE_EDDSA,
                              msg, msg_len, sig64, 64, &sig_len);
    psa_destroy_key(key_id);
    if (status != PSA_SUCCESS || sig_len != 64) {
        ESP_LOGE(KD_TAG, "ed25519 psa_sign_message failed: %d", (int)status);
        return -1;
    }
    return 0;
}

int ed25519_verify_msg(const uint8_t *sig, size_t sig_len,
                       const uint8_t *msg, size_t msg_len,
                       const uint8_t pubkey[32]) {
    psa_status_t status;
    psa_key_id_t key_id = 0;
    psa_key_attributes_t attrs = PSA_KEY_ATTRIBUTES_INIT;

    psa_crypto_init();
    psa_set_key_type(&attrs, PSA_KEY_TYPE_ECC_PUBLIC_KEY(PSA_ECC_FAMILY_TWISTED_EDWARDS));
    psa_set_key_bits(&attrs, 255);
    psa_set_key_usage_flags(&attrs, PSA_KEY_USAGE_VERIFY_MESSAGE);
    psa_set_key_algorithm(&attrs, PSA_ALG_PURE_EDDSA);

    status = psa_import_key(&attrs, pubkey, 32, &key_id);
    if (status != PSA_SUCCESS) {
        ESP_LOGE(KD_TAG, "ed25519 verify psa_import failed: %d", (int)status);
        return -1;
    }
    status = psa_verify_message(key_id, PSA_ALG_PURE_EDDSA,
                                msg, msg_len, sig, sig_len);
    psa_destroy_key(key_id);
    return (status == PSA_SUCCESS) ? 0 : -1;
}

/* --- x25519 --- */

int derive_x25519_key(const uint8_t seed[64], uint32_t index,
                      uint8_t privkey[32], uint8_t pubkey[32]) {
    /* m/44'/102002'/<index>'/0'/0' — all hardened (SLIP-0010) */
    const uint32_t path[5] = {
        44u | BIP32_HARDENED_FLAG,
        102002u | BIP32_HARDENED_FLAG,
        index | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
    };
    uint8_t derived_seed[32];
    psa_status_t status;
    psa_key_id_t key_id = 0;
    psa_key_attributes_t attrs = PSA_KEY_ATTRIBUTES_INIT;
    size_t pub_len = 0;

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

    if (slip10_derive_seed(seed, path, 5, derived_seed) != 0) {
        return -1;
    }

    /* The SLIP-0010 derived 32 bytes IS the x25519 private key.
     * PSA imports it and exports the public key (PSA handles clamping). */
    memcpy(privkey, derived_seed, 32);
    memset(derived_seed, 0, sizeof(derived_seed));

    psa_crypto_init();
    psa_set_key_type(&attrs, PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY));
    psa_set_key_bits(&attrs, 255);
    psa_set_key_usage_flags(&attrs, PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE);
    psa_set_key_algorithm(&attrs, PSA_ALG_ECDH);

    status = psa_import_key(&attrs, privkey, 32, &key_id);
    if (status != PSA_SUCCESS) {
        ESP_LOGE(KD_TAG, "x25519 psa_import failed: %d", (int)status);
        memset(privkey, 0, 32);
        return -1;
    }
    status = psa_export_public_key(key_id, pubkey, 32, &pub_len);
    psa_destroy_key(key_id);
    if (status != PSA_SUCCESS || pub_len != 32) {
        ESP_LOGE(KD_TAG, "x25519 psa_export_public failed: %d", (int)status);
        memset(privkey, 0, 32);
        return -1;
    }
    return 0;
}

/* --- ML-DSA-65 --- */

int derive_ml_dsa_65_key(const uint8_t seed[64], uint32_t index,
                         uint8_t *pk, uint8_t *sk) {
    /* m/44'/102003'/<index>'/0'/0' — all hardened (SLIP-0010) -> 32-byte seed */
    const uint32_t path[5] = {
        44u | BIP32_HARDENED_FLAG,
        102003u | BIP32_HARDENED_FLAG,
        index | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
    };
    uint8_t pq_seed[32];

    if (seed == NULL || pk == NULL || sk == NULL) {
        return -1;
    }

    if (slip10_derive_seed(seed, path, 5, pq_seed) != 0) {
        return -1;
    }

    int ret = fw_pq_ml_dsa_65_keygen(pq_seed, pk, sk);
    memset(pq_seed, 0, sizeof(pq_seed));
    return ret;
}

/* --- SLH-DSA-128s --- */

int derive_slh_dsa_128s_key(const uint8_t seed[64], uint32_t index,
                            uint8_t *pk, uint8_t *sk) {
    /* m/44'/102004'/<index>'/0'/0' — all hardened (SLIP-0010) -> 32-byte seed */
    const uint32_t path[5] = {
        44u | BIP32_HARDENED_FLAG,
        102004u | BIP32_HARDENED_FLAG,
        index | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
    };
    uint8_t pq_seed[32];

    if (seed == NULL || pk == NULL || sk == NULL) {
        return -1;
    }

    if (slip10_derive_seed(seed, path, 5, pq_seed) != 0) {
        return -1;
    }

    ESP_LOGW(KD_TAG, "SLH-DSA-128s keygen: this takes 5-30 seconds on ESP32");
    int ret = fw_pq_slh_dsa_128s_keygen(pq_seed, pk, sk);
    memset(pq_seed, 0, sizeof(pq_seed));
    return ret;
}

/* --- ML-KEM-768 --- */

int derive_ml_kem_768_key(const uint8_t seed[64], uint32_t index,
                          uint8_t *pk, uint8_t *sk) {
    /* m/44'/102005'/<index>'/0'/0' — all hardened (SLIP-0010) -> 32-byte seed */
    const uint32_t path[5] = {
        44u | BIP32_HARDENED_FLAG,
        102005u | BIP32_HARDENED_FLAG,
        index | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
        0u | BIP32_HARDENED_FLAG,
    };
    uint8_t pq_seed[32];

    if (seed == NULL || pk == NULL || sk == NULL) {
        return -1;
    }

    if (slip10_derive_seed(seed, path, 5, pq_seed) != 0) {
        return -1;
    }

    int ret = fw_pq_ml_kem_768_keygen(pq_seed, pk, sk);
    memset(pq_seed, 0, sizeof(pq_seed));
    return ret;
}