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

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

#if defined(AURORA_PLATFORM_WIN32)
    #include <wincrypt.h>
#endif

namespace Aurora::RNG
{
#if (defined(AURORA_PLATFORM_LINUX) || defined(AURORA_PLATFORM_ANDRIOD) || defined(AURORA_PLATFORM_BSD) || defined(VENDOR_CONSOLE_SONY) || defined(VENDOR_CONSOLE_NINTENDO))
    #define _UNIX_LIKE_DEVRAND
#endif

#if defined(_UNIX_LIKE_DEVRAND)
    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)
        {
            fclose(gDevURand);
            gDevURand = NULL;
        }
    }

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

#elif defined(AURORA_PLATFORM_WIN32)
    
    #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 unsigned long RngWin32(AuUInt8 *buf, unsigned long 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 unsigned long RngStdC(AuUInt8 *buf, unsigned long 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 unsigned long RngGetBytes(AuUInt8 *out, unsigned long outlen)
    {
        unsigned long x;
    #if defined(AURORA_PLATFORM_WIN32)
        x = RngWin32(out, outlen);
        if (x != 0)
        {
            return x;
        }
    #elif defined(_UNIX_LIKE_DEVRAND)
        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, AuUInt length)
    {
        auto bytes = RngGetBytes(reinterpret_cast<AuUInt8 *>(in), length);
        SysAssertExp(bytes == length, "Couldn't consume {} RNG bytes", length);
    }

    void Init()
    {
        InitRandPlatform();
    }

    void Release()
    {
    #if defined(_UNIX_LIKE_DEVRAND)
        fclose(gDevURand);
    #elif defined(AURORA_PLATFORM_WIN32) && defined(USE_OLD_NTCRYPT)
        CryptReleaseContext(gCryptoProv, 0);
    #endif
    }
}