#include "SigScanner/SigScanner.h"
#include "RuntimeOffsets.h"
#include "Offsets.h"

#include <algorithm>
#include <cctype>
#include <cstring>
#include <sstream>

#include <spdlog/spdlog.h>

// -----------------------------------------------------------------------
// Constants
// -----------------------------------------------------------------------

// Large sequential reads via Process::read() reliably fail on DayZ's module
// because its code/data sections are not contiguous in physical memory.
// Scatter reads handle non-contiguous pages natively and zero-fill failures,
// so we scan using one scatter entry per 4 KB OS page (matching VMM granularity).
static constexpr size_t kPageSize        = 0x1000;            // 4 KB — one OS page per scatter entry
static constexpr size_t kPagesPerBatch   = 256;               // 256 pages = 1 MB per scatter call
static constexpr size_t kBatchSize       = kPagesPerBatch * kPageSize;
static constexpr size_t kOverlap         = 512;               // cross-boundary safety (unchanged)
static constexpr size_t kDefaultScanSize = 64 * 1024 * 1024; // fallback if PE read fails

// -----------------------------------------------------------------------
// ParsePattern
// -----------------------------------------------------------------------
std::vector<SigScanner::PatternByte>
SigScanner::ParsePattern(const std::string& hex)
{
    std::vector<PatternByte> out;
    std::istringstream ss(hex);
    std::string token;
    while (ss >> token) {
        if (token == "?" || token == "??") {
            out.push_back({ 0x00, true });
        } else {
            uint8_t byte = static_cast<uint8_t>(std::stoul(token, nullptr, 16));
            out.push_back({ byte, false });
        }
    }
    return out;
}

// -----------------------------------------------------------------------
// FindFirst
// -----------------------------------------------------------------------
std::optional<size_t> SigScanner::FindFirst(
    const std::vector<uint8_t>&     haystack,
    const std::vector<PatternByte>& pattern,
    size_t                          startPos)
{
    if (pattern.empty()) return std::nullopt;
    const size_t patLen = pattern.size();
    if (haystack.size() < patLen) return std::nullopt;
    const size_t limit = haystack.size() - patLen;

    for (size_t i = startPos; i <= limit; ++i) {
        bool ok = true;
        for (size_t j = 0; j < patLen; ++j) {
            if (!pattern[j].wildcard && haystack[i + j] != pattern[j].value) {
                ok = false;
                break;
            }
        }
        if (ok) return i;
    }
    return std::nullopt;
}

// -----------------------------------------------------------------------
// ReadU32
// -----------------------------------------------------------------------
uint32_t SigScanner::ReadU32(const std::vector<uint8_t>& bytes, size_t offset)
{
    uint32_t v = 0;
    std::memcpy(&v, bytes.data() + offset, sizeof(v));
    return v;
}

// -----------------------------------------------------------------------
// ResolveRip  (MovCs)
// target = matchRva + instrLen + sign_extend(disp32)
// -----------------------------------------------------------------------
uint64_t SigScanner::ResolveRip(uint64_t matchRva, int instrLen, uint32_t disp32Raw)
{
    const int32_t  disp   = static_cast<int32_t>(disp32Raw);
    const int64_t  target = static_cast<int64_t>(matchRva)
                          + static_cast<int64_t>(instrLen)
                          + static_cast<int64_t>(disp);
    return static_cast<uint64_t>(target);
}

// -----------------------------------------------------------------------
// ExtractMovReg
// Find the first wildcard byte in the pattern; read uint32 from that
// position in the original chunk.  For DayZ, the 4 wildcard bytes that
// immediately follow the opcode prefix are the struct member offset.
// -----------------------------------------------------------------------
uint32_t SigScanner::ExtractMovReg(const std::vector<uint8_t>&     chunk,
                                    size_t                          matchPos,
                                    const std::vector<PatternByte>& pattern)
{
    for (size_t i = 0; i < pattern.size(); ++i) {
        if (pattern[i].wildcard) {
            return ReadU32(chunk, matchPos + i);
        }
    }
    return 0;
}

// -----------------------------------------------------------------------
// ReadModuleSize
// -----------------------------------------------------------------------
size_t SigScanner::ReadModuleSize(VmmAccessor& mem, uint32_t pid, uint64_t moduleBase)
{
    uint16_t magic = 0;
    if (!mem.TryReadValue<uint16_t>(pid, moduleBase, magic) || magic != 0x5A4D) return 0;

    uint32_t e_lfanew = 0;
    if (!mem.TryReadValue<uint32_t>(pid, moduleBase + 0x3C, e_lfanew)) return 0;

    uint32_t peSig = 0;
    if (!mem.TryReadValue<uint32_t>(pid, moduleBase + e_lfanew, peSig) || peSig != 0x00004550) return 0;

    // SizeOfImage: PE_header + 4 (PE sig) + 20 (COFF) + 0x38 (optional header offset)
    uint32_t sizeOfImage = 0;
    if (!mem.TryReadValue<uint32_t>(pid, moduleBase + e_lfanew + 24 + 0x38, sizeOfImage)) return 0;

    return static_cast<size_t>(sizeOfImage);
}

// -----------------------------------------------------------------------
// Scan
// -----------------------------------------------------------------------
SigScanner::ScanResult
SigScanner::Scan(VmmAccessor& mem, uint32_t pid, uint64_t moduleBase)
{
    ScanResult result;

    // ---- Determine scan range ----
    size_t scanSize = ReadModuleSize(mem, pid, moduleBase);
    if (scanSize == 0) {
        spdlog::warn("SigScanner: PE header unreadable — defaulting to {} MB",
                     kDefaultScanSize / (1024 * 1024));
        scanSize = kDefaultScanSize;
    } else {
        spdlog::debug("SigScanner: SizeOfImage = {} MB", scanSize / (1024 * 1024));
    }

    // ====================================================================
    // Pattern table — DayZ 1.29
    // ====================================================================

    // ---- MovCs: RIP-relative globals ----
    // Modbase::World  (48 8B 05 [disp32] 48 8D 54 24 ?? 48 8B 48 30)
    auto patWorld = ParsePattern(
        "48 8B 05 ? ? ? ? 48 8D 54 24 ? 48 8B 48 30");

    // Modbase::Network  (48 8D 0D [disp32] E8 ?? ?? ?? ?? 48 8B 1D ?? ?? ?? ?? 84 C0 74 20 41)
    auto patNetMgr = ParsePattern(
        "48 8D 0D ? ? ? ? E8 ? ? ? ? 48 8B 1D ? ? ? ? 84 C0 74 20 41");

    // ---- MovReg: World struct member offsets ----
    // World::NearEntList  (48 8B 83 [off32] 49 8B 14 06 48 3B D5)
    auto patNear = ParsePattern(
        "48 8B 83 ? ? ? ? 49 8B 14 06 48 3B D5");

    // World::FarEntList  (48 8B 83 [off32] 49 8B 0C 06 48 3B CD ...)
    auto patFar = ParsePattern(
        "48 8B 83 ? ? ? ? 49 8B 0C 06 48 3B CD 74 17 80 B9 ? ? ? ? ? 75 0E");

    // World::BulletList  (48 8B 83 [off32] 49 8B CF 48 03 0C F8)
    auto patBullet = ParsePattern(
        "48 8B 83 ? ? ? ? 49 8B CF 48 03 0C F8");

    // World::Camera  (4C 8B 83 [off32] 4C 8B 11 48 89 70 08)
    auto patCamera = ParsePattern(
        "4C 8B 83 ? ? ? ? 4C 8B 11 48 89 70 08");

    // World::ItemList  (48 89 85 [off32] 48 8D 15 ?? ?? ?? ?? 48 89 9D ?? ?? ?? ?? 48 8D 8D)
    auto patItem = ParsePattern(
        "48 89 85 ? ? ? ? 48 8D 15 ? ? ? ? 48 89 9D ? ? ? ? 48 8D 8D");

    // ---- MovReg: entity struct offsets ----
    // Human::VisualState  (48 8B 9F [off32] 49 8B CE FF 90 ?? ?? ?? ?? 8B 10)
    auto patVS = ParsePattern(
        "48 8B 9F ? ? ? ? 49 8B CE FF 90 ? ? ? ? 8B 10");

    // Human::IsDead  (0F B6 9A [off32] 48 89 7C 24 20 41 0F B6 B8 ?? ?? ?? ?? FF 90 B0 00)
    auto patDead = ParsePattern(
        "0F B6 9A ? ? ? ? 48 89 7C 24 20 41 0F B6 B8 ? ? ? ? FF 90 B0 00");

    // DayZPlayer::Inventory  (48 8B 8B [off32] 48 8B 01 FF 90 ?? ?? ?? ?? EB 02)
    auto patInv = ParsePattern(
        "48 8B 8B ? ? ? ? 48 8B 01 FF 90 ? ? ? ? EB 02");

    // DayZPlayerInventory::Hands  (48 8B 8B [off32] 48 8B F8 48 85 C9)
    auto patHands = ParsePattern(
        "48 8B 8B ? ? ? ? 48 8B F8 48 85 C9");

    // ====================================================================
    // Chunk scan — scatter-based, one 4 KB scatter entry per OS page.
    // ====================================================================
    std::vector<uint8_t> chunk;
    chunk.reserve(kBatchSize);

    std::vector<VmmAccessor::ScatterEntry> pages;
    pages.reserve(kPagesPerBatch);

    size_t offset = 0;
    while (offset < scanSize) {
        // Early-out once everything is found
        if (result.baseWorld        &&
            result.baseNetworkManager &&
            result.worldNearEntList &&
            result.worldFarEntList  &&
            result.worldBulletList  &&
            result.worldCamera      &&
            result.worldItemList    &&
            result.humanVisualState &&
            result.humanIsDead      &&
            result.playerInventory  &&
            result.inventoryHands)
        {
            break;
        }

        const size_t batchBytes = std::min(kBatchSize, scanSize - offset);
        const size_t numPages   = (batchBytes + kPageSize - 1) / kPageSize;

        // Zero-initialise: failed scatter pages stay zero and won't match any pattern.
        chunk.assign(batchBytes, 0);
        pages.clear();
        for (size_t p = 0; p < numPages; ++p) {
            const size_t pageOff = p * kPageSize;
            const size_t rdLen   = std::min(kPageSize, batchBytes - pageOff);
            pages.push_back({ moduleBase + offset + pageOff,
                              chunk.data() + pageOff,
                              rdLen });
        }
        mem.ScatterRead(pid, pages); // zero-fills any pages that fail to read

        // readSize drives the overlap-advance at the bottom of the loop.
        const size_t readSize = batchBytes;

        // -- MovCs --
        if (!result.baseWorld) {
            if (auto pos = FindFirst(chunk, patWorld)) {
                result.baseWorld = ResolveRip(offset + *pos, 7, ReadU32(chunk, *pos + 3));
            }
        }
        if (!result.baseNetworkManager) {
            if (auto pos = FindFirst(chunk, patNetMgr)) {
                result.baseNetworkManager = ResolveRip(offset + *pos, 7, ReadU32(chunk, *pos + 3));
            }
        }

        // -- MovReg (World struct) --
        if (!result.worldNearEntList) {
            if (auto pos = FindFirst(chunk, patNear)) {
                result.worldNearEntList = ExtractMovReg(chunk, *pos, patNear);
            }
        }
        if (!result.worldFarEntList) {
            if (auto pos = FindFirst(chunk, patFar)) {
                result.worldFarEntList = ExtractMovReg(chunk, *pos, patFar);
            }
        }
        if (!result.worldBulletList) {
            if (auto pos = FindFirst(chunk, patBullet)) {
                result.worldBulletList = ExtractMovReg(chunk, *pos, patBullet);
            }
        }
        if (!result.worldCamera) {
            if (auto pos = FindFirst(chunk, patCamera)) {
                result.worldCamera = ExtractMovReg(chunk, *pos, patCamera);
            }
        }
        if (!result.worldItemList) {
            if (auto pos = FindFirst(chunk, patItem)) {
                result.worldItemList = ExtractMovReg(chunk, *pos, patItem);
            }
        }

        // -- MovReg (entity struct) --
        if (!result.humanVisualState) {
            if (auto pos = FindFirst(chunk, patVS)) {
                result.humanVisualState = ExtractMovReg(chunk, *pos, patVS);
            }
        }
        if (!result.humanIsDead) {
            if (auto pos = FindFirst(chunk, patDead)) {
                result.humanIsDead = ExtractMovReg(chunk, *pos, patDead);
            }
        }
        if (!result.playerInventory) {
            if (auto pos = FindFirst(chunk, patInv)) {
                result.playerInventory = ExtractMovReg(chunk, *pos, patInv);
            }
        }
        if (!result.inventoryHands) {
            if (auto pos = FindFirst(chunk, patHands)) {
                result.inventoryHands = ExtractMovReg(chunk, *pos, patHands);
            }
        }

        if (readSize <= kOverlap) break;
        offset += readSize - kOverlap;
    }

    // ---- Report ----
    auto log = [](const char* name, auto val) {
        if (val) spdlog::info("SigScanner: {:35s} = 0x{:X}", name, *val);
        else     spdlog::warn("SigScanner: {:35s} NOT FOUND", name);
    };
    log("base::world",           result.baseWorld);
    log("base::network_manager", result.baseNetworkManager);
    log("World::NearEntityList", result.worldNearEntList);
    log("World::FarEntityList",  result.worldFarEntList);
    log("World::BulletTable",    result.worldBulletList);
    log("World::Camera",         result.worldCamera);
    log("World::ItemList",       result.worldItemList);
    log("Common::VisualState",   result.humanVisualState);
    log("Player::IsDead",        result.humanIsDead);
    log("Inventory::Base",       result.playerInventory);
    log("Inventory::Hands",      result.inventoryHands);

    return result;
}

// -----------------------------------------------------------------------
// Apply  — write results into RuntimeOffsets, log each change
// -----------------------------------------------------------------------
void SigScanner::Apply(const ScanResult& result)
{
    auto applyU64 = [](const char* name, uint64_t& dst, std::optional<uint64_t> src) {
        if (!src) return;
        if (*src != dst) {
            spdlog::info("SigScanner::Apply: {} 0x{:X} -> 0x{:X}", name, dst, *src);
            dst = *src;
        }
    };
    auto applyU32 = [](const char* name, uint64_t& dst, std::optional<uint32_t> src) {
        if (!src) return;
        if (static_cast<uint64_t>(*src) != dst) {
            spdlog::info("SigScanner::Apply: {} 0x{:X} -> 0x{:X}", name, dst, *src);
            dst = static_cast<uint64_t>(*src);
        }
    };

    applyU64("Base::World",          RuntimeOffsets::Base::World,          result.baseWorld);
    applyU64("Base::NetworkManager", RuntimeOffsets::Base::NetworkManager, result.baseNetworkManager);

    applyU32("World::NearEntityList", RuntimeOffsets::World::NearEntityList, result.worldNearEntList);
    applyU32("World::FarEntityList",  RuntimeOffsets::World::FarEntityList,  result.worldFarEntList);
    applyU32("World::BulletTable",    RuntimeOffsets::World::BulletTable,    result.worldBulletList);
    applyU32("World::Camera",         RuntimeOffsets::World::Camera,         result.worldCamera);
    applyU32("World::ItemList",       RuntimeOffsets::World::ItemList,       result.worldItemList);

    // Derive companion offsets that always follow the primary by +8
    if (result.worldNearEntList) {
        RuntimeOffsets::World::NearTableSize = RuntimeOffsets::World::NearEntityList + 8;
    }
    if (result.worldFarEntList) {
        RuntimeOffsets::World::FarTableSize = RuntimeOffsets::World::FarEntityList + 8;
    }
    if (result.worldBulletList) {
        RuntimeOffsets::World::BulletCount = RuntimeOffsets::World::BulletTable + 8;
    }
    if (result.worldItemList) {
        RuntimeOffsets::World::ItemTableAllocCount = RuntimeOffsets::World::ItemList + 8;
        RuntimeOffsets::World::ItemTableValidCount = RuntimeOffsets::World::ItemList + 16;
    }

    applyU32("Common::VisualState", RuntimeOffsets::Common::VisualState, result.humanVisualState);
    applyU32("Player::IsDead",      RuntimeOffsets::Player::IsDead,      result.humanIsDead);
    applyU32("Inventory::Base",     RuntimeOffsets::Inventory::Base,     result.playerInventory);
    applyU32("Inventory::Hands",    RuntimeOffsets::Inventory::Hands,    result.inventoryHands);
}
