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

    File: AuCpuId.NT.cpp
    Date: 2022-1-25
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "AuHWInfo.hpp"
#include "AuCpuInfo.hpp"
#include "AuCpuInfo.NT.hpp"

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

namespace Aurora::HWInfo
{
    AUKN_SYM AuUInt32 gWin32CPUSetLookupTable[512] {};

    static bool SetWindows10CpuSetInfoSlow()
    {
        SYSTEM_CPU_SET_INFORMATION cpuSetInfo[128];
        SYSTEM_LOGICAL_PROCESSOR_INFORMATION sysinfo[128];
        DWORD dwLength {};

        if (!pGetSystemCpuSetInformation)
        {
            return false;
        }

        if (!pGetSystemCpuSetInformation(cpuSetInfo, sizeof(cpuSetInfo), &dwLength, 0, 0))
        {
            return false;
        }

        struct CpuInfo
        {
            AuList<AuUInt8> low;
            AuList<AuPair<AuUInt8, CpuBitId>> server;
            CpuBitId mask;
        };

        AuBST<AuUInt8, CpuInfo> cpuThreads;

        AuUInt8 uThreadCount = AuUInt8(dwLength / sizeof(decltype(*cpuSetInfo)));

        AuUInt8 uBestClass {};
        bool bHasBestClass {};

        if (uThreadCount)
        {
            AuUInt8 countTable[255] {};
            for (AU_ITERATE_N(i, uThreadCount))
            {
                countTable[cpuSetInfo[i].CpuSet.EfficiencyClass]++;
            }

            AuUInt8 uBestClassCounter {};
            for (AU_ITERATE_N(i, AuArraySize(countTable)))
            {
                auto uCurrentCount = countTable[i];

                if (uCurrentCount <= uBestClassCounter)
                {
                    continue;
                }

                if (!i)
                {
                    continue;
                }

                uBestClassCounter = uCurrentCount;
                bHasBestClass = true;
                uBestClass = i;
            }
        }

        for (AU_ITERATE_N(i, uThreadCount))
        {
            auto &idx = cpuThreads[cpuSetInfo[i].CpuSet.CoreIndex];

            //  Win7 KAFFINITY = u64 affinity masks 
            //  Windows 10 + seems to be ((this->group + 1ul) * 0x100ul) + index
            //  Windows internals says... 
            //      ULONG sets[] = { 0x100, 0x101, 0x102, 0x103 };
            //      ::SetProcessDefaultCpuSets/SetThreadSelectedCpuSets(::GetCurrentProcess(), sets, _countof(sets));
            //  (useless)
            //  People generally isolate group and keep logical processors in a different set, kinda worthless for bitwise math

            SysAssert(cpuSetInfo[i].CpuSet.LogicalProcessorIndex < 64);

        #if defined(_AU_MASSIVE_CPUID)
            SysAssert(cpuSetInfo[i].CpuSet.Group < 4);
        #else
            SysAssert(cpuSetInfo[i].CpuSet.Group < 2);
        #endif

            AuUInt8 id = AuUInt8(cpuSetInfo[i].CpuSet.LogicalProcessorIndex /*no greater than 64*/ + (cpuSetInfo[i].CpuSet.Group * 64));
            auto cpuId = CpuBitId(id);

        #if defined(AURORA_RUNTIME_USE_FAST_CPUSETS)
            auto sets = cpuId.ToCpuSets();
            SysAssert(sets.size() == 1);
            SysAssert(sets[0] == cpuSetInfo[i].CpuSet.Id);
        #endif

            // Future proofing for later NT versions 
            SysAssert(id < AuArraySize(gWin32CPUSetLookupTable));
            gWin32CPUSetLookupTable[id] = cpuSetInfo[i].CpuSet.Id;

            idx.server.push_back(AuMakePair(cpuSetInfo[i].CpuSet.EfficiencyClass, cpuId));
            idx.low.push_back(id);
            idx.mask.Add(cpuId);
            gCpuInfo.maskAllCores.Add(cpuId);
        }

        for (const auto &[cpuId, coreIds] : cpuThreads)
        {
            AuUInt64 shortMask {};
            for (const auto &[eClass, id] : coreIds.server)
            {
                if (bHasBestClass &&
                    eClass != uBestClass)
                {
                    gCpuInfo.maskECores.Add(id);
                }
                else
                {
                    gCpuInfo.maskPCores.Add(id);
                    gCpuInfo.pCoreTopology.push_back(id);
                }
            }

            for (const auto &id : coreIds.low)
            {
                shortMask |= AuUInt64(1) << AuUInt64(id);
            }

            gCpuInfo.coreTopology.push_back(coreIds.mask);
            gCpuInfo.threadTopology.push_back(shortMask);
        }

        gCpuInfo.uSocket = 1;
        gCpuInfo.uThreads = uThreadCount;
        gCpuInfo.uCores = AuUInt8(cpuThreads.size());

        if (!pGetLogicalProcessorInformation)
        {
            return true;
        }

        if (!pGetLogicalProcessorInformation(sysinfo, &dwLength))
        {
            return true;
        }

        gCpuInfo.uSocket = 0;
        dwLength /= sizeof(*sysinfo);

        for (auto i = 0u; i < dwLength; i++)
        {
            if (sysinfo[i].Relationship == RelationProcessorPackage)
            {
                gCpuInfo.uSocket++;
            }

            if (sysinfo[i].Relationship == RelationCache)
            {
                switch (sysinfo[i].Cache.Level)
                {
                case 1:
                    gCpuInfo.dwCacheLine = AuMax<AuUInt32>(gCpuInfo.dwCacheLine, sysinfo[i].Cache.LineSize);
                    gCpuInfo.dwCacheL1 = AuMax<AuUInt32>(gCpuInfo.dwCacheL1, sysinfo[i].Cache.Size);
                    break;
                case 2:
                    gCpuInfo.dwCacheL2 = AuMax<AuUInt32>(gCpuInfo.dwCacheL2, sysinfo[i].Cache.Size);
                    break;
                case 3:
                    gCpuInfo.dwCacheL3 = AuMax<AuUInt32>(gCpuInfo.dwCacheL3, sysinfo[i].Cache.Size);
                    break;
                }
            }
        }

        return true;
    }

    static bool SetWindowsXPSp3ExtendedInformation()
    {
        SYSTEM_LOGICAL_PROCESSOR_INFORMATION sysinfo[128];
        DWORD dwLength = AuArraySize(sysinfo) * sizeof(*sysinfo);

        if (!pGetLogicalProcessorInformation)
        {
            return false;
        }

        if (!pGetLogicalProcessorInformation(sysinfo, &dwLength))
        {
            return false;
        }

        dwLength /= sizeof(*sysinfo);

        gCpuInfo.uSocket = 0;
        gCpuInfo.uCores = 0;
        gCpuInfo.uThreads = 0;

        bool sparse = false;
        bool hasHTCores = false;
        for (auto i = 0u; i < dwLength; i++)
        {
            if (sysinfo[i].Relationship == RelationProcessorCore)
            {
                auto mask = sysinfo[i].ProcessorMask;

                gCpuInfo.uCores++;
                gCpuInfo.threadTopology.push_back(mask);

                CpuBitId serverId;
                serverId.lower = mask;
                gCpuInfo.coreTopology.push_back(serverId);
                gCpuInfo.maskAllCores.Add(serverId);

                auto uThreadsInCore = AuPopCnt(mask);
                gCpuInfo.uThreads += uThreadsInCore;

                hasHTCores |= (uThreadsInCore == 2);

                if (hasHTCores && (uThreadsInCore == 1))
                {
                    AuUInt8 idx {};
                    while (serverId.CpuBitScanForward(idx, idx))
                    {
                        gCpuInfo.maskECores.Add(idx);
                        idx++;
                    }
                }
                else
                {
                    AuUInt8 idx {};
                    while (serverId.CpuBitScanForward(idx, idx))
                    {
                        gCpuInfo.maskPCores.Add(idx);
                        idx++;
                    }

                    gCpuInfo.pCoreTopology.push_back(serverId);
                }
            }
            else if (sysinfo[i].Relationship == RelationProcessorPackage)
            {
                gCpuInfo.uSocket++;
            }
            else if (sysinfo[i].Relationship == RelationCache)
            {
                switch (sysinfo[i].Cache.Level)
                {
                case 1:
                    gCpuInfo.dwCacheLine = AuMax<AuUInt32>(gCpuInfo.dwCacheLine, sysinfo[i].Cache.LineSize);
                    gCpuInfo.dwCacheL1 = AuMax<AuUInt32>(gCpuInfo.dwCacheL1, sysinfo[i].Cache.Size);
                    break;
                case 2:
                    gCpuInfo.dwCacheL2 = AuMax<AuUInt32>(gCpuInfo.dwCacheL2, sysinfo[i].Cache.Size);
                    break;
                case 3:
                    gCpuInfo.dwCacheL3 = AuMax<AuUInt32>(gCpuInfo.dwCacheL3, sysinfo[i].Cache.Size);
                    break;
                }
            }
        }

        gCpuInfo.bMaskMTContig = !sparse;
        gCpuInfo.bMaskMTHalf = sparse;
        return true;
    }

    static void SetBasicWindowsXPInformation()
    {
        SYSTEM_INFO sysinfo;
        GetSystemInfo(&sysinfo);

        gCpuInfo.uSocket = 1;
        gCpuInfo.uCores = 1;
        gCpuInfo.uThreads = sysinfo.dwNumberOfProcessors;

        if (sysinfo.dwNumberOfProcessors & 1)
        {
            for (AU_ITERATE_N(i, gCpuInfo.uThreads))
            {
                auto mask = 1 << i;

                gCpuInfo.maskPCores.SetBit(i);

                gCpuInfo.threadTopology.push_back(mask);

                CpuBitId coreId;
                coreId.lower = mask;
                gCpuInfo.coreTopology.push_back(coreId);
                gCpuInfo.pCoreTopology.push_back(coreId);
                gCpuInfo.maskAllCores.Add(coreId);
            }

            gCpuInfo.uCores = gCpuInfo.uThreads;
        }
        else
        {
            for (AU_ITERATE_N(i, gCpuInfo.uThreads))
            {
                if (i & 1)
                {
                    continue;
                }

                gCpuInfo.maskPCores.SetBit(i);
                gCpuInfo.maskPCores.SetBit(i + 1);

                auto maskA = 1u << i;
                auto maskB = 1u << (i + 1);
                auto maskC = maskA | maskB;

                gCpuInfo.threadTopology.push_back(maskA);
                gCpuInfo.threadTopology.push_back(maskB);

                CpuBitId coreId;
                coreId.lower = maskC;
                gCpuInfo.coreTopology.push_back(coreId);
                gCpuInfo.pCoreTopology.push_back(coreId);
                gCpuInfo.maskAllCores.Add(coreId);
            }

            gCpuInfo.uCores = gCpuInfo.uThreads / 2;

            gCpuInfo.bMaskMTHalf = true;
        }
    }

    void SetCpuTopologyNT()
    {
        if (SetWindows10CpuSetInfoSlow())
        {
            return;
        }

        if (SetWindowsXPSp3ExtendedInformation())
        {
            return;
        }

        SetBasicWindowsXPInformation();
    }
}