/***
    Copyright (C) 2021 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: RNG.cpp
    Date: 2021-6-11
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "RNG.hpp"

#if defined(AURORA_PLATFORM_WIN32)
    #include <wincrypt.h>
    #include <bcrypt.h>
#elif defined(AURORA_IS_MODERNNT_DERIVED)
    #include <bcrypt.h>
    #undef USE_OLD_NTCRYPT
#endif

#if defined(AURORA_IS_POSIX_DERIVED)
    #include <cstdio>
#endif

#include "WELL.hpp"

namespace Aurora::RNG
{
    static WELLRand gWellRand;
    static RandomUnique_t gFastDevice;

#if defined(AURORA_IS_POSIX_DERIVED)
    static FILE *gDevURand;

    static void InitRandPlatform()
    {
        gDevURand = fopen("/dev/urandom", "rb");
        if (gDevURand == NULL)
        {
            gDevURand = fopen("/dev/random", "rb");
        }

        if (!gDevURand)
        {
            return;
        }

        if (setvbuf(gDevURand, NULL, _IONBF, 0) != 0)
        {
        #if 0
            fclose(gDevURand);
            gDevURand = NULL;
        #endif
        }
    }

    static AuUInt32 RngUnix(AuUInt8 *buf, AuUInt32 len)
    {
        return fread(buf, 1, (size_t)len, gDevURand);
    }

#elif defined(AURORA_IS_MODERNNT_DERIVED)
    
    #if defined(USE_OLD_NTCRYPT)
    static HCRYPTPROV gCryptoProv;
    #endif

    static void InitRandPlatform()
    {
    #if defined(USE_OLD_NTCRYPT)
        if (!CryptAcquireContext(&gCryptoProv, NULL, MS_DEF_PROV, PROV_RSA_FULL,
            (CRYPT_VERIFYCONTEXT | CRYPT_MACHINE_KEYSET)))
        {
            if (!CryptAcquireContext(&gCryptoProv, NULL, MS_DEF_PROV, PROV_RSA_FULL,
                CRYPT_VERIFYCONTEXT | CRYPT_MACHINE_KEYSET | CRYPT_NEWKEYSET))
            {
                gCryptoProv = 0;
            }
        }
    #endif
    }

    static AuUInt32 RngWin32(AuUInt8 *buf, AuUInt32 len)
    {
    #if defined(USE_OLD_NTCRYPT)
        if (!gCryptoProv)
        {
            return 0;
        }

        if (!CryptGenRandom(gCryptoProv, len, buf))
        {
            return 0;
        }

        return len;
    #else
        
        #if !defined(AURORA_DONT_PREFER_WIN32_USERLAND_AES_RNG)
            if (BCryptGenRandom(BCRYPT_RNG_ALG_HANDLE, reinterpret_cast<PUCHAR>(buf), len, 0) == 0)
            {
                return len;
            }
        #endif

        if (BCryptGenRandom(NULL, reinterpret_cast<PUCHAR>(buf), len, BCRYPT_USE_SYSTEM_PREFERRED_RNG) == 0)
        {
            return len;
        }

        return 0;
    #endif
    }

#else

    static void InitRandPlatform()
    {
    }

#endif

    static AuUInt32 RngStdC(AuUInt8 *buf, AuUInt32 len)
    {
        clock_t t1;
        int l, acc, bits, a, b;

        l = len;
        acc = a = b = 0;
        bits = 8;

        while (len--)
        {
            while (bits--) // for each bit in byte
            {
                do
                {
                    t1 = clock();
                    while (t1 == clock()) // spin within the resolution of 1 C clock() tick
                    {
                        a ^= 1; // flip
                    }

                    t1 = clock();
                    while (t1 == clock())
                    {
                        b ^= 1; // flip
                    }
                } while (a == b); // ensure theres enough entropy for a deviation to occur
                acc = (acc << 1) | a; // push the first bit state
            }

            *buf++ = acc;

            acc = 0;
            bits = 8;
        }

        return l;
    }

    static AuUInt32 RngGetBytes(AuUInt8 *out, AuUInt32 outlen)
    {
        AuUInt32 x;
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        x = RngWin32(out, outlen);
        if (x != 0)
        {
            return x;
        }
    #elif defined(AURORA_IS_POSIX_DERIVED)
        x = RngUnix(out, outlen);
        if (x != 0)
        {
            return x;
        }
    #endif
        x = RngStdC(out, outlen);
        if (x != 0)
        {
            return x;
        }
        return 0;
    }

    AUKN_SYM void ReadSecureRNG(void *in, AuUInt32 length)
    {
        AuUInt32 offset;
        AuUInt8 *headPtr;
        
        headPtr = reinterpret_cast<AuUInt8 *>(in);
        offset  = 0;

        while (offset != length)
        {
            auto req   = length - offset;
            auto bytes = RngGetBytes(headPtr + offset, req);
            SysAssertExp(bytes, "Couldn't consume {} RNG bytes", req);
            offset += bytes;
        }
    }
    
    AUKN_SYM void ReadFastRNG(void *in, AuUInt32 length)
    {
        WELL_NextBytes(&gWellRand, in, length);
    }

    AUKN_SYM AuString ReadString(AuUInt32 length, ERngStringCharacters type)
    {
        return gFastDevice->NextString(length, type);
    }

    AUKN_SYM void     RngString(char *string, AuUInt32 length, ERngStringCharacters type)
    {
        gFastDevice->NextString(string, length, type);
    }

    AUKN_SYM AuUInt8  RngByte()
    {
        return gFastDevice->NextByte();
    }

    AUKN_SYM bool     RngBoolean()
    {
        return gFastDevice->NextBoolean();
    }

    AUKN_SYM AuUInt32 RngU32()
    {
        return gFastDevice->NextU32();
    }

    AUKN_SYM AuUInt32 RngU32(AuUInt32 min, AuUInt32 max)
    {
        return gFastDevice->NextU32(min, max);
    }

    AUKN_SYM AuUInt64 RngU64()
    {
        return gFastDevice->NextU64();
    }

    AUKN_SYM AuInt32  RngInt(AuInt32 min, AuInt32 max)
    {
        return gFastDevice->NextInt(min, max);
    }

    AUKN_SYM double   RngDecimal()
    {
        return gFastDevice->NextDecimal();
    }

    AUKN_SYM float    RngNumber(float min, float max)
    {
        return gFastDevice->NextNumber(min, max);
    }

    AUKN_SYM AuUInt32 RngIndex(AuUInt32 count /* = max + 1*/)
    {
        return gFastDevice->NextIndex(count);
    }

    static void InitFastRng()
    {
        AuArray<AuUInt8, 64> maxEntropy;
        RngArray<false, AuUInt8>(maxEntropy.data(), 64);
        gWellRand = WELL_SeedRandBig64(maxEntropy);
        gFastDevice = RandomUnique(RandomDef {false});
    }

    void Init()
    {
        InitRandPlatform();
        InitFastRng();
    }

    void Release()
    {
    #if defined(AURORA_IS_POSIX_DERIVED)
        if (gDevURand != NULL)
        fclose(gDevURand);
    #elif defined(AURORA_IS_MODERNNT_DERIVED) && defined(USE_OLD_NTCRYPT)
        CryptReleaseContext(gCryptoProv, 0);
    #endif
    }
}