#include "random.h"
#include "fs/vfs.h"
#include "libk/stdio.h"
#include "arch/x86_64/sys/tsc.h"
#include "arch/x86_64/sys/pit.h"
#include "errno.h"
#include <stdbool.h>
#include <string.h>
#include "limine.h"

/* ChaCha20 core (public-domain, 20-round) */
static inline void chacha20_quarterround(uint32_t *a, uint32_t *b, uint32_t *c, uint32_t *d) {
    *a += *b; *d ^= *a; *d = (*d << 16) | (*d >> 16);
    *c += *d; *b ^= *c; *b = (*b << 12) | (*b >> 20);
    *a += *b; *d ^= *a; *d = (*d << 8) | (*d >> 24);
    *c += *d; *b ^= *c; *b = (*b << 7) | (*b >> 25);
}

static void chacha20_block(uint32_t state[16], uint8_t output[64]) {
    uint32_t x[16];
    for (int i = 0; i < 16; ++i) x[i] = state[i];

    for (int i = 0; i < 10; ++i) {  /* 10 double-rounds = 20 rounds */
        /* column rounds */
        chacha20_quarterround(&x[0], &x[4], &x[8], &x[12]);
        chacha20_quarterround(&x[1], &x[5], &x[9], &x[13]);
        chacha20_quarterround(&x[2], &x[6], &x[10], &x[14]);
        chacha20_quarterround(&x[3], &x[7], &x[11], &x[15]);
        /* diagonal rounds */
        chacha20_quarterround(&x[0], &x[5], &x[10], &x[15]);
        chacha20_quarterround(&x[1], &x[6], &x[11], &x[12]);
        chacha20_quarterround(&x[2], &x[7], &x[8], &x[13]);
        chacha20_quarterround(&x[3], &x[4], &x[9], &x[14]);
    }

    for (int i = 0; i < 16; ++i) {
        x[i] += state[i];
        ((uint32_t*)output)[i] = x[i];
    }
}

/* Global ChaCha20 PRNG state */
static uint8_t rng_key[32];
static uint64_t rng_counter = 0;
static bool rng_seeded = false;

static void rng_generate(uint8_t *buf, size_t len) {
    size_t i = 0;
    while (len > 0) {
        uint32_t state[16];
        static const uint32_t sigma[4] = {0x61707865, 0x3320646e, 0x79622d32, 0x6b206574};

        state[0] = sigma[0]; state[1] = sigma[1]; state[2] = sigma[2]; state[3] = sigma[3];
        for (int j = 0; j < 8; ++j)
            state[4 + j] = ((uint32_t*)rng_key)[j];

        state[12] = (uint32_t)rng_counter;
        state[13] = (uint32_t)(rng_counter >> 32);
        state[14] = 0xdeadbeef;   /* fixed nonce (PRNG only) */
        state[15] = 0xbeefdead;

        uint8_t block[64];
        chacha20_block(state, block);

        size_t copy = (len < 64) ? len : 64;
        for (size_t k = 0; k < copy; ++k)
            buf[i + k] = block[k];

        i += copy;
        len -= copy;
        rng_counter++;
    }
}

static void rng_add_entropy(const uint8_t *data, size_t len) {
    if (len == 0 || !data) return;

    for (size_t i = 0; i < len; ++i)
        rng_key[i % 32] ^= data[i];

    /* stir (forward secrecy) */
    uint8_t discard[64];
    rng_generate(discard, 64);

    rng_seeded = true;
}

int random_read(void* buf, size_t len) {
    if (len == 0) return 0;
    if (!buf) return -EFAULT;
    if (!rng_seeded) return -EAGAIN;   /* should never happen after init */

    rng_generate((uint8_t*)buf, len);
    return (int)len;
}

int random_write(const void* buf, size_t len) {
    if (len == 0) return 0;
    if (!buf) return -EFAULT;

    rng_add_entropy((const uint8_t*)buf, len);
    return (int)len;
}

int getrandom(void *buf, size_t buflen, unsigned int flags)
{
    if (!buf) return -EFAULT;
    if (buflen == 0) return 0;

    if (flags & GRND_INSECURE) {
        // maybe fall back to weaker RNG?
    }

    if (!rng_seeded) {
        if (flags & GRND_NONBLOCK)
            return -EAGAIN;
        // could block here in the future
        return -EAGAIN;
    }

    rng_generate((uint8_t*)buf, buflen);
    return (int)buflen;

}

void random_init(void) {
    printf("Initializing ChaCha20-based CSPRNG for /dev/random and /dev/urandom...\n");

    uint8_t entropy[64] = {0};
    size_t epos = 0;

    /* === SEEDING SOURCES === */

    /* 1. Limine boot timestamp (wall-clock) */
    extern volatile struct limine_date_at_boot_request boot_request;
    if (boot_request.response) {
        uint64_t ts = boot_request.response->timestamp;
        if (epos + sizeof(ts) <= sizeof(entropy)) {
            memcpy(entropy + epos, &ts, sizeof(ts));
            epos += sizeof(ts);
        }
    }

    /* 2. Multiple TSC samples (high-resolution jitter) */
    for (int i = 0; i < 4; ++i) {
        uint64_t tsc_val = rdtsc();
        if (epos + sizeof(tsc_val) <= sizeof(entropy)) {
            memcpy(entropy + epos, &tsc_val, sizeof(tsc_val));
            epos += sizeof(tsc_val);
        }
    }

    /* 3. PIT tick counter */
    uint64_t pit_ticks = PIT_GetTicks();
    if (epos + sizeof(pit_ticks) <= sizeof(entropy)) {
        memcpy(entropy + epos, &pit_ticks, sizeof(pit_ticks));
        epos += sizeof(pit_ticks);
    }

    /* 4. RDRAND (hardware TRNG) if present */
    {
        uint32_t eax, ebx, ecx, edx;
        asm volatile("cpuid" : "=a"(eax), "=b"(ebx), "=c"(ecx), "=d"(edx) : "a"(1));
        if (ecx & (1u << 30)) {
            printf("RDRAND detected - using hardware RNG for seeding\n");
            for (int i = 0; i < 4 && epos + 8 <= sizeof(entropy); ++i) {
                uint64_t rd;
                asm volatile("1: rdrand %0\n\tjnc 1b" : "=r"(rd));
                memcpy(entropy + epos, &rd, sizeof(rd));
                epos += sizeof(rd);
            }
        }
    }

    /* 5. Unix seconds from PIT (time-based entropy) */
    extern uint64_t g_Unixseconds;
    uint64_t unix_sec = g_Unixseconds;
    if (epos + sizeof(unix_sec) <= sizeof(entropy)) {
        memcpy(entropy + epos, &unix_sec, sizeof(unix_sec));
        epos += sizeof(unix_sec);
    }

    /* 6. Final TSC padding (extra jitter) */
    while (epos < 32) {
        uint64_t t = rdtsc();
        memcpy(entropy + epos, &t, (epos + 8 <= 32) ? 8 : (32 - epos));
        epos += (epos + 8 <= 32) ? 8 : (32 - epos);
    }

    /* Initialise ChaCha20 key from collected entropy */
    memcpy(rng_key, entropy, 32);
    rng_counter = rdtsc();   /* random starting counter */
    rng_seeded = true;

    printf("ChaCha20 RNG seeded with %zu bytes of entropy (TSC, PIT, RDRAND, boot time, unix time, etc.)\n", epos);

    /* Register both devices (open as "/dev/random" and "/dev/urandom") */
    VFS_RegisterCharDev("/dev/random", random_read, random_write);
    VFS_RegisterCharDev("/dev/urandom", random_read, random_write);

    printf("Registered chardevs: /dev/random and /dev/urandom (ChaCha20 CSPRNG)\n");
}