/* pq_drbg.c — Deterministic PRNG for post-quantum key generation.
 *
 * Implements a SHAKE-256-based deterministic PRNG that replaces PQClean's
 * randombytes() function. This allows deterministic PQ key generation from
 * a mnemonic-derived seed: same seed → same randombytes output sequence.
 *
 * The PRNG is simple: 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. This is not a NIST SP 800-90A
 * compliant DRBG, but it is deterministic and sufficient for PQ keygen
 * (which only needs the output to be uniformly distributed, which SHAKE
 * provides).
 *
 * Security argument: SHAKE-256 is a XOF (extendable output function) based
 * on Keccak. Given a 256-bit seed, the output is computationally
 * indistinguishable from random. The counter prevents reuse of output
 * blocks. This is the same approach used by PQM4 for deterministic test
 * vectors.
 */
#include <string.h>
#include <openssl/evp.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, but we use 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) {
    EVP_MD_CTX *ctx;
    unsigned char seed_block[32 + 8]; /* seed + counter (8 bytes LE) */
    size_t outlen = sizeof(g_buffer);

    /* Build input: seed || counter (little-endian) */
    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);

    ctx = EVP_MD_CTX_new();
    if (ctx == NULL) {
        return;
    }

    if (EVP_DigestInit_ex(ctx, EVP_shake256(), NULL) != 1) {
        EVP_MD_CTX_free(ctx);
        return;
    }
    if (EVP_DigestUpdate(ctx, seed_block, (size_t)g_seed_len + 8) != 1) {
        EVP_MD_CTX_free(ctx);
        return;
    }
    if (EVP_DigestFinalXOF(ctx, g_buffer, outlen) != 1) {
        EVP_MD_CTX_free(ctx);
        return;
    }

    EVP_MD_CTX_free(ctx);

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

    /* Copy seed (up to 32 bytes; if shorter, zero-pad) */
    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;
}

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

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

/* --- randombytes() — the function PQClean calls --- */

int randombytes(unsigned char *buf, size_t len) {
    return pq_drbg_randombytes(buf, len);
}
