/* pq_crypto_firmware.c — Post-quantum crypto wrappers for ESP32 firmware.
 *
 * Wraps the PQClean algorithm API (via the pqclean component) with
 * firmware-friendly functions that handle the deterministic DRBG setup
 * for keygen and provide clean sign/verify/encaps/decaps interfaces.
 *
 * The PQ key buffers are large (ML-DSA-65 priv = 4032 bytes, SLH-DSA-128s
 * sig = 7856 bytes). Callers must allocate these on the heap or as static
 * buffers — stack allocation on ESP32 (8KB task stack default) will overflow
 * for the larger buffers.
 */
#include "pq_crypto_firmware.h"
#include "pqclean.h"

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

/* --- Key generation (deterministic from seed) --- */

int fw_pq_ml_dsa_65_keygen(const uint8_t seed[32],
                            uint8_t *pk, uint8_t *sk) {
    if (seed == NULL || pk == NULL || sk == NULL) {
        return -1;
    }

    /* Initialize the deterministic DRBG with the mnemonic-derived seed */
    pq_drbg_init(seed, 32);

    /* Run PQClean keygen — randombytes() draws from the DRBG */
    int ret = crypto_sign_keypair(pk, sk);

    /* Wipe the DRBG state — the seed material is sensitive */
    pq_drbg_zeroize();

    return ret;
}

int fw_pq_slh_dsa_128s_keygen(const uint8_t seed[32],
                               uint8_t *pk, uint8_t *sk) {
    if (seed == NULL || pk == NULL || sk == NULL) {
        return -1;
    }

    pq_drbg_init(seed, 32);

    /* WARNING: This call takes 5-30 seconds on ESP32 due to the
     * hypertree construction (7 layers of WOTS+ + Merkle trees). */
    int ret = slh_dsa_128s_crypto_sign_keypair(pk, sk);

    pq_drbg_zeroize();

    return ret;
}

int fw_pq_ml_kem_768_keygen(const uint8_t seed[32],
                             uint8_t *pk, uint8_t *sk) {
    if (seed == NULL || pk == NULL || sk == NULL) {
        return -1;
    }

    pq_drbg_init(seed, 32);

    int ret = crypto_kem_keypair(pk, sk);

    pq_drbg_zeroize();

    return ret;
}

/* --- Signing --- */

int fw_pq_ml_dsa_65_sign(uint8_t *sig, size_t *siglen,
                          const uint8_t *m, size_t mlen,
                          const uint8_t *sk) {
    if (sig == NULL || siglen == NULL || m == NULL || sk == NULL) {
        return -1;
    }
    return crypto_sign(sig, siglen, m, mlen, sk);
}

int fw_pq_ml_dsa_65_verify(const uint8_t *sig, size_t siglen,
                            const uint8_t *m, size_t mlen,
                            const uint8_t *pk) {
    if (sig == NULL || m == NULL || pk == NULL) {
        return -1;
    }
    /* crypto_sign_open expects (m_out, mlen_out, sm, smlen, pk) where sm
     * is the signed message. For detached signatures we reconstruct: the
     * PQClean API uses crypto_sign_open with sm = sig || m. */
    /* For firmware use, we provide a simple verify by re-signing is not
     * possible (non-deterministic). The PQClean crypto_sign_open expects
     * the concatenated sig||msg format. Callers should use the PQClean
     * API directly for verification, or we build the sm buffer here. */
    uint8_t *sm = (uint8_t *)malloc(siglen + mlen);
    if (sm == NULL) {
        return -1;
    }
    memcpy(sm, sig, siglen);
    memcpy(sm + siglen, m, mlen);

    uint8_t *m_out = (uint8_t *)malloc(mlen);
    if (m_out == NULL) {
        free(sm);
        return -1;
    }
    size_t mlen_out = 0;

    int ret = crypto_sign_open(m_out, &mlen_out, sm, siglen + mlen, pk);

    free(sm);
    free(m_out);

    return ret;
}

int fw_pq_slh_dsa_128s_sign(uint8_t *sig, size_t *siglen,
                              const uint8_t *m, size_t mlen,
                              const uint8_t *sk) {
    if (sig == NULL || siglen == NULL || m == NULL || sk == NULL) {
        return -1;
    }
    /* WARNING: This call takes 5-30 seconds on ESP32. */
    return slh_dsa_128s_crypto_sign(sig, siglen, m, mlen, sk);
}

int fw_pq_slh_dsa_128s_verify(const uint8_t *sig, size_t siglen,
                                const uint8_t *m, size_t mlen,
                                const uint8_t *pk) {
    if (sig == NULL || m == NULL || pk == NULL) {
        return -1;
    }
    uint8_t *sm = (uint8_t *)malloc(siglen + mlen);
    if (sm == NULL) {
        return -1;
    }
    memcpy(sm, sig, siglen);
    memcpy(sm + siglen, m, mlen);

    uint8_t *m_out = (uint8_t *)malloc(mlen);
    if (m_out == NULL) {
        free(sm);
        return -1;
    }
    size_t mlen_out = 0;

    int ret = slh_dsa_128s_crypto_sign_open(m_out, &mlen_out, sm,
                                             siglen + mlen, pk);

    free(sm);
    free(m_out);

    return ret;
}

/* --- KEM --- */

int fw_pq_ml_kem_768_encaps(uint8_t *ct, uint8_t *ss, const uint8_t *pk) {
    if (ct == NULL || ss == NULL || pk == NULL) {
        return -1;
    }
    /* encaps uses real randomness (hardware RNG) — the DRBG is not
     * initialized, so randombytes() falls back to esp_fill_random(). */
    return crypto_kem_enc(ct, ss, pk);
}

int fw_pq_ml_kem_768_decaps(uint8_t *ss, const uint8_t *ct, const uint8_t *sk) {
    if (ss == NULL || ct == NULL || sk == NULL) {
        return -1;
    }
    return crypto_kem_dec(ss, ct, sk);
}
