/* pq_drbg_firmware.c — Deterministic PRNG for PQ key generation on ESP32.
 *
 * Same algorithm as the host's src/pq_drbg.c but uses the crypto backend
 * abstraction (which resolves to mbedtls on ESP32) for SHAKE-256 instead
 * of OpenSSL EVP. This allows deterministic PQ key generation from a
 * mnemonic-derived seed: same seed -> same randombytes output sequence.
 *
 * The PRNG: SHAKE-256(seed || counter) produces a stream of pseudo-random
 * bytes. The counter is a 64-bit little-endian integer that increments
 * each time we need more output.
 */
#include <string.h>
#include <stdlib.h>
#include "crypto_backend.h"

/* --- DRBG state --- */

static unsigned char g_seed[32];
static int g_seed_len = 0;
static uint64_t g_counter = 0;
static unsigned char g_buffer[168]; /* SHAKE-256 rate = 136, 168 for safety */
static size_t g_buffer_pos = sizeof(g_buffer);
static int g_initialized = 0;

/* --- internal: squeeze more bytes from SHAKE-256 --- */

static void drbg_refill(void) {
    unsigned char seed_block[32 + 8]; /* seed + counter (8 bytes LE) */

    memcpy(seed_block, g_seed, (size_t)g_seed_len);
    seed_block[g_seed_len + 0] = (unsigned char)(g_counter & 0xFF);
    seed_block[g_seed_len + 1] = (unsigned char)((g_counter >> 8) & 0xFF);
    seed_block[g_seed_len + 2] = (unsigned char)((g_counter >> 16) & 0xFF);
    seed_block[g_seed_len + 3] = (unsigned char)((g_counter >> 24) & 0xFF);
    seed_block[g_seed_len + 4] = (unsigned char)((g_counter >> 32) & 0xFF);
    seed_block[g_seed_len + 5] = (unsigned char)((g_counter >> 40) & 0xFF);
    seed_block[g_seed_len + 6] = (unsigned char)((g_counter >> 48) & 0xFF);
    seed_block[g_seed_len + 7] = (unsigned char)((g_counter >> 56) & 0xFF);

    crypto_backend_shake256(seed_block, (size_t)g_seed_len + 8,
                            g_buffer, sizeof(g_buffer));

    g_counter++;
    g_buffer_pos = 0;
}

/* --- public API --- */

void pq_drbg_init(const unsigned char *seed, size_t seed_len) {
    if (seed == NULL || seed_len == 0) {
        return;
    }

    memset(g_seed, 0, sizeof(g_seed));
    if (seed_len > sizeof(g_seed)) {
        seed_len = sizeof(g_seed);
    }
    memcpy(g_seed, seed, seed_len);
    g_seed_len = (int)sizeof(g_seed); /* always use 32-byte seed (zero-padded) */

    g_counter = 0;
    g_buffer_pos = sizeof(g_buffer);
    g_initialized = 1;
}

void pq_drbg_zeroize(void) {
    crypto_backend_cleanse(g_seed, sizeof(g_seed));
    crypto_backend_cleanse(g_buffer, sizeof(g_buffer));
    g_seed_len = 0;
    g_counter = 0;
    g_buffer_pos = sizeof(g_buffer);
    g_initialized = 0;
}

/* Returns 1 if the DRBG has been initialized (keygen mode), 0 otherwise.
 * Used by randombytes_mbedtls.c to decide between deterministic DRBG and
 * hardware RNG. */
int pq_drbg_is_initialized(void) {
    return g_initialized;
}

/* pq_drbg_randombytes is called by randombytes() below. */
int pq_drbg_randombytes(unsigned char *buf, size_t len) {
    if (buf == NULL || !g_initialized) {
        return -1;
    }

    while (len > 0) {
        size_t avail;
        size_t to_copy;

        if (g_buffer_pos >= sizeof(g_buffer)) {
            drbg_refill();
            if (g_buffer_pos >= sizeof(g_buffer)) {
                return -1; /* refill failed */
            }
        }

        avail = sizeof(g_buffer) - g_buffer_pos;
        to_copy = (len < avail) ? len : avail;
        memcpy(buf, g_buffer + g_buffer_pos, to_copy);
        g_buffer_pos += to_copy;
        buf += to_copy;
        len -= to_copy;
    }

    return 0;
}
