// SPDX-License-Identifier: GPL-3.0-or-later
#pragma once

#include <Arduino.h>
#include <mbedtls/ecdh.h>
#include <mbedtls/entropy.h>
#include <mbedtls/ctr_drbg.h>
#include <mbedtls/md.h>
#include <mbedtls/sha256.h>

#define ECDH_PUBKEY_LEN  65
#define AES_KEY_LEN      32
#define PSK_LEN          32
#define AES_PKT_SIZE     29
#define ECDH_PKT_SIZE    97      // pubkey(65) + HMAC(32)

class CryptoECDH {
private:
    mbedtls_ecdh_context ecdh;
    mbedtls_entropy_context entropy;
    mbedtls_ctr_drbg_context ctr_drbg;
    bool initialized = false;

    CryptoECDH(const CryptoECDH&) = delete;
    CryptoECDH& operator=(const CryptoECDH&) = delete;

public:
    CryptoECDH() = default;

    uint8_t public_key[ECDH_PUBKEY_LEN];
    size_t pub_len = 0;
    uint8_t shared_secret[32];
    uint8_t aes_key[AES_KEY_LEN];

    bool init() {
        mbedtls_ecdh_init(&ecdh);
        mbedtls_entropy_init(&entropy);
        mbedtls_ctr_drbg_init(&ctr_drbg);

        const char *pers = "sylda_ecdh";
        int ret = mbedtls_ctr_drbg_seed(&ctr_drbg, mbedtls_entropy_func, &entropy,
                                        (const unsigned char *)pers, strlen(pers));
        if (ret != 0) {
            mbedtls_ecdh_free(&ecdh);
            mbedtls_entropy_free(&entropy);
            mbedtls_ctr_drbg_free(&ctr_drbg);
            return false;
        }

        ret = mbedtls_ecp_group_load(&ecdh.grp, MBEDTLS_ECP_DP_SECP256R1);
        if (ret != 0) {
            mbedtls_ecdh_free(&ecdh);
            mbedtls_entropy_free(&entropy);
            mbedtls_ctr_drbg_free(&ctr_drbg);
            return false;
        }

        initialized = true;
        return true;
    }

    bool generateKeys() {
        if (!initialized) return false;
        int ret = mbedtls_ecdh_gen_public(&ecdh.grp, &ecdh.d, &ecdh.Q,
                                          mbedtls_ctr_drbg_random, &ctr_drbg);
        if (ret != 0) return false;

        ret = mbedtls_ecp_point_write_binary(&ecdh.grp, &ecdh.Q,
                                             MBEDTLS_ECP_PF_UNCOMPRESSED,
                                             &pub_len, public_key, ECDH_PUBKEY_LEN);
        return (ret == 0 && pub_len == ECDH_PUBKEY_LEN);
    }

    bool computeSharedSecret(const uint8_t *peerPubKey, size_t peerPubLen) {
        if (!initialized || peerPubLen != ECDH_PUBKEY_LEN) return false;

        int ret = mbedtls_ecp_point_read_binary(&ecdh.grp, &ecdh.Qp,
                                                peerPubKey, peerPubLen);
        if (ret != 0) return false;

        size_t olen;
        ret = mbedtls_ecdh_compute_shared(&ecdh.grp, &ecdh.z, &ecdh.Qp, &ecdh.d,
                                          mbedtls_ctr_drbg_random, &ctr_drbg);
        if (ret != 0) return false;

        ret = mbedtls_mpi_write_binary(&ecdh.z, shared_secret, 32);
        if (ret != 0) return false;

        mbedtls_sha256(shared_secret, 32, aes_key, 0);

        return true;
    }

    static void signPubkey(const uint8_t *psk, const uint8_t *pubkey, uint8_t *hmacOut) {
        const mbedtls_md_info_t *mdInfo = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
        mbedtls_md_hmac(mdInfo, psk, PSK_LEN, pubkey, ECDH_PUBKEY_LEN, hmacOut);
    }

    static bool verifyPubkey(const uint8_t *psk, const uint8_t *pubkey, const uint8_t *hmac) {
        uint8_t computed[32];
        signPubkey(psk, pubkey, computed);
        return memcmp(computed, hmac, 32) == 0;
    }

    void free() {
        if (initialized) {
            mbedtls_ecdh_free(&ecdh);
            mbedtls_ctr_drbg_free(&ctr_drbg);
            mbedtls_entropy_free(&entropy);
            memset(shared_secret, 0, sizeof(shared_secret));
            memset(aes_key, 0, sizeof(aes_key));
            memset(public_key, 0, sizeof(public_key));
            initialized = false;
        }
    }
};