// 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>
#include <tlhelp32.h>

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


const char* Memory::Float32ToHexBytes(float value) {
    static char bytes[4]; // buffer persistant (évite les problèmes de scope)
    std::memcpy(bytes, &value, sizeof(float));
    return bytes; // pointeur vers les 4 octets bruts
}

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(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::WaitForModule(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::AOBScan(
    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 (!WaitForModule(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;
}

PVOID Memory::SetupOrClearHardwareBreakPointForAllThreads(uintptr_t targetAddress, PVOID vehHandle, bool enable, PVECTORED_EXCEPTION_HANDLER pVEH, int hwIndex)
{
    DWORD pid = GetCurrentProcessId();
    HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if (snapshot == INVALID_HANDLE_VALUE) return nullptr;

    THREADENTRY32 te;
    te.dwSize = sizeof(te);

    // Add VectoredExceptionHandler
    if (enable && !vehHandle && pVEH)
    {
        vehHandle = AddVectoredExceptionHandler(1, pVEH);
    }

    if (Thread32First(snapshot, &te))
    {
        do
        {
            if (te.th32OwnerProcessID != pid) continue;

            HANDLE hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, te.th32ThreadID);
            if (!hThread) continue;

            CONTEXT ctx = {};
            ctx.ContextFlags = CONTEXT_DEBUG_REGISTERS;

            if (GetThreadContext(hThread, &ctx))
            {
                if (enable)
                {
                    switch (hwIndex) {
                    case 0: ctx.Dr0 = targetAddress; break; // Set Hardware breakpoint #1
                    case 1: ctx.Dr1 = targetAddress; break; // Set Hardware breakpoint #2
                    case 2: ctx.Dr2 = targetAddress; break; // Set Hardware breakpoint #3
                    case 3: ctx.Dr3 = targetAddress; break; // Set Hardware breakpoint #4
                    default: break;
                    }
                    ctx.Dr7 |= (1ULL << (hwIndex * 2)); // activate hardware breakpoint
                }
                else
                {
                    switch (hwIndex) {
                    case 0: ctx.Dr0 = 0; break; // Unset Hardware breakpoint #1
                    case 1: ctx.Dr1 = 0; break; // Unset Hardware breakpoint #2
                    case 2: ctx.Dr2 = 0; break; // Unset Hardware breakpoint #3
                    case 3: ctx.Dr3 = 0; break; // Unset Hardware breakpoint #4
                    default: break;
                    }
                    ctx.Dr7 &= ~(1ULL << (hwIndex * 2)); // deactivate hardware breakpoint
                }

                SetThreadContext(hThread, &ctx);
            }

            CloseHandle(hThread);

        } while (Thread32Next(snapshot, &te));
    }

    CloseHandle(snapshot);

    // Remove VectoredExceptionHandler
    if (!enable && vehHandle)
    {
        RemoveVectoredExceptionHandler(vehHandle);
        vehHandle = nullptr;
    }

    return vehHandle;
}

