#include "efuse_ecdsa.h"

#include "sdkconfig.h"
#include <esp_efuse.h>
#include <esp_log.h>
#include <mbedtls/ecdsa.h>
#include <mbedtls/ecp.h>
#include <stdbool.h>
#include <stdint.h>
#include <string.h>

#include "ecdsa/ecdsa_alt.h"

static const char *TAG = "efuse_ecdsa";

/* We use EFUSE_BLK_KEY1 (block 5) for the ECDSA key, leaving KEY0 free
 * for secure boot / TEE keys. */
#define ECDSA_EFUSE_BLOCK  EFUSE_BLK_KEY1

/* -------------------------------------------------------------------------
 * Public API
 * ---------------------------------------------------------------------- */

bool efuse_ecdsa_key_provisioned(void)
{
    esp_efuse_block_t blk;
    return esp_efuse_find_purpose(ESP_EFUSE_KEY_PURPOSE_ECDSA_KEY, &blk);
}

bool efuse_ecdsa_provision_key(const uint8_t key[32])
{
    if (efuse_ecdsa_key_provisioned()) {
        ESP_LOGW(TAG, "ECDSA key already provisioned – skipping");
        return true;
    }

    if (!esp_efuse_key_block_unused(ECDSA_EFUSE_BLOCK)) {
        ESP_LOGE(TAG, "eFuse key block %d is already in use", ECDSA_EFUSE_BLOCK);
        return false;
    }

    esp_err_t err = esp_efuse_write_key(ECDSA_EFUSE_BLOCK,
                                        ESP_EFUSE_KEY_PURPOSE_ECDSA_KEY,
                                        key, 32);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "esp_efuse_write_key failed: %s", esp_err_to_name(err));
        return false;
    }

    ESP_LOGI(TAG, "ECDSA P-256 key burned into eFuse block %d", ECDSA_EFUSE_BLOCK);
    return true;
}

bool efuse_ecdsa_get_pubkey(uint8_t pub_x[32], uint8_t pub_y[32])
{
    esp_efuse_block_t blk;
    if (!esp_efuse_find_purpose(ESP_EFUSE_KEY_PURPOSE_ECDSA_KEY, &blk)) {
        ESP_LOGE(TAG, "No eFuse block with ECDSA_KEY purpose found");
        return false;
    }

    esp_ecdsa_pk_conf_t conf = {
        .grp_id     = MBEDTLS_ECP_DP_SECP256R1,
        .efuse_block = blk,
        .load_pubkey = true,
    };

    mbedtls_pk_context pk;
    if (esp_ecdsa_set_pk_context(&pk, &conf) != 0) {
        ESP_LOGE(TAG, "esp_ecdsa_set_pk_context failed");
        return false;
    }

    mbedtls_ecp_keypair *kp = mbedtls_pk_ec(pk);
    int ret = 0;
    ret |= mbedtls_mpi_write_binary(&kp->MBEDTLS_PRIVATE(Q).MBEDTLS_PRIVATE(X), pub_x, 32);
    ret |= mbedtls_mpi_write_binary(&kp->MBEDTLS_PRIVATE(Q).MBEDTLS_PRIVATE(Y), pub_y, 32);

    mbedtls_pk_free(&pk);

    if (ret != 0) {
        ESP_LOGE(TAG, "Failed to export public key coordinates");
        return false;
    }
    return true;
}

bool efuse_ecdsa_sign(const uint8_t digest[32],
                      uint8_t r_out[32], uint8_t s_out[32])
{
    esp_efuse_block_t blk;
    if (!esp_efuse_find_purpose(ESP_EFUSE_KEY_PURPOSE_ECDSA_KEY, &blk)) {
        ESP_LOGE(TAG, "No eFuse block with ECDSA_KEY purpose found");
        return false;
    }

    mbedtls_ecdsa_context ctx;
    mbedtls_ecdsa_init(&ctx);
    mbedtls_ecp_group_load(&ctx.MBEDTLS_PRIVATE(grp), MBEDTLS_ECP_DP_SECP256R1);

    mbedtls_mpi key_mpi;
    esp_ecdsa_privkey_load_mpi(&key_mpi, blk);

    mbedtls_mpi r, s;
    mbedtls_mpi_init(&r);
    mbedtls_mpi_init(&s);

    int ret = mbedtls_ecdsa_sign(&ctx.MBEDTLS_PRIVATE(grp),
                                 &r, &s, &key_mpi,
                                 digest, 32, NULL, NULL);
    if (ret != 0) {
        ESP_LOGE(TAG, "mbedtls_ecdsa_sign failed: -0x%04X", (unsigned)-ret);
        goto out;
    }

    ret |= mbedtls_mpi_write_binary(&r, r_out, 32);
    ret |= mbedtls_mpi_write_binary(&s, s_out, 32);

out:
    mbedtls_mpi_free(&r);
    mbedtls_mpi_free(&s);
    mbedtls_mpi_free(&key_mpi);
    mbedtls_ecdsa_free(&ctx);
    return ret == 0;
}