#include "machine_code.h"

#include <algorithm>
#include <cstring>
#include <sstream>
#include <vector>

#ifdef _WIN32
    #ifndef WIN32_LEAN_AND_MEAN
        #define WIN32_LEAN_AND_MEAN
    #endif
    #include <windows.h>
    #include <winioctl.h>
    #include <intrin.h>
#elif defined(__APPLE__)
    #include <sys/sysctl.h>
    #include <sys/mount.h>
    #include <sys/stat.h>
#else // Linux
    #if defined(__x86_64__) || defined(__i386__)
        #include <cpuid.h>
    #endif
    #include <sys/statfs.h>
    #include <unistd.h>
#endif

namespace reginfo {

namespace {

// === SHA-256 implementation ===

static const unsigned int SHA256_K[64] = {
    0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
    0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
    0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
    0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
    0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
    0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
    0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
    0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2
};

static inline unsigned int rotr(unsigned int x, unsigned int n) {
    return (x >> n) | (x << (32 - n));
}

std::string sha256(const std::string& input) {
    // Pre-processing: pad the message
    std::vector<unsigned char> msg(input.begin(), input.end());
    unsigned long long bitLen = msg.size() * 8;
    msg.push_back(0x80);
    while ((msg.size() + 8) % 64 != 0)
        msg.push_back(0x00);
    for (int i = 7; i >= 0; --i)
        msg.push_back(static_cast<unsigned char>((bitLen >> (i * 8)) & 0xFF));

    // Initial hash values
    unsigned int H[8] = {
        0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a,
        0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19
    };

    // Process each 512-bit chunk
    for (size_t chunk = 0; chunk < msg.size(); chunk += 64) {
        unsigned int W[64];
        for (int i = 0; i < 16; ++i) {
            size_t off = chunk + i * 4;
            W[i] = (static_cast<unsigned int>(msg[off]) << 24)
                 | (static_cast<unsigned int>(msg[off + 1]) << 16)
                 | (static_cast<unsigned int>(msg[off + 2]) << 8)
                 | static_cast<unsigned int>(msg[off + 3]);
        }
        for (int i = 16; i < 64; ++i) {
            unsigned int s0 = rotr(W[i - 15], 7) ^ rotr(W[i - 15], 18) ^ (W[i - 15] >> 3);
            unsigned int s1 = rotr(W[i - 2], 17) ^ rotr(W[i - 2], 19) ^ (W[i - 2] >> 10);
            W[i] = W[i - 16] + s0 + W[i - 7] + s1;
        }

        unsigned int a = H[0], b = H[1], c = H[2], d = H[3];
        unsigned int e = H[4], f = H[5], g = H[6], h = H[7];

        for (int i = 0; i < 64; ++i) {
            unsigned int S1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25);
            unsigned int ch = (e & f) ^ (~e & g);
            unsigned int temp1 = h + S1 + ch + SHA256_K[i] + W[i];
            unsigned int S0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22);
            unsigned int maj = (a & b) ^ (a & c) ^ (b & c);
            unsigned int temp2 = S0 + maj;

            h = g; g = f; f = e; e = d + temp1;
            d = c; c = b; b = a; a = temp1 + temp2;
        }

        H[0] += a; H[1] += b; H[2] += c; H[3] += d;
        H[4] += e; H[5] += f; H[6] += g; H[7] += h;
    }

    // Format as 64 hex characters
    char hex[65];
    snprintf(hex, sizeof(hex),
        "%08X%08X%08X%08X%08X%08X%08X%08X",
        H[0], H[1], H[2], H[3], H[4], H[5], H[6], H[7]);
    return std::string(hex);
}

#ifdef _WIN32

std::string getCpuIdWin32() {
	SetThreadAffinityMask(GetCurrentThread(), 1);
    int cpuInfo[4] = {0};
    __cpuid(cpuInfo, 1);
    unsigned int cpuSerial = (unsigned int)((cpuInfo[0])^ (cpuInfo[1])  ^  (cpuInfo[2])^ (cpuInfo[3]));

    __cpuid(cpuInfo, 0x80000000);
    unsigned int extMax = static_cast<unsigned int>(cpuInfo[0]);
    std::string brand;
    if (extMax >= 0x80000004) {
        for (unsigned int i = 0x80000002; i <= 0x80000004; ++i) {
            __cpuid(cpuInfo, static_cast<int>(i));
            brand.append(reinterpret_cast<char*>(cpuInfo), 16);
        }
        // Trim trailing nulls/spaces
        while (!brand.empty() && (brand.back() == '\0' || brand.back() == ' '))
            brand.pop_back();
    }
    return brand + "|" + std::to_string(cpuSerial);
}

std::string getSystemDiskSerialWin32() {
    DWORD serialNumber = 0;
    char systemDrive[4] = {0};
    GetEnvironmentVariableA("SystemDrive", systemDrive, sizeof(systemDrive));
    if (systemDrive[0] == '\0')
        strcpy_s(systemDrive, "C:");

    std::string rootPath = std::string(systemDrive) + "\\";
    GetVolumeInformationA(rootPath.c_str(), nullptr, 0, &serialNumber,
                          nullptr, nullptr, nullptr, 0);

    DWORD physSerial = 0;
    std::string diskPath = std::string("\\\\.\\") + systemDrive[0] + ":";
    HANDLE hDevice = CreateFileA(diskPath.c_str(), 0,
        FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, 0, nullptr);
    if (hDevice != INVALID_HANDLE_VALUE) {
        STORAGE_DEVICE_NUMBER devNum;
        DWORD bytesReturned = 0;
        if (DeviceIoControl(hDevice, IOCTL_STORAGE_GET_DEVICE_NUMBER,
                nullptr, 0, &devNum, sizeof(devNum), &bytesReturned, nullptr)) {
            physSerial = devNum.DeviceNumber ^ devNum.PartitionNumber;
        }
        CloseHandle(hDevice);
    }

    return std::to_string(serialNumber) + "|" + std::to_string(physSerial);
}

#elif defined(__APPLE__)

std::string getCpuIdMac() {
    char buffer[256];
    size_t bufLen = sizeof(buffer);
    if (sysctlbyname("machdep.cpu.brand_string", buffer, &bufLen, nullptr, 0) == 0) {
        return std::string(buffer, bufLen - 1);
    }
    return "0";
}

std::string getSystemDiskSerialMac() {
    struct statfs fs;
    if (statfs("/", &fs) == 0) {
        unsigned long long diskId = 0;
        memcpy(&diskId, &fs.f_fsid, sizeof(fs.f_fsid));
        return std::to_string(diskId);
    }
    return "0";
}

#else // Linux

std::string getCpuIdLinux() {
    std::string result;
#if defined(__x86_64__) || defined(__i386__)
    unsigned int eax, ebx, ecx, edx;
    if (__get_cpuid(1, &eax, &ebx, &ecx, &edx)) {
        unsigned int cpuSerial = eax ^ ebx ^ ecx ^ edx;
        result = std::to_string(cpuSerial);
    }
#endif
    FILE* fp = fopen("/proc/cpuinfo", "r");
    if (fp) {
        char line[256];
        std::string modelName;
        while (fgets(line, sizeof(line), fp)) {
            if (strstr(line, "model name") || strstr(line, "Model Name")) {
                const char* colon = strchr(line, ':');
                if (colon) modelName = colon + 1;
                break;
            }
        }
        fclose(fp);
        if (!modelName.empty()) {
            while (!modelName.empty() && (modelName.back() == '\n' || modelName.back() == ' '))
                modelName.pop_back();
            result += "|" + modelName;
        }
    }
    return result.empty() ? "0" : result;
}

std::string getSystemDiskSerialLinux() {
    struct statfs fs;
    if (statfs("/", &fs) == 0) {
        unsigned long long diskId = 0;
        memcpy(&diskId, &fs.f_fsid, sizeof(fs.f_fsid));
        return std::to_string(diskId);
    }
    return "0";
}

#endif

} // anonymous namespace

std::string getMachineCode() {
#ifdef _WIN32
    std::string raw = getCpuIdWin32() + "|" + getSystemDiskSerialWin32();
#elif defined(__APPLE__)
    std::string raw = getCpuIdMac() + "|" + getSystemDiskSerialMac();
#else
    std::string raw = getCpuIdLinux() + "|" + getSystemDiskSerialLinux();
#endif
    return sha256(raw);
}

} // namespace reginfo