// 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;

uint8_t* Memory::GetOffsetFromOpcode(uint8_t* opcode, int extraOffset)
{
    if (!opcode)
        return nullptr;

    int32_t disp = 0;
    std::memcpy(&disp, opcode, sizeof(int32_t));

    if (disp < 0)
        return nullptr; // optionnel : gérer ou pas les offsets négatifs

    // Retourne l'adresse "offsetée" (base + disp)
    return opcode + 4 + disp + extraOffset; // +4 car disp32 fait 4 octets
}

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
}

std::vector<std::uint8_t> Memory::ReadBytes(const void* addr, std::size_t size) {
    std::vector<std::uint8_t> buffer(size);
    std::memcpy(buffer.data(), addr, size);
    return buffer;
}

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);
    }
}

MODULEINFO Memory::WaitForModule(const std::string& module_name, int timeoutMs, int intervalMs)
{
    const HANDLE hProc = GetCurrentProcess();
    MODULEINFO modInfo{};

    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) {
                        if (GetModuleInformation(hProc, hMods[i], &modInfo, sizeof(modInfo)))
                            return modInfo;
                    }
                }
            }
        }

        Sleep(intervalMs);
    }

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

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) {

    bool found = false;

    _log = log;
    HANDLE hProc = GetCurrentProcess();
    MODULEINFO modInfo{};
    HMODULE targetModule = nullptr;
    // Get module when name is specidifed
    if (!(module_name.empty() || module_name == "*")) {
        if (_log) _log->info("Module name: {}", module_name);
        MODULEINFO modinfo = WaitForModule(module_name);
        if (modinfo.lpBaseOfDll == nullptr) {
            if (_log) _log->warn("Skipping AOB scan because module '{}' is unavailable.", module_name);
            return nullptr;
        }
    }
    // Fallback to determine module loaded
    if (!found || module_name.empty() || module_name == "*") {
        char exeBuf[MAX_PATH] = { 0 };
        DWORD exeLen = GetModuleFileNameA(nullptr, exeBuf, MAX_PATH);
        std::string exeName = (exeLen > 0) ? std::string(exeBuf, exeBuf + exeLen) : std::string();
        size_t pos = exeName.find_last_of("\\/");

        if (pos != std::string::npos) exeName = exeName.substr(pos + 1);
        if (_log && !exeName.empty()) _log->info("Module name: {}", exeName);

        targetModule = GetModuleHandleA(nullptr);
        if (!targetModule || !GetModuleInformation(hProc, targetModule, &modInfo, sizeof(modInfo))) {
            if (_log) _log->error("Failed to find main module.");
            return nullptr;
        }
    }
    // Convert AOB string into vector 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)));
    }

    if (pattern_bytes.empty()) {
        if (_log) _log->warn("Empty AOB pattern passed.");
        return nullptr;
    }
    // Logging scanning area
    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 scan
    MEMORY_BASIC_INFORMATION mbi{};
    for (uint8_t* current = base; current < base + size;)
    {
        if (!VirtualQuery(current, &mbi, sizeof(mbi)))
            break;

        bool isCommitted = (mbi.State & MEM_COMMIT) != 0;
        bool hasAccess = (mbi.Protect & protect_flags) != 0;
        bool isNoAccess = (mbi.Protect & PAGE_NOACCESS) != 0;
        bool isGuard = (mbi.Protect & PAGE_GUARD) != 0;

        if (isCommitted && hasAccess && !isNoAccess && !isGuard) {
            uint8_t* regionBase = reinterpret_cast<uint8_t*>(mbi.BaseAddress);
            size_t regionSize = mbi.RegionSize;

            for (size_t i = 0; i <= regionSize - pattern_bytes.size(); ++i)
            {
                bool match = true;
                for (size_t j = 0; j < pattern_bytes.size(); ++j)
                {
                    if (pattern_bytes[j] != -1 && regionBase[i + j] != static_cast<uint8_t>(pattern_bytes[j])) {
                        match = false;
                        break;
                    }
                }

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

        current = reinterpret_cast<uint8_t*>(mbi.BaseAddress) + mbi.RegionSize;
    }
    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;
}