#include "VmmAccessor.h"

// VmmAccessor.h includes <Windows.h> with WIN32_LEAN_AND_MEAN, which strips
// <winternl.h>.  vmmdll.h needs NTSTATUS from <winternl.h>, so pull it in
// explicitly here before including vmmdll.h.
#include <winternl.h>

// vmmdll.h is needed here only for VMMDLL_ConfigSet / VMMDLL_OPT_REFRESH_ALL.
#include "vmmdll.h"

#include <algorithm>  // std::find
#include <spdlog/spdlog.h>

// -------------------------------------------------------------------------
// Initialize
// -------------------------------------------------------------------------

bool VmmAccessor::Initialize(bool useMemoryMap) {
    std::unique_lock<std::shared_mutex> lk(m_accessMutex);
    m_process.reset();
    m_attachedPid = 0;
    spdlog::trace("VmmAccessor: initialising DMA (useMemoryMap={})...", useMemoryMap);
    m_dma = std::make_unique<DMA>(useMemoryMap);
    if (!m_dma->is_initialized()) {
        spdlog::error("VmmAccessor: DMA initialisation failed.");
        return false;
    }
    spdlog::info("VmmAccessor: DMA initialised.");
    return true;
}

// -------------------------------------------------------------------------
// ForceRefresh
// -------------------------------------------------------------------------

void VmmAccessor::ForceRefresh() {
    if (!IsInitialized()) return;
    VMMDLL_ConfigSet(m_dma->handle.get(), VMMDLL_OPT_REFRESH_ALL, 1);
    // VMMDLL_OPT_REFRESH_ALL resets per-process DTB back to the kernel default,
    // losing the fix_cr3 result.  Re-apply immediately so subsequent reads still
    // use the correct UserDirectoryTableBase paging context.
    if (m_process && !m_processName.empty()) {
        if (!m_process->fix_cr3(m_processName)) {
            spdlog::warn("VmmAccessor: fix_cr3 re-apply failed after ForceRefresh.");
        }
    }
}

// -------------------------------------------------------------------------
// TryFindProcess
// -------------------------------------------------------------------------

bool VmmAccessor::TryFindProcess(const std::string& processName,
                                  uint32_t&          outPid) {
    if (!IsInitialized()) return false;

    DWORD pid = m_dma->get_process_id(processName);
    if (pid == 0) return false;

    // Already attached to this exact PID — reuse the existing Process handle.
    // WaitForSession re-attaches every 3 s while waiting for spawn; without this
    // early-out we'd needlessly tear down a valid handle and re-run the CR3 fix
    // (and log it) on every one of those ticks.
    if (m_process && m_attachedPid == static_cast<uint32_t>(pid)) {
        outPid = m_attachedPid;
        return true;
    }

    // (Re-)create Process bound to this PID so all subsequent reads use it.
    // Unique lock: blocks the camera thread's reads while the handle is swapped
    // so it can never dereference a half-destroyed Process.
    std::unique_lock<std::shared_mutex> lk(m_accessMutex);
    m_process     = std::make_unique<Process>(*m_dma, processName);
    m_attachedPid = static_cast<uint32_t>(pid);
    m_processName = processName;
    outPid        = m_attachedPid;

    // Apply the CR3 fix so reads use UserDirectoryTableBase instead of the
    // kernel DTB.  Required for games (including DayZ) that run user-mode
    // code under a separate paging context.  No-op if already on the right DTB.
    if (!m_process->fix_cr3(processName)) {
        spdlog::warn("VmmAccessor: fix_cr3 failed for '{}' — reads may be unreliable.", processName);
    } else {
        spdlog::info("VmmAccessor: CR3 fix applied for '{}'.", processName);
    }

    spdlog::trace("VmmAccessor: attached to '{}' PID={}", processName, m_attachedPid);
    return true;
}

// -------------------------------------------------------------------------
// TryGetModuleBase
// -------------------------------------------------------------------------

bool VmmAccessor::TryGetModuleBase(uint32_t           /*pid*/,
                                    const std::string& moduleName,
                                    uint64_t&          outBase) {
    if (!IsAttached()) return false;

    uint64_t base = m_process->get_base_address(moduleName);
    if (base == 0) {
        spdlog::warn("VmmAccessor: module '{}' base not found.", moduleName);
        return false;
    }
    spdlog::trace("VmmAccessor: module '{}' base=0x{:X}", moduleName, base);
    outBase = base;
    return true;
}

// -------------------------------------------------------------------------
// IsProcessAlive
// -------------------------------------------------------------------------

bool VmmAccessor::IsProcessAlive(uint32_t           pid,
                                  const std::string& processName) {
    if (!IsInitialized()) return false;

    DWORD livePid = m_dma->get_process_id(processName);
    return livePid != 0 && static_cast<uint32_t>(livePid) == pid;
}

// -------------------------------------------------------------------------
// ReadRaw (private)
// -------------------------------------------------------------------------

bool VmmAccessor::ReadRaw(uint32_t /*pid*/, uint64_t address,
                           void* buf, size_t size) {
    if (!buf || size == 0) return false;
    std::shared_lock<std::shared_mutex> lk(m_accessMutex);
    if (!m_process) return false;
    bool ok = m_process->read(address, buf, size);
    if (ok) m_bytesReadTotal.fetch_add(size, std::memory_order_relaxed);
    return ok;
}

// -------------------------------------------------------------------------
// TryReadPointer
// -------------------------------------------------------------------------

bool VmmAccessor::TryReadPointer(uint32_t pid, uint64_t address,
                                   uint64_t& outPointer) {
    if (!MemoryValidation::IsValidUserAddress(address)) return false;

    uint64_t ptr = 0;
    if (!ReadRaw(pid, address, &ptr, sizeof(ptr))) return false;

    if (!MemoryValidation::IsValidUserAddress(ptr)) return false;

    outPointer = ptr;
    return true;
}

// -------------------------------------------------------------------------
// ScatterRead
// -------------------------------------------------------------------------

bool VmmAccessor::ScatterRead(uint32_t pid, std::vector<ScatterEntry>& entries) {
    if (entries.empty()) return false;
    std::shared_lock<std::shared_mutex> lk(m_accessMutex);
    if (!m_process) return false;

    VMMDLL_SCATTER_HANDLE hScatter = m_process->create_scatter(static_cast<DWORD>(pid));
    if (!hScatter) return false;

    for (auto& e : entries) {
        if (e.buffer && e.size > 0 && MemoryValidation::IsValidUserAddress(e.address))
            m_process->add_read_scatter(hScatter, e.address, e.buffer, e.size);
    }

    m_process->execute_scatter(hScatter, static_cast<DWORD>(pid));
    m_process->close_scatter(hScatter);

    // Account for all scheduled bytes (VolkDMA zero-fills failures; total is still representative).
    size_t totalBytes = 0;
    for (const auto& e : entries) totalBytes += e.size;
    m_bytesReadTotal.fetch_add(totalBytes, std::memory_order_relaxed);
    m_scatterOpsTotal.fetch_add(1, std::memory_order_relaxed);
    return true;
}

// -------------------------------------------------------------------------
// ReadBytes
// -------------------------------------------------------------------------

bool VmmAccessor::ReadBytes(uint32_t              pid,
                             uint64_t              address,
                             size_t                count,
                             std::vector<uint8_t>& outBytes) {
    if (count == 0) {
        outBytes.clear();
        return true;
    }
    outBytes.resize(count);
    if (!ReadRaw(pid, address, outBytes.data(), count)) {
        outBytes.clear();
        return false;
    }
    return true;
}

// -------------------------------------------------------------------------
// DecodeArmaBytes (static, private)
// -------------------------------------------------------------------------

std::string VmmAccessor::DecodeArmaBytes(const std::vector<uint8_t>& bytes) {
    if (bytes.empty()) return {};

    // Determine effective length (stop at first null terminator).
    size_t len = 0;
    while (len < bytes.size() && bytes[len] != 0) ++len;
    if (len == 0) return {};

    // --- Attempt 1: validate as UTF-8 ---
    auto isValidUtf8 = [](const uint8_t* data, size_t n) -> bool {
        size_t i = 0;
        while (i < n) {
            uint8_t b = data[i];
            int seqLen = 0;
            if      ((b & 0x80) == 0x00) seqLen = 1;
            else if ((b & 0xE0) == 0xC0) seqLen = 2;
            else if ((b & 0xF0) == 0xE0) seqLen = 3;
            else if ((b & 0xF8) == 0xF0) seqLen = 4;
            else return false;
            if (i + seqLen > n) return false;
            for (int j = 1; j < seqLen; ++j) {
                if ((data[i + j] & 0xC0) != 0x80) return false;
            }
            i += seqLen;
        }
        return true;
    };

    if (isValidUtf8(bytes.data(), len)) {
        return std::string(reinterpret_cast<const char*>(bytes.data()), len);
    }

    // --- Attempt 2: treat as Windows CP1251, convert to UTF-8 ---
    int wideLen = MultiByteToWideChar(1251 /*CP1251*/, 0,
                                       reinterpret_cast<LPCCH>(bytes.data()),
                                       static_cast<int>(len),
                                       nullptr, 0);
    if (wideLen <= 0) {
        return std::string(reinterpret_cast<const char*>(bytes.data()), len);
    }

    std::wstring wide(static_cast<size_t>(wideLen), L'\0');
    MultiByteToWideChar(1251, 0,
                        reinterpret_cast<LPCCH>(bytes.data()),
                        static_cast<int>(len),
                        wide.data(), wideLen);

    int utf8Len = WideCharToMultiByte(CP_UTF8, 0,
                                       wide.data(), wideLen,
                                       nullptr, 0,
                                       nullptr, nullptr);
    if (utf8Len <= 0) {
        return std::string(reinterpret_cast<const char*>(bytes.data()), len);
    }

    std::string utf8(static_cast<size_t>(utf8Len), '\0');
    WideCharToMultiByte(CP_UTF8, 0,
                        wide.data(), wideLen,
                        utf8.data(), utf8Len,
                        nullptr, nullptr);
    return utf8;
}

// -------------------------------------------------------------------------
// ReadArmaString
// -------------------------------------------------------------------------

std::string VmmAccessor::ReadArmaString(uint32_t pid, uint64_t address,
                                          size_t maxLength) {
    if (!MemoryValidation::IsValidUserAddress(address)) return {};

    // Read the 16-bit length field at address + 0x8
    uint16_t strLen = 0;
    if (!ReadRaw(pid, address + 0x8, &strLen, sizeof(strLen))) return {};
    if (strLen == 0) return {};

    size_t readLen = static_cast<size_t>(strLen);
    if (readLen > maxLength) readLen = maxLength;

    // Read the character data at address + 0x10
    std::vector<uint8_t> buf(readLen);
    if (!ReadRaw(pid, address + 0x10, buf.data(), readLen)) return {};

    return DecodeArmaBytes(buf);
}

// -------------------------------------------------------------------------
// ReadCString
// -------------------------------------------------------------------------

// -------------------------------------------------------------------------
// GetStats
// -------------------------------------------------------------------------

DmaStats VmmAccessor::GetStats() const {
    using Clock = std::chrono::steady_clock;

    const uint64_t bytes = m_bytesReadTotal.load(std::memory_order_relaxed);
    const uint64_t ops   = m_scatterOpsTotal.load(std::memory_order_relaxed);

    const auto now = Clock::now();
    const double dt = std::chrono::duration<double>(now - m_statsLastSample).count();

    // Recompute at most 10 Hz to avoid thrashing on high-frequency overlay calls.
    if (dt >= 0.1) {
        const double bytesDelta = static_cast<double>(bytes - m_statsLastBytes);
        const double opsDelta   = static_cast<double>(ops   - m_statsLastOps);
        m_statsMBps  = static_cast<float>(bytesDelta / dt / (1024.0 * 1024.0));
        m_statsOpsPs = static_cast<float>(opsDelta   / dt);
        m_statsLastSample = now;
        m_statsLastBytes  = bytes;
        m_statsLastOps    = ops;
    }

    DmaStats s;
    s.attached     = IsAttached();
    s.readMBps     = m_statsMBps;
    s.scatterOpsPs = m_statsOpsPs;
    s.totalGB      = static_cast<float>(bytes) / (1024.f * 1024.f * 1024.f);
    s.totalOps     = ops;
    return s;
}

std::string VmmAccessor::ReadCString(uint32_t pid, uint64_t address,
                                       size_t maxLength) {
    if (!MemoryValidation::IsValidUserAddress(address)) return {};
    if (maxLength == 0 || !IsAttached()) return {};

    std::vector<uint8_t> buf(maxLength);
    if (!ReadRaw(pid, address, buf.data(), maxLength)) {
        // Partial reads: try progressively smaller chunks down to 1 page.
        // If even that fails, give up.
        if (maxLength <= 64 || !ReadRaw(pid, address, buf.data(), 64)) {
            return {};
        }
        buf.resize(64);
    }

    // Truncate at null terminator if present.
    auto nullPos = std::find(buf.begin(), buf.end(), uint8_t{0});
    buf.erase(nullPos, buf.end());
    if (buf.empty()) return {};

    return DecodeArmaBytes(buf);
}
