#include "HWIDCollector.h"
#include <windows.h>
#include <wincrypt.h>
#include <iphlpapi.h>
#include <wbemidl.h>
#include <comdef.h>
#include <cstdio>
#include <cstring>
#include <vector>

#pragma comment(lib, "advapi32.lib")
#pragma comment(lib, "iphlpapi.lib")
#pragma comment(lib, "wbemuuid.lib")
#pragma comment(lib, "ole32.lib")
#pragma comment(lib, "oleaut32.lib")

// ─── Singleton ────────────────────────────────────────

HWIDCollector* HWIDCollector::Get()
{
    static HWIDCollector instance;
    return &instance;
}

// ─── SHA-256 ──────────────────────────────────────────

std::string HWIDCollector::SHA256(const std::string& input)
{
    if (input.empty()) return {};

    HCRYPTPROV hProv = 0;
    HCRYPTHASH hHash = 0;
    std::string result;

    if (!CryptAcquireContextA(&hProv, nullptr, nullptr, PROV_RSA_AES, CRYPT_VERIFYCONTEXT))
        return {};

    if (!CryptCreateHash(hProv, CALG_SHA_256, 0, 0, &hHash))
    {
        CryptReleaseContext(hProv, 0);
        return {};
    }

    if (!CryptHashData(hHash, reinterpret_cast<const BYTE*>(input.data()),
                       static_cast<DWORD>(input.size()), 0))
    {
        CryptDestroyHash(hHash);
        CryptReleaseContext(hProv, 0);
        return {};
    }

    BYTE hash[32];
    DWORD hashLen = sizeof(hash);
    if (CryptGetHashParam(hHash, HP_HASHVAL, hash, &hashLen, 0))
    {
        char hex[65];
        for (int i = 0; i < 32; ++i)
            snprintf(hex + i * 2, 3, "%02x", hash[i]);
        result.assign(hex, 64);
    }

    CryptDestroyHash(hHash);
    CryptReleaseContext(hProv, 0);
    return result;
}

// ─── CPU: Registry ProcessorNameString ────────────────

std::string HWIDCollector::CollectCPUId()
{
    HKEY hKey;
    if (RegOpenKeyExA(HKEY_LOCAL_MACHINE,
            "HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0",
            0, KEY_READ, &hKey) != ERROR_SUCCESS)
    {
        printf("[HWID] CPU: registry key not found\n");
        return {};
    }

    char name[256] = {};
    DWORD size = sizeof(name);
    if (RegQueryValueExA(hKey, "ProcessorNameString", nullptr, nullptr,
            reinterpret_cast<BYTE*>(name), &size) != ERROR_SUCCESS)
    {
        printf("[HWID] CPU: ProcessorNameString not found\n");
        RegCloseKey(hKey);
        return {};
    }
    RegCloseKey(hKey);

    std::string cpuName(name);
    // Trim trailing whitespace/newlines
    while (!cpuName.empty() && (cpuName.back() == ' ' || cpuName.back() == '\n' || cpuName.back() == '\r'))
        cpuName.pop_back();

    std::string hash = SHA256(cpuName);
    printf("[HWID] CPU: %s -> %s\n", cpuName.c_str(), hash.c_str());
    return hash;
}

// ─── Board: WMI Win32_BaseBoard.SerialNumber ──────────

std::string HWIDCollector::CollectBoardSerial()
{
    std::string result;

    HRESULT hr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    bool comInitialized = SUCCEEDED(hr) || hr == S_FALSE || hr == RPC_E_CHANGED_MODE;
    // note: RPC_E_CHANGED_MODE means COM was already initialized with a different mode

    if (hr == RPC_E_CHANGED_MODE) comInitialized = false; // can't use safely

    IWbemLocator* pLoc = nullptr;
    IWbemServices* pSvc = nullptr;
    IEnumWbemClassObject* pEnum = nullptr;

    if (!comInitialized) goto cleanup;

    hr = CoCreateInstance(CLSID_WbemLocator, nullptr, CLSCTX_INPROC_SERVER,
                          IID_IWbemLocator, reinterpret_cast<void**>(&pLoc));
    if (FAILED(hr) || !pLoc) { printf("[HWID] Board: CoCreateInstance failed 0x%lx\n", hr); goto cleanup; }

    hr = pLoc->ConnectServer(_bstr_t(L"ROOT\\CIMV2"), nullptr, nullptr, nullptr,
                             0, nullptr, nullptr, &pSvc);
    if (FAILED(hr) || !pSvc) { printf("[HWID] Board: ConnectServer failed 0x%lx\n", hr); goto cleanup; }

    // Set proxy security
    CoSetProxyBlanket(pSvc, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, nullptr,
                      RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE, nullptr, EOAC_NONE);

    hr = pSvc->ExecQuery(
        _bstr_t(L"WQL"),
        _bstr_t(L"SELECT SerialNumber FROM Win32_BaseBoard"),
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
        nullptr, &pEnum);
    if (FAILED(hr) || !pEnum) { printf("[HWID] Board: ExecQuery failed 0x%lx\n", hr); goto cleanup; }

    {
        IWbemClassObject* pObj = nullptr;
        ULONG returned = 0;
        hr = pEnum->Next(WBEM_INFINITE, 1, &pObj, &returned);
        if (SUCCEEDED(hr) && pObj)
        {
            VARIANT vt;
            VariantInit(&vt);
            hr = pObj->Get(L"SerialNumber", 0, &vt, nullptr, nullptr);
            if (SUCCEEDED(hr) && vt.vt == VT_BSTR && vt.bstrVal)
            {
                std::wstring ws(vt.bstrVal);
                result = std::string(ws.begin(), ws.end());
            }
            VariantClear(&vt);
            pObj->Release();
        }
    }

    // Fallback: try Win32_BIOS
    if (result.empty() || result == "To be filled by O.E.M.")
    {
        if (pEnum) { pEnum->Release(); pEnum = nullptr; }
        hr = pSvc->ExecQuery(
            _bstr_t(L"WQL"),
            _bstr_t(L"SELECT SerialNumber FROM Win32_BIOS"),
            WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
            nullptr, &pEnum);
        if (SUCCEEDED(hr) && pEnum)
        {
            IWbemClassObject* pObj = nullptr;
            ULONG returned = 0;
            hr = pEnum->Next(WBEM_INFINITE, 1, &pObj, &returned);
            if (SUCCEEDED(hr) && pObj)
            {
                VARIANT vt;
                VariantInit(&vt);
                hr = pObj->Get(L"SerialNumber", 0, &vt, nullptr, nullptr);
                if (SUCCEEDED(hr) && vt.vt == VT_BSTR && vt.bstrVal)
                {
                    std::wstring ws(vt.bstrVal);
                    result = std::string(ws.begin(), ws.end());
                }
                VariantClear(&vt);
                pObj->Release();
            }
        }
    }

cleanup:
    if (pEnum) pEnum->Release();
    if (pSvc) pSvc->Release();
    if (pLoc) pLoc->Release();
    if (comInitialized) CoUninitialize();

    std::string hash = SHA256(result);
    printf("[HWID] Board: '%s' -> %s\n", result.c_str(), hash.c_str());
    return hash;
}

// ─── Disk: GetVolumeInformationA ──────────────────────

std::string HWIDCollector::CollectVolumeSerial()
{
    DWORD serial = 0;
    if (!GetVolumeInformationA("C:\\", nullptr, 0, &serial, nullptr, nullptr, nullptr, 0))
    {
        printf("[HWID] Disk: GetVolumeInformation failed: 0x%lx\n", GetLastError());
        return {};
    }

    char buf[16];
    snprintf(buf, sizeof(buf), "%08lx", serial);
    std::string serialStr(buf);

    std::string hash = SHA256(serialStr);
    printf("[HWID] Disk: %s -> %s\n", serialStr.c_str(), hash.c_str());
    return hash;
}

// ─── MAC: GetAdaptersInfo ─────────────────────────────

std::string HWIDCollector::CollectMacAddress()
{
    ULONG bufLen = sizeof(IP_ADAPTER_INFO);
    std::vector<BYTE> buffer(bufLen);
    PIP_ADAPTER_INFO pAdapterInfo = reinterpret_cast<PIP_ADAPTER_INFO>(buffer.data());

    DWORD ret = GetAdaptersInfo(pAdapterInfo, &bufLen);
    if (ret == ERROR_BUFFER_OVERFLOW)
    {
        buffer.resize(bufLen);
        pAdapterInfo = reinterpret_cast<PIP_ADAPTER_INFO>(buffer.data());
        ret = GetAdaptersInfo(pAdapterInfo, &bufLen);
    }

    if (ret != ERROR_SUCCESS)
    {
        printf("[HWID] MAC: GetAdaptersInfo failed: 0x%lx\n", ret);
        return {};
    }

    // Find first physical adapter (Ethernet or Wi-Fi, skip loopback/tunnel)
    for (PIP_ADAPTER_INFO p = pAdapterInfo; p; p = p->Next)
    {
        // Skip loopback and tunnel adapters
        if (p->Type == MIB_IF_TYPE_LOOPBACK) continue;
        if (p->Type == MIB_IF_TYPE_OTHER) continue;
        if (p->AddressLength == 0) continue;

        char mac[18];
        snprintf(mac, sizeof(mac), "%02x:%02x:%02x:%02x:%02x:%02x",
                 p->Address[0], p->Address[1], p->Address[2],
                 p->Address[3], p->Address[4], p->Address[5]);

        std::string macStr(mac);
        std::string hash = SHA256(macStr);
        printf("[HWID] MAC: %s (%s) -> %s\n", p->Description, macStr.c_str(), hash.c_str());
        return hash;
    }

    printf("[HWID] MAC: no physical adapter found\n");
    return {};
}

// ─── Public API ───────────────────────────────────────

const HWIDComponents& HWIDCollector::Collect()
{
    if (m_collected)
        return m_components;

    printf("[HWID] Collecting hardware identifiers...\n");

    m_components.cpu = CollectCPUId();
    m_components.board = CollectBoardSerial();
    m_components.disk = CollectVolumeSerial();
    m_components.mac = CollectMacAddress();

    // Valid if we have at least CPU + Board (core 80%)
    m_components.valid = !m_components.cpu.empty() && !m_components.board.empty();
    m_collected = true;

    printf("[HWID] Collection complete (valid=%d)\n", m_components.valid);
    return m_components;
}
