#include "memory/pe_fingerprint.hpp"

#include <cstring>
#include <fstream>
#include <vector>

namespace pe_fingerprint {

namespace {

// Hằng số round SHA-256 (NIST FIPS 180-4, mục 4.2.2)
constexpr uint32_t kSha256K[64] = {
    0x428a2f98u, 0x71374491u, 0xb5c0fbcfu, 0xe9b5dba5u, 0x3956c25bu, 0x59f111f1u, 0x923f82a4u, 0xab1c5ed5u,
    0xd807aa98u, 0x12835b01u, 0x243185beu, 0x550c7dc3u, 0x72be5d74u, 0x80deb1feu, 0x9bdc06a7u, 0xc19bf174u,
    0xe49b69c1u, 0xefbe4786u, 0x0fc19dc6u, 0x240ca1ccu, 0x2de92c6fu, 0x4a7484aau, 0x5cb0a9dcu, 0x76f988dau,
    0x983e5152u, 0xa831c66du, 0xb00327c8u, 0xbf597fc7u, 0xc6e00bf3u, 0xd5a79147u, 0x06ca6351u, 0x14292967u,
    0x27b70a85u, 0x2e1b2138u, 0x4d2c6dfcu, 0x53380d13u, 0x650a7354u, 0x766a0abbu, 0x81c2c92eu, 0x92722c85u,
    0xa2bfe8a1u, 0xa81a664bu, 0xc24b8b70u, 0xc76c51a3u, 0xd192e819u, 0xd6990624u, 0xf40e3585u, 0x106aa070u,
    0x19a4c116u, 0x1e376c08u, 0x2748774cu, 0x34b0bcb5u, 0x391c0cb3u, 0x4ed8aa4au, 0x5b9cca4fu, 0x682e6ff3u,
    0x748f82eeu, 0x78a5636fu, 0x84c87814u, 0x8cc70208u, 0x90befffau, 0xa4506cebu, 0xbef9a3f7u, 0xc67178f2u
};

inline uint32_t RotateRight(uint32_t value, uint32_t bits) {
    return (value >> bits) | (value << (32u - bits));
}

// Bảng cấu trúc PE đọc thủ công (không phụ thuộc alignment của windows.h)
#pragma pack(push, 1)
struct RawDosHeader {
    uint16_t e_magic;
    uint8_t padding[58];
    uint32_t e_lfanew;
};

struct RawCoffHeader {
    uint32_t signature;
    uint16_t machine;
    uint16_t numberOfSections;
    uint32_t timeDateStamp;
    uint32_t pointerToSymbolTable;
    uint32_t numberOfSymbols;
    uint16_t sizeOfOptionalHeader;
    uint16_t characteristics;
};

struct RawSectionHeader {
    char name[8];
    uint32_t virtualSize;
    uint32_t virtualAddress;
    uint32_t sizeOfRawData;
    uint32_t pointerToRawData;
    uint32_t pointerToRelocations;
    uint32_t pointerToLinenumbers;
    uint16_t numberOfRelocations;
    uint16_t numberOfLinenumbers;
    uint32_t characteristics;
};
#pragma pack(pop)

static_assert(sizeof(RawDosHeader) == 64, "DOS header phai dung 64 bytes");
static_assert(sizeof(RawCoffHeader) == 24, "COFF header phai dung 24 bytes");
static_assert(sizeof(RawSectionHeader) == 40, "Section header phai dung 40 bytes");

// Offset trong Optional Header PE32+ (magic 0x20B)
constexpr uint16_t kPe32PlusMagic = 0x20Bu;
constexpr size_t kSizeOfImageOffsetInOptHeader = 56;
constexpr size_t kSizeOfHeadersOffsetInOptHeader = 60;

// FNV-1a 64-bit
constexpr uint64_t kFnvOffsetBasis = 0xcbf29ce484222325ull;
constexpr uint64_t kFnvPrime = 0x100000001b3ull;

inline uint64_t Fnv1a(uint64_t hash, const void* data, size_t length) {
    const auto* bytes = static_cast<const uint8_t*>(data);
    for (size_t i = 0; i < length; ++i) {
        hash ^= bytes[i];
        hash *= kFnvPrime;
    }
    return hash;
}

} // namespace

// ==========================================================
// SHA-256 (self-contained)
// ==========================================================
void Sha256Context::Transform(const uint8_t* block) {
    uint32_t w[64];
    for (int t = 0; t < 16; ++t) {
        w[t] = (static_cast<uint32_t>(block[t * 4 + 0]) << 24) |
               (static_cast<uint32_t>(block[t * 4 + 1]) << 16) |
               (static_cast<uint32_t>(block[t * 4 + 2]) << 8) |
               (static_cast<uint32_t>(block[t * 4 + 3]));
    }
    for (int t = 16; t < 64; ++t) {
        const uint32_t s0 = RotateRight(w[t - 15], 7) ^ RotateRight(w[t - 15], 18) ^ (w[t - 15] >> 3);
        const uint32_t s1 = RotateRight(w[t - 2], 17) ^ RotateRight(w[t - 2], 19) ^ (w[t - 2] >> 10);
        w[t] = w[t - 16] + s0 + w[t - 7] + s1;
    }

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

    for (int t = 0; t < 64; ++t) {
        const uint32_t S1 = RotateRight(e, 6) ^ RotateRight(e, 11) ^ RotateRight(e, 25);
        const uint32_t ch = (e & f) ^ ((~e) & g);
        const uint32_t temp1 = h + S1 + ch + kSha256K[t] + w[t];
        const uint32_t S0 = RotateRight(a, 2) ^ RotateRight(a, 13) ^ RotateRight(a, 22);
        const uint32_t maj = (a & b) ^ (a & c) ^ (b & c);
        const uint32_t temp2 = S0 + maj;

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

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

void Sha256Context::Update(const uint8_t* data, size_t length) {
    m_totalBytes += length;
    size_t offset = 0;

    if (m_bufferLength > 0) {
        const size_t need = 64 - m_bufferLength;
        const size_t take = (length < need) ? length : need;
        std::memcpy(m_buffer + m_bufferLength, data, take);
        m_bufferLength += take;
        offset += take;
        if (m_bufferLength == 64) {
            Transform(m_buffer);
            m_bufferLength = 0;
        }
    }

    while (offset + 64 <= length) {
        Transform(data + offset);
        offset += 64;
    }

    if (offset < length) {
        const size_t remaining = length - offset;
        std::memcpy(m_buffer, data + offset, remaining);
        m_bufferLength = remaining;
    }
}

std::array<uint8_t, 32> Sha256Context::Final() {
    const uint64_t bitLength = m_totalBytes * 8ull;

    uint8_t tail[128] = {};
    size_t tailLength = m_bufferLength;
    if (tailLength > 0) {
        std::memcpy(tail, m_buffer, tailLength);
    }
    tail[tailLength++] = 0x80;

    // Không còn đủ 8 byte cho trường độ dài trong block 64 byte -> dùng 2 block
    const size_t blockCount = (tailLength <= 56) ? 64 : 128;
    for (int i = 0; i < 8; ++i) {
        tail[blockCount - 1 - static_cast<size_t>(i)] =
            static_cast<uint8_t>((bitLength >> (8 * i)) & 0xFFu);
    }
    for (size_t off = 0; off < blockCount; off += 64) {
        Transform(tail + off);
    }

    std::array<uint8_t, 32> digest{};
    for (int i = 0; i < 8; ++i) {
        digest[i * 4 + 0] = static_cast<uint8_t>((m_state[i] >> 24) & 0xFFu);
        digest[i * 4 + 1] = static_cast<uint8_t>((m_state[i] >> 16) & 0xFFu);
        digest[i * 4 + 2] = static_cast<uint8_t>((m_state[i] >> 8) & 0xFFu);
        digest[i * 4 + 3] = static_cast<uint8_t>(m_state[i] & 0xFFu);
    }
    return digest;
}

std::array<uint8_t, 32> Sha256(const uint8_t* data, size_t length) {
    Sha256Context ctx;
    if (data != nullptr && length > 0) {
        ctx.Update(data, length);
    }
    return ctx.Final();
}

bool Sha256File(const std::string& path, std::array<uint8_t, 32>& outHash) {
    std::ifstream file(path, std::ios::binary);
    if (!file.is_open()) {
        return false;
    }

    Sha256Context ctx;
    std::vector<uint8_t> chunk(1024 * 1024);
    while (file) {
        file.read(reinterpret_cast<char*>(chunk.data()), static_cast<std::streamsize>(chunk.size()));
        const std::streamsize got = file.gcount();
        if (got > 0) {
            ctx.Update(chunk.data(), static_cast<size_t>(got));
        }
    }

    outHash = ctx.Final();
    return true;
}

std::string ToHex(const uint8_t* data, size_t length) {
    static const char* kDigits = "0123456789abcdef";
    std::string out;
    out.reserve(length * 2);
    for (size_t i = 0; i < length; ++i) {
        out.push_back(kDigits[(data[i] >> 4) & 0x0F]);
        out.push_back(kDigits[data[i] & 0x0F]);
    }
    return out;
}


// ==========================================================
// ClientBuildId
// ==========================================================
std::string ClientBuildId::ShortId() const {
    static const char* kDigits = "0123456789ABCDEF";
    std::string out;
    out.reserve(8);
    for (int shift = 28; shift >= 0; shift -= 4) {
        out.push_back(kDigits[(timeDateStamp >> shift) & 0x0Fu]);
    }
    return out;
}

uint64_t ClientBuildId::SectionLayoutHash() const {
    uint64_t hash = kFnvOffsetBasis;
    hash = Fnv1a(hash, &numberOfSections, sizeof(numberOfSections));
    hash = Fnv1a(hash, &sizeOfImage, sizeof(sizeOfImage));
    for (const auto& section : sections) {
        hash = Fnv1a(hash, section.name.data(), section.name.size());
        hash = Fnv1a(hash, &section.virtualAddress, sizeof(section.virtualAddress));
        hash = Fnv1a(hash, &section.virtualSize, sizeof(section.virtualSize));
        hash = Fnv1a(hash, &section.rawSize, sizeof(section.rawSize));
        hash = Fnv1a(hash, &section.characteristics, sizeof(section.characteristics));
    }
    return hash;
}

bool ClientBuildId::SameBuildAs(const ClientBuildId& other) const {
    if (!valid || !other.valid) {
        return false;
    }
    return timeDateStamp == other.timeDateStamp &&
           sizeOfImage == other.sizeOfImage &&
           SectionLayoutHash() == other.SectionLayoutHash();
}

const SectionInfo* ClientBuildId::FindSection(const std::string& sectionName) const {
    for (const auto& section : sections) {
        if (section.name == sectionName) {
            return &section;
        }
    }
    return nullptr;
}

// ==========================================================
// Nạp header + bảng section của tệp PE
// ==========================================================
bool LoadPeHeaders(const std::string& path, ClientBuildId& outBuild, std::string& outError) {
    outBuild = ClientBuildId{};
    outError.clear();

    std::ifstream file(path, std::ios::binary);
    if (!file.is_open()) {
        outError = "Khong mo duoc tep: " + path;
        return false;
    }

    file.seekg(0, std::ios::end);
    const auto fileSize = static_cast<uint64_t>(file.tellg());
    file.seekg(0, std::ios::beg);
    outBuild.fileSize = fileSize;
    outBuild.path = path;

    RawDosHeader dos{};
    file.read(reinterpret_cast<char*>(&dos), sizeof(dos));
    if (file.gcount() != static_cast<std::streamsize>(sizeof(dos)) || dos.e_magic != IMAGE_DOS_SIGNATURE) {
        outError = "DOS header khong hop le (magic != 'MZ')";
        return false;
    }
    if (dos.e_lfanew == 0 || static_cast<uint64_t>(dos.e_lfanew) + sizeof(RawCoffHeader) > fileSize) {
        outError = "e_lfanew tro ra ngoai pham vi tep";
        return false;
    }

    file.seekg(dos.e_lfanew, std::ios::beg);
    RawCoffHeader coff{};
    file.read(reinterpret_cast<char*>(&coff), sizeof(coff));
    if (file.gcount() != static_cast<std::streamsize>(sizeof(coff))) {
        outError = "Khong doc du COFF header";
        return false;
    }
    if (coff.signature != IMAGE_NT_SIGNATURE) {
        outError = "PE signature khong hop le";
        return false;
    }
    if (coff.machine != IMAGE_FILE_MACHINE_AMD64) {
        outError = "Machine khong phai AMD64 (POE2 la tien trinh x64)";
        return false;
    }
    if (coff.numberOfSections == 0 || coff.numberOfSections > 96) {
        outError = "So luong section bat thuong";
        return false;
    }
    if (coff.sizeOfOptionalHeader < sizeof(uint16_t)) {
        outError = "sizeOfOptionalHeader qua nho";
        return false;
    }

    outBuild.machine = coff.machine;
    outBuild.numberOfSections = coff.numberOfSections;
    outBuild.timeDateStamp = coff.timeDateStamp;

    // Optional header PE32+: đọc magic rồi đọc SizeOfImage / SizeOfHeaders theo offset cố định
    const auto optionalStart = static_cast<uint64_t>(dos.e_lfanew) + sizeof(RawCoffHeader);
    uint16_t optionalMagic = 0;
    file.seekg(optionalStart, std::ios::beg);
    file.read(reinterpret_cast<char*>(&optionalMagic), sizeof(optionalMagic));
    if (optionalMagic != kPe32PlusMagic) {
        outError = "Optional header khong phai PE32+ (magic != 0x20B)";
        return false;
    }
    if (optionalStart + kSizeOfHeadersOffsetInOptHeader + sizeof(uint32_t) > fileSize) {
        outError = "Optional header tro ra ngoai pham vi tep";
        return false;
    }
    file.seekg(optionalStart + kSizeOfImageOffsetInOptHeader, std::ios::beg);
    file.read(reinterpret_cast<char*>(&outBuild.sizeOfImage), sizeof(outBuild.sizeOfImage));
    file.seekg(optionalStart + kSizeOfHeadersOffsetInOptHeader, std::ios::beg);
    file.read(reinterpret_cast<char*>(&outBuild.sizeOfHeaders), sizeof(outBuild.sizeOfHeaders));

    // Bảng section
    const auto sectionTableStart = optionalStart + coff.sizeOfOptionalHeader;
    file.seekg(sectionTableStart, std::ios::beg);
    outBuild.sections.reserve(coff.numberOfSections);
    for (uint16_t i = 0; i < coff.numberOfSections; ++i) {
        RawSectionHeader raw{};
        file.read(reinterpret_cast<char*>(&raw), sizeof(raw));
        if (file.gcount() != static_cast<std::streamsize>(sizeof(raw))) {
            outError = "Bang section bi cat ngan (thieu section #" + std::to_string(i) + ")";
            return false;
        }

        SectionInfo info;
        info.name = std::string(raw.name, strnlen(raw.name, sizeof(raw.name)));
        info.virtualAddress = raw.virtualAddress;
        info.virtualSize = raw.virtualSize;
        info.rawSize = raw.sizeOfRawData;
        info.rawPointer = raw.pointerToRawData;
        info.characteristics = raw.characteristics;
        outBuild.sections.push_back(std::move(info));
    }

    outBuild.valid = true;
    return true;
}

bool LoadPeHeadersWithHash(const std::string& path, ClientBuildId& outBuild, std::string& outError) {
    if (!LoadPeHeaders(path, outBuild, outError)) {
        return false;
    }
    if (!Sha256File(path, outBuild.fileSha256)) {
        outError = "Khong bam duoc SHA-256 cho tep: " + path;
        return false;
    }
    outBuild.sha256Computed = true;
    return true;
}

bool LocatePoe2Executable(std::string& outPath) {
    // Đồng bộ với src/assistant_tool/client_launcher.py DEFAULT_STANDALONE_PATHS
    static const char* kCandidates[] = {
        "C:\\Program Files (x86)\\Grinding Gear Games\\Path of Exile 2\\PathOfExile.exe",
        "C:\\Program Files (x86)\\Grinding Gear Games\\Path of Exile 2\\PathOfExile_x64.exe",
        "C:\\Program Files\\Grinding Gear Games\\Path of Exile 2\\PathOfExile.exe",
        "C:\\Program Files (x86)\\Steam\\steamapps\\common\\Path of Exile 2\\PathOfExile.exe",
        "C:\\Program Files (x86)\\Steam\\steamapps\\common\\Path of Exile 2\\PathOfExileSteam.exe",
        "D:\\Games\\Path of Exile 2\\PathOfExile.exe"
    };

    for (const char* candidate : kCandidates) {
        const DWORD attrs = GetFileAttributesA(candidate);
        if (attrs != INVALID_FILE_ATTRIBUTES && (attrs & FILE_ATTRIBUTE_DIRECTORY) == 0) {
            outPath = candidate;
            return true;
        }
    }
    return false;
}

} // namespace pe_fingerprint