// MemoryScanner.cpp : Définit les fonctions de la bibliothèque statique.
//

#include "Memory.hpp"
#include <psapi.h>
#include <sstream>
#include <spdlog/spdlog.h>
#include <spdlog/sinks/basic_file_sink.h>
#include <iomanip>

static std::shared_ptr<spdlog::logger> _log;
std::unordered_map<void*, Memory::PatchInfo> Memory::patches;

void Memory::PatchBytes(void* address, const char* bytes, size_t len)
{
    auto it = patches.find(address);
    if (it == patches.end())
    {
        // If a patch doesn't exist, create a new one.
        PatchInfo info;
        info.address = address;
        info.originalBytes.resize(len);
        memcpy(info.originalBytes.data(), address, len);
        // Store the patch info.
        patches[address] = info;
    }

    // Patch the bytes.
    DWORD oldProtect;
    VirtualProtect(address, len, PAGE_EXECUTE_READWRITE, &oldProtect);
    memcpy(address, bytes, len);
    VirtualProtect(address, len, oldProtect, &oldProtect);
}

//void Memory::RestoreBytes(uintptr_t address)
void Memory::RestoreBytes(void *address)
{
    auto it = patches.find(address);
    if (it != patches.end())
    {
        // Restore the original bytes.
        const auto& info = it->second;
        DWORD oldProtect;
        VirtualProtect(info.address, info.originalBytes.size(), PAGE_EXECUTE_READWRITE, &oldProtect);
        memcpy(info.address, info.originalBytes.data(), info.originalBytes.size());
        VirtualProtect(info.address, info.originalBytes.size(), oldProtect, &oldProtect);

        // Remove the patch info.
        patches.erase(it);
    }
}

bool Memory::wait_for_module(const std::string& module_name, int timeoutMs = 15000, int intervalMs = 500)
{
    const HANDLE hProc = GetCurrentProcess();

    for (int waited = 0; waited < timeoutMs; waited += intervalMs)
    {
        HMODULE hMods[1024];
        DWORD cbNeeded;

        if (EnumProcessModules(hProc, hMods, sizeof(hMods), &cbNeeded))
        {
            for (unsigned int i = 0; i < (cbNeeded / sizeof(HMODULE)); ++i)
            {
                char modName[MAX_PATH];
                if (GetModuleBaseNameA(hProc, hMods[i], modName, sizeof(modName)))
                {
                    if (_stricmp(modName, module_name.c_str()) == 0)
                    {
                        return true;
                    }
                }
            }
        }

        Sleep(intervalMs);
    }

    if (_log) _log->warn("Timeout: module '{}' not found in process after {} ms.", module_name, timeoutMs);
    return false;
}

std::string Memory::byteToHexEscaped(const BYTE byte) {
    std::ostringstream oss;
    oss << "\\x" << std::uppercase << std::hex << std::setw(2)
        << std::setfill('0') << static_cast<int>(byte);
    return oss.str();
}

uint8_t* Memory::aob_scan(
    const std::string& module_name,
    const std::string& signature,
    DWORD protect_flags = PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_READWRITE | PAGE_EXECUTE_WRITECOPY,
    std::shared_ptr<spdlog::logger> log) {

    _log = log;
    if (!wait_for_module(module_name))
    {
        if (log) log->warn("Skipping AOB scan because module '{}' is unavailable.", module_name);
        return nullptr;
    }

    // Convert signature to bytes
    std::vector<int> pattern_bytes;
    std::istringstream stream(signature);
    std::string byte_str;
    while (stream >> byte_str)
    {
        if (byte_str == "??" || byte_str == "?")
            pattern_bytes.push_back(-1);
        else
            pattern_bytes.push_back(static_cast<int>(std::strtol(byte_str.c_str(), nullptr, 16)));
    }

    HMODULE hMods[1024];
    DWORD cbNeeded;
    HANDLE hProc = GetCurrentProcess();

    if (!EnumProcessModules(hProc, hMods, sizeof(hMods), &cbNeeded))
    {
        spdlog::error("EnumProcessModules failed.");
        return nullptr;
    }

    for (unsigned int i = 0; i < (cbNeeded / sizeof(HMODULE)); ++i)
    {
        char modName[MAX_PATH];
        if (GetModuleBaseNameA(hProc, hMods[i], modName, sizeof(modName)))
        {
            if (_stricmp(modName, module_name.c_str()) == 0)
            {
                MODULEINFO modInfo;
                if (!GetModuleInformation(hProc, hMods[i], &modInfo, sizeof(modInfo)))
                {
                    if (log) log->error("GetModuleInformation failed for '{}'", module_name);
                    return nullptr;
                }

                uint8_t* base = reinterpret_cast<uint8_t*>(modInfo.lpBaseOfDll);
                size_t size = modInfo.SizeOfImage;
                if (log) log->info("Scanning memory region: 0x{:X} - 0x{:X}", reinterpret_cast<uintptr_t>(base), reinterpret_cast<uintptr_t>(base + size));

                MEMORY_BASIC_INFORMATION mbi{};
                for (uint8_t* current = base; current < base + size;)
                {
                    if (!VirtualQuery(current, &mbi, sizeof(mbi)))
                        break;

                    if ((mbi.State & MEM_COMMIT) && (mbi.Protect & protect_flags))
                    {
                        for (size_t i = 0; i <= mbi.RegionSize - pattern_bytes.size(); ++i)
                        {
                            bool match = true;
                            for (size_t j = 0; j < pattern_bytes.size(); ++j)
                            {
                                if (pattern_bytes[j] != -1 && current[i + j] != static_cast<uint8_t>(pattern_bytes[j]))
                                {
                                    match = false;
                                    break;
                                }
                            }

                            if (match)
                            {
                                uint8_t* result = current + i;
                                return result;
                            }
                        }
                    }

                    current = reinterpret_cast<uint8_t*>(mbi.BaseAddress) + mbi.RegionSize;
                }

                if (log) log->warn("No AOB match found in module '{}'.", module_name);
                return nullptr;
            }
        }
    }

    if (log) log->warn("Module '{}' unexpectedly disappeared during scan.", module_name);
    return nullptr;
}