#include "rand.h"
#include "ch32fun.h"
#include "ch32v003hw.h"
#include <stdint.h>

#define BUILD_SEED                                                             \
    ((uint64_t)(__TIME__[0]) * (uint64_t)(__TIME__[1]) *                       \
     (uint64_t)(__TIME__[3]) * (uint64_t)(__TIME__[4]) *                       \
     (uint64_t)(__TIME__[6]) * (uint64_t)(__TIME__[7]))

#define FLASH_SEED_ADDR ((uintptr_t *)0x08003700) // PRNG state storage
#define PRNG_SAVE_INTERVAL 50 // Save every 1000 calls to prand()

static uint64_t seed = BUILD_SEED;

// Initialize this to something close to interval
static int prand_counter = PRNG_SAVE_INTERVAL - 10;

uint64_t prand() {
    seed = seed * 6364136223846793005ULL + 1;

    if (++prand_counter >= PRNG_SAVE_INTERVAL) {
        rand_save_to_flash();
        prand_counter = 0;
    }

    return seed;
}

uint64_t prand_range(uint64_t min, uint64_t max) {
    return min + (prand() % (max - min + 1));
}

void sprand(uint64_t s) {
    if (s) {
        seed = s;
    } else {
        rand_reseed();
    }
}

void rand_reseed() {
    uint64_t stored_seed = *(volatile uint64_t *)FLASH_SEED_ADDR;

    if (stored_seed == 0 || stored_seed == 0xFFFFFFFFFFFFFFFFULL) {
        seed = BUILD_SEED;
    } else {
        seed = stored_seed;
    }
}

// See:
// https://github.com/cnlohr/ch32v003fun/blob/2ac62072272f2ccd2122e688a9e0566de3976a94/examples/flashtest/flashtest.c
void rand_save_to_flash() {
    FLASH->KEYR = FLASH_KEY1; // Unlock flash
    FLASH->KEYR = FLASH_KEY2;

    FLASH->MODEKEYR = FLASH_KEY1; // Unlock programming mode
    FLASH->MODEKEYR = FLASH_KEY2;

    // Erase the flash page
    FLASH->CTLR = CR_PAGE_ER;
    FLASH->ADDR = (intptr_t)FLASH_SEED_ADDR;
    FLASH->CTLR = CR_STRT_Set | CR_PAGE_ER;
    while (FLASH->STATR & FLASH_STATR_BSY); // Wait for erase

    // Write new seed
    FLASH->CTLR = CR_PAGE_PG;
    FLASH->CTLR = CR_BUF_RST | CR_PAGE_PG;
    FLASH->ADDR = (intptr_t)FLASH_SEED_ADDR;

    ((uint32_t *)FLASH_SEED_ADDR)[0] = (uint32_t)seed;
    ((uint32_t *)FLASH_SEED_ADDR)[1] = (uint32_t)(seed >> 32);

    FLASH->CTLR = CR_PAGE_PG | FLASH_CTLR_BUF_LOAD;
    while (FLASH->STATR & FLASH_STATR_BSY); // Wait for completion

    FLASH->CTLR = CR_PAGE_PG | CR_STRT_Set; // Commit write
    while (FLASH->STATR & FLASH_STATR_BSY); // Wait for completion
}