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

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

#if defined(AURORA_IS_BSD_DERIVED)
    #include <sys/types.h>
    #include <sys/sysctl.h>
#endif

#if defined(AURORA_IS_POSIX_DERIVED)
    #include <stdlib.h>
    #include <sys/sysinfo.h>
#endif

#if defined(AURORA_COMPILER_CLANG) || defined(AURORA_IS_POSIX_DERIVED)
    #include <cpuid.h>
#endif

namespace Aurora::HWInfo
{
    static CpuInfo gCpuInfo;

    union CPUIdContext
    {
        struct
        {
            AuUInt32 eax;
            AuUInt32 ebx;
            AuUInt32 ecx;
            AuUInt32 edx;
        };
        int regs[4];
    };

#if defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86)
    #if defined(AURORA_COMPILER_MSVC)
        static CPUIdContext cpuid(AuUInt32 a)
        {
            CPUIdContext context;
            __cpuid(context.regs, a);
            return context;
        }
    #elif defined(AURORA_COMPILER_CLANG) || defined(AURORA_COMPILER_GCC)
        static CPUIdContext cpuid(AuUInt32 a)
        {
            CPUIdContext context;
            __get_cpuid(a, &context.eax, &context.ebx, &context.ecx, &context.edx);
            return context;
        }
    #else
        static CPUIdContext cpuid(AuUInt32 a)
        {
            return {};
        }
    #endif
#else
    static CPUIdContext cpuid(AuUInt32 a)
    {
        return {};
    }
#endif

    bool CpuId::SSE3()
    {
        return AuTestBit(f_1_ECX, 0);
    }

    bool CpuId::PCLMULQDQ()
    {
        return AuTestBit(f_1_ECX, 1);
    }

    bool CpuId::MONITOR()
    {
        return AuTestBit(f_1_ECX, 3);
    }

    bool CpuId::SSSE3()
    {
        return AuTestBit(f_1_ECX, 9);
    }

    bool CpuId::FMA()
    {
        return AuTestBit(f_1_ECX, 12);
    }

    bool CpuId::CMPXCHG16B()
    {
        return AuTestBit(f_1_ECX, 13);
    }

    bool CpuId::SSE41()
    {
        return AuTestBit(f_1_ECX, 19);
    }

    bool CpuId::SSE42()
    {
        return AuTestBit(f_1_ECX, 20);
    }

    bool CpuId::MOVBE()
    {
        return AuTestBit(f_1_ECX, 22);
    }

    bool CpuId::POPCNT()
    {
        return AuTestBit(f_1_ECX, 23);
    }

    bool CpuId::AES()
    {
        return AuTestBit(f_1_ECX, 25);
    }

    bool CpuId::XSAVE()
    {
        return AuTestBit(f_1_ECX, 26);
    }

    bool CpuId::OSXSAVE()
    {
        return AuTestBit(f_1_ECX, 27);
    }

    bool CpuId::AVX()
    {
        return AuTestBit(f_1_ECX, 28);
    }

    bool CpuId::F16C()
    {
        return AuTestBit(f_1_ECX, 29);
    }

    bool CpuId::RDRAND()
    {
        return AuTestBit(f_1_ECX, 30);
    }

    bool CpuId::MSR()
    {
        return AuTestBit(f_1_EDX, 5);
    }

    bool CpuId::CX8()
    {
        return AuTestBit(f_1_EDX, 8);
    }

    bool CpuId::SEP()
    {
        return AuTestBit(f_1_EDX, 11);
    }

    bool CpuId::CMOV()
    {
        return AuTestBit(f_1_EDX, 15);
    }

    bool CpuId::CLFSH()
    {
        return AuTestBit(f_1_EDX, 19);
    }

    bool CpuId::MMX()
    {
        return AuTestBit(f_1_EDX, 23);
    }

    bool CpuId::FXSR()
    {
        return AuTestBit(f_1_EDX, 24);
    }

    bool CpuId::SSE()
    {
        return AuTestBit(f_1_EDX, 25);
    }

    bool CpuId::SSE2()
    {
        return AuTestBit(f_1_EDX, 26);
    }

    bool CpuId::FSGSBASE()
    {
        return AuTestBit(f_7_EBX, 0);
    }

    bool CpuId::BMI1()
    {
        return AuTestBit(f_7_EBX, 3);
    }

    bool CpuId::HLE()
    {
        return isIntel && AuTestBit(f_7_EBX, 4);
    }

    bool CpuId::AVX2()
    {
        return AuTestBit(f_7_EBX, 5);
    }

    bool CpuId::BMI2()
    {
        return AuTestBit(f_7_EBX, 8);
    }

    bool CpuId::ERMS()
    {
        return AuTestBit(f_7_EBX, 9);
    }

    bool CpuId::INVPCID()
    {
        return AuTestBit(f_7_EBX, 10);
    }

    bool CpuId::RTM()
    {
        return isIntel && AuTestBit(f_7_EBX, 11);
    }

    bool CpuId::AVX512F()
    {
        return AuTestBit(f_7_EBX, 16);
    }

    bool CpuId::RDSEED()
    {
        return AuTestBit(f_7_EBX, 18);
    }

    bool CpuId::ADX()
    {
        return AuTestBit(f_7_EBX, 19);
    }

    bool CpuId::AVX512PF()
    {
        return AuTestBit(f_7_EBX, 26);
    }

    bool CpuId::AVX512ER()
    {
        return AuTestBit(f_7_EBX, 27);
    }

    bool CpuId::AVX512CD()
    {
        return AuTestBit(f_7_EBX, 28);
    }

    bool CpuId::SHA()
    {
        return AuTestBit(f_7_EBX, 29);
    }

    bool CpuId::PREFETCHWT1()
    {
        return AuTestBit(f_7_ECX, 0);
    }

    bool CpuId::LAHF()
    {
        return AuTestBit(f_81_ECX, 0);
    }

    bool CpuId::LZCNT()
    {
        return isIntel && AuTestBit(f_81_ECX, 5);
    }

    bool CpuId::ABM()
    {
        return isAMD && AuTestBit(f_81_ECX, 5);
    }

    bool CpuId::SSE4a()
    {
        return isAMD && AuTestBit(f_81_ECX, 6);
    }

    bool CpuId::XOP()
    {
        return isAMD && AuTestBit(f_81_ECX, 11);
    }

    bool CpuId::TBM()
    {
        return isAMD && AuTestBit(f_81_ECX, 21);
    }

    bool CpuId::SYSCALL()
    {
        return isIntel && AuTestBit(f_81_EDX, 11);
    }

    bool CpuId::MMXEXT()
    {
        return isAMD && AuTestBit(f_81_EDX, 22);
    }

    bool CpuId::RDTSCP()
    {
        return isIntel && AuTestBit(f_81_EDX, 27);
    }

    bool CpuId::_3DNOWEXT()
    {
        return isAMD && AuTestBit(f_81_EDX, 30);
    }

    bool CpuId::_3DNOW()
    {
        return isAMD && AuTestBit(f_81_EDX, 31);
    }

    AUKN_SYM const CpuInfo &GetCPUInfo()
    {
        return gCpuInfo;
    }

    static void SetCpuId()
    {
        // Credit: https://docs.microsoft.com/en-us/cpp/intrinsics/cpuid-cpuidex?view=msvc-160
    #if defined(AURORA_ARCH_X64) || defined(AURORA_ARCH_X86)
        AuList<CPUIdContext> data;
        AuList<CPUIdContext> extdata;

        auto cpuInfo = cpuid(0);
        auto nIds = cpuInfo.eax;

        for (int i = 0; i <= nIds; ++i)
        {
            data.push_back(cpuid(i));
        }

        char vendor[0x20];
        memset(vendor, 0, sizeof(vendor));
        *reinterpret_cast<int*>(vendor)     = cpuInfo.ebx;
        *reinterpret_cast<int*>(vendor + 4) = cpuInfo.edx;
        *reinterpret_cast<int*>(vendor + 8) = cpuInfo.ecx;
        
        gCpuInfo.cpuId.vendor = vendor;

        if (gCpuInfo.cpuId.vendor == "GenuineIntel")
        {
            gCpuInfo.cpuId.isIntel = true;
        }
        else if (gCpuInfo.cpuId.vendor == "AuthenticAMD")
        {
            gCpuInfo.cpuId.isAMD = true;
        }

        // load bitset with flags for function 0x00000001
        if (nIds >= 1)
        {
            gCpuInfo.cpuId.f_1_ECX = data[1].ecx;
            gCpuInfo.cpuId.f_1_EDX = data[1].edx;
        }

        // load bitset with flags for function 0x00000007
        if (nIds >= 7)
        {
            gCpuInfo.cpuId.f_7_EBX = data[7].ebx;
            gCpuInfo.cpuId.f_7_ECX = data[7].ecx;
        }

        // gets the number of the highest valid extended ID.
        auto cpui = cpuid(0x80000000);
        auto nExIds = cpui.eax;

        char brand[0x40];
        memset(brand, 0, sizeof(brand));

        for (int i = 0x80000000; i <= nExIds; ++i)
        {
            extdata.push_back(cpuid(i));
        }

        // load bitset with flags for function 0x80000001
        if (nExIds >= 0x80000001)
        {
            gCpuInfo.cpuId.f_81_ECX = extdata[1].ecx;
            gCpuInfo.cpuId.f_81_EDX = extdata[1].edx;
        }

        // Interpret CPU brand string if reported
        if (nExIds >= 0x80000004)
        {
            memcpy(brand, &extdata[2], sizeof(cpui));
            memcpy(brand + 16, &extdata[3], sizeof(cpui));
            memcpy(brand + 32, &extdata[4], sizeof(cpui));
            gCpuInfo.cpuId.brand = brand;
        }
    #endif
    }

    static void SetCpuTopology()
    {
    #if defined(AURORA_IS_MODERNNT_DERIVED)
        SYSTEM_LOGICAL_PROCESSOR_INFORMATION sysinfo[128];
        DWORD length = AuArraySize(sysinfo) * sizeof(*sysinfo);
        
        if (!GetLogicalProcessorInformation(sysinfo, &length))
        {
            SYSTEM_INFO sysinfo;
            GetSystemInfo(&sysinfo);
            gCpuInfo.socket  = 1;
            gCpuInfo.cores   = 1;
            gCpuInfo.threads = sysinfo.dwNumberOfProcessors;
            return;
        }

        length /= sizeof(*sysinfo);
        
        gCpuInfo.socket  = 0;
        gCpuInfo.cores   = 0;
        gCpuInfo.threads = 0;

        for (auto i = 0; i < length; i++)
        {
            if (sysinfo[i].Relationship == RelationProcessorCore)
            {
                //TODO: fuck it, if some macro fuckery, use popcnt on x86
                //we just need to count the bits. first it was just two BitScanForwards. discontiguous cores fucked things up so now we have a loop just to count a few bits.
                gCpuInfo.cores++;

                auto mask = sysinfo[i].ProcessorMask;
                unsigned long offset {}, tmp;
                while (offset != (sizeof(offset) * 8)) 
                {
                    // Count the index to a 1
                    if (BitScanForward(&tmp, mask >> offset) == 0) break; // mask was zero, end of scan
                    offset += tmp;
                    
                    // Count the 1's by inverting the bitmap and counting to 1 
                    BitScanForward(&tmp, ~(mask >> offset));
                    offset += tmp;
                    
                    // Increment threads by the bits set in 
                    gCpuInfo.threads += tmp;
                }
            }
            else if (sysinfo[i].Relationship == RelationProcessorPackage)
            {
                gCpuInfo.socket++;
            }
        }

    #elif defined(AURORA_IS_BSD_DERIVED)

        auto opt = QueryBsdHwStat(HW_AVAILCPU);
        if (opt.value_or(0) < 1)
        {
            opt = QueryBsdHwStat(HW_NCPU);
        }

        gCpuInfo.socket  = 1;
        gCpuInfo.cores   = 1;
        gCpuInfo.threads = opt.value_or(1);

    #elif defined(AURORA_IS_LINUX_DERIVED)
        
        gCpuInfo.socket  = 1;
        gCpuInfo.cores   = 1;
        gCpuInfo.threads = get_nprocs();

    #elif defined(AURORA_IS_POSIX_DERIVED)
        
        gCpuInfo.socket  = 1;
        gCpuInfo.cores   = 1;
        gCpuInfo.threads = sysconf(_SC_NPROCESSORS_ONLN);

    #endif
    }

    void InitCpuInfo()
    {
        gCpuInfo.cpuArch = Aurora::Build::kCurrentArchitecture;
        SetCpuId();
        SetCpuTopology();
    }
}