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

    File: CpuId.Nt.cpp
    Date: 2022-1-25
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "HWInfo.hpp"
#include "CpuInfo.hpp"
#include "CpuInfo.NT.hpp"

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

namespace Aurora::HWInfo
{
    static bool IsWindowsLTSC()
    {
        OSVERSIONINFOEXW osvi = {sizeof(osvi), 0, 0, 0, 0, {0}, 0, 0, VER_SUITE_ENTERPRISE, 0};
        DWORDLONG        const dwlConditionMask = VerSetConditionMask(0, VER_SUITENAME, VER_EQUAL);

        return !VerifyVersionInfoW(&osvi, VER_SUITENAME, dwlConditionMask);
    }

    static bool TrySetNtCpuSetInfoSlowExtended()
    {
        SYSTEM_CPU_SET_INFORMATION cpuSetInfo[128];
        SYSTEM_LOGICAL_PROCESSOR_INFORMATION sysinfo[128];
        DWORD length = {};

        if (!GetSystemCpuSetInformation(cpuSetInfo, sizeof(cpuSetInfo), &length, 0, 0))
        {
            return false;
        }

        struct CpuInfo
        {
            AuList<AuUInt8> low;
            AuList<CpuBitId> server;
            CpuBitId mask;
        };
        AuBST<AuUInt8, CpuInfo> cpuThreads;
        AuUInt8 cpuCount;

        cpuCount = length / sizeof(decltype(*cpuSetInfo));

        for (int i = 0; i < cpuCount; i++)
        {
            auto &idx = cpuThreads[cpuSetInfo[i].CpuSet.CoreIndex + cpuSetInfo[i].CpuSet.Group];
            AuUInt8 id = AuUInt8(cpuSetInfo[i].CpuSet.LogicalProcessorIndex + cpuSetInfo[i].CpuSet.Group);
            auto cpuId = CpuBitId(id);
            idx.server.push_back(cpuId);
            idx.low.push_back(id);
            idx.mask.Add(cpuId);
        }

        for (const auto &[cpuId, coreIds] : cpuThreads)
        {
            AuUInt64 shortMask {};
            for (const auto &id : coreIds.server)
            {
                // TODO (scar):
                if (false)
                {
                    gCpuInfo.maskECores.Add(id);
                }
            }

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

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

        gCpuInfo.socket = 1;
        gCpuInfo.threads = cpuCount;
        gCpuInfo.cores = cpuThreads.size();

        if (!GetLogicalProcessorInformation(sysinfo, &length))
        {
            return true;
        }

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

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

        return true;
    }

    void SetCpuTopologyNT()
    {
        SYSTEM_LOGICAL_PROCESSOR_INFORMATION sysinfo[128];
        DWORD length = AuArraySize(sysinfo) * sizeof(*sysinfo);

        if (SWInfo::IsWindows10OrGreater() || IsWindowsServer() || IsWindowsLTSC())
        {
            if (TrySetNtCpuSetInfoSlowExtended())
            {
                return;
            }
        }

        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;

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

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

                CpuBitId serverId;
                serverId.lower = mask;
                gCpuInfo.serverTopology.push_back(mask);

                // 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.
                int counter {};
                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;

                    if (counter++) sparse = true;

                    // Increment threads by the bits set in 
                    gCpuInfo.threads += tmp;
                }
            }
            else if (sysinfo[i].Relationship == RelationProcessorPackage)
            {
                gCpuInfo.socket++;
            }
        }

        gCpuInfo.maskMTContig = !sparse;
        gCpuInfo.maskMTHalf = sparse;
    }
}