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

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

#if defined(AURORA_PLATFORM_WIN32)
    #include <bcrypt.h>

    #if !defined(BCRYPT_RNG_ALG_HANDLE)
        #define BCRYPT_RNG_ALG_HANDLE ((BCRYPT_ALG_HANDLE) 0x00000081)
    #endif

    #define AURORA_RNG_HAS_OLD_WINCRYPT
#endif

#if defined(AURORA_IS_MODERNNT_DERIVED)
    #include <bcrypt.h>
#endif

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

#if defined(AURORA_IS_POSIX_DERIVED)
    #include <cstdio>
    #include <fcntl.h>
#endif

#if defined(AURORA_IS_BSD_DERIVED)
    #include <stdlib.h>
    #define AURORA_RND_HAS_ARC4
#endif

#include "AuWELL.hpp"

namespace Aurora::RNG
{
#if defined(AURORA_IS_POSIX_DERIVED)
    static int gDevURand { -1 };

    void EntropyInit()
    {
        gDevURand = ::open("/dev/urandom", O_RDONLY | O_CLOEXEC);
        if (gDevURand <= 0)
        {
            gDevURand = ::open("/dev/random", O_RDONLY | O_CLOEXEC);
        }

        if (gDevURand <= 0)
        {
            gDevURand = -1;
        }
    }

    static AuUInt32 RngUnix(AuUInt8 *pBuf, AuUInt32 uLen)
    {
        if (gDevURand == -1)
        {
            return 0;
        }

        int iRet = ::read(gDevURand, pBuf, uLen);
        if (iRet <= 0)
        {
            return 0;
        }

        return AuUInt32(iRet);
    }

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

    void EntropyInit()
    {
    #if defined(AURORA_RNG_HAS_OLD_WINCRYPT)
        if (pBCryptGenRandom)
        {
            return;
        }

        if (pRtlGenRandom)
        {
            return;
        }

        if (pCryptAcquireContextW)
        {
            if (!pCryptAcquireContextW(&gCryptoProv, NULL,
                                       MS_DEF_PROV_W, PROV_RSA_FULL,
                                       (CRYPT_VERIFYCONTEXT | CRYPT_MACHINE_KEYSET)))
            {
                if (!pCryptAcquireContextW(&gCryptoProv, NULL,
                                           MS_DEF_PROV_W, PROV_RSA_FULL,
                                           CRYPT_VERIFYCONTEXT | CRYPT_MACHINE_KEYSET | CRYPT_NEWKEYSET))
                {
                    gCryptoProv = 0;
                }
            }
        }
    #endif
    }

    static AuUInt32 RngWin32(AuUInt8 *pBuf, AuUInt32 uLen)
    {
        if (pRtlGenRandom)
        {
            if (pRtlGenRandom(pBuf, uLen))
            {
                return uLen;
            }
        }

        if (pBCryptGenRandom)
        {
            if (AuSwInfo::IsWindows10OrGreater() || 
                AuBuild::kCurrentPlatform != AuBuild::EPlatform::ePlatformWin32)
            {
                if (pBCryptGenRandom(BCRYPT_RNG_ALG_HANDLE, reinterpret_cast<PUCHAR>(pBuf), uLen, 0) == 0)
                {
                    return uLen;
                }
            }

            if (pBCryptGenRandom(NULL, reinterpret_cast<PUCHAR>(pBuf), uLen, BCRYPT_USE_SYSTEM_PREFERRED_RNG) == 0)
            {
                return uLen;
            }
        }

    #if defined(AURORA_RNG_HAS_OLD_WINCRYPT)
        if (gCryptoProv)
        {
            if (pCryptGenRandom(gCryptoProv, uLen, pBuf))
            {
                return uLen;
            }
        }
    #endif

        return 0;
    }

#else

    void EntropyInit()
    {

    }

#endif

    static AuUInt64 RngTimeClock()
    {
    #if 0
        return clock(); // [...] resolution of 1 posix clock() tick (this is usually 1,000,000 == CLOCKS_PER_SEC per posix, but it can be 1k)
    #else
        static AuUInt64 gTimeHolder;
        static AuUInt64 gTime = (gTimeHolder = AuTime::SteadyClockJiffies()) >= 10'000'000ull ?
            1 :
            AuPageRoundUp<AuUInt64>((10'000'000ull / gTimeHolder), 1024);
        return AuTime::SteadyClockNS() * gTime / 10'000ull;
    #endif
    }

    static AuUInt32 RngTimeBased(AuUInt8 *pBuf, AuUInt32 uLen)
    {
        AuUInt64 t1;
        int acc, bits, a, b, c;

        acc = a = b = c = 0;
        bits = 8;

        void *pASLRSeed = (void *)&RngTimeBased;

        for (AU_ITERATE_N(uOffsetInByteStream, uLen))
        {
            for (AU_ITERATE_N(uMultiplePassesForTheFunOfIt, 3))
            {
                while (bits--) // for each bit in byte
                {
                    do
                    {
                        t1 = RngTimeClock();
                        while (t1 == RngTimeClock()) // spin within 1 microseconds (*)
                        {
                            a ^= 1; // flip
                        }

                        t1 = RngTimeClock();
                        while (t1 == RngTimeClock()) // spin within 1 microseconds
                        {
                            b ^= 1; // flip
                        }

                        // Update 2023/Sept: Now x10 because we dont have enough timer res on some platforms. 
                        // This is was generally never a CPU limitation or API limitation until Windows XP. Check software or motherboard...
                        // ...but still, let's continue.
                        // Deadlocking preprocess init, because for example the motherboard is failing or a VM configuration changed, is pretty dumb.
                        // (ask me how i know)
                        // Update oct: it scales with whatever frequency user-land clocks can be trusted with, per the respective kernels' HAL (viso, kshared user data, etc).
                    }
                    while (a == b); // ensure theres enough entropy for a deviation to occur
                    acc = (acc << 1) | a; // push the first bit state
                }

                c = AuHashCode(acc) ^
                    (AuHashCode(c) * kFnv1MagicPrime32) ^
                    (a ? AuHashCode(pASLRSeed) : 0);
                    
                acc = 0;
                bits = 8;
            }

            *pBuf++ = AuUInt8(c);
        }

        return uLen;
    }

    AuUInt32 RngGetBytes(AuUInt8 *pBuffer, AuUInt32 uBytes)
    {
        AuUInt32 x;
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        x = RngWin32(pBuffer, uBytes);
        if (x != 0)
        {
            return x;
        }
    #elif defined(AURORA_IS_POSIX_DERIVED)
        #if defined(AURORA_IS_LINUX_DERIVED)
            x = sys_getrandom(pBuffer, uBytes);
            if (x > 0)
            {
                return x;
            }
        #endif
        #if defined(AURORA_RND_HAS_ARC4)
            arc4random_buf(pBuffer, uBytes);
            return uBytes;
        #endif
        x = RngUnix(pBuffer, uBytes);
        if (x != 0)
        {
            return x;
        }
    #endif
        x = RngTimeBased(pBuffer, uBytes);
        if (x != 0)
        {
            return x;
        }
        return 0;
    }

    void EntropyDeinit()
    {
    #if defined(AURORA_IS_POSIX_DERIVED)
        if (gDevURand > 0)
        {
            ::close(gDevURand);
            gDevURand = -1;
        }
    #elif defined(AURORA_IS_MODERNNT_DERIVED) && defined(AURORA_RNG_HAS_OLD_WINCRYPT)
        if (pBCryptGenRandom)
        {
            return;
        }

        if (pRtlGenRandom)
        {
            return;
        }

        if (pCryptReleaseContext)
        {
            pCryptReleaseContext(gCryptoProv, 0);
        }
    #endif
    }
}