#include "Readers/FarEntityListReader.h"

#include <algorithm>
#include <cstring>

#include "Offsets.h"
#include "RuntimeOffsets.h"
#include "Memory/MemoryValidation.h"
#include "Memory/VmmAccessor.h"

// -------------------------------------------------------------------------
// Constructor
// -------------------------------------------------------------------------

FarEntityListReader::FarEntityListReader(int batchSize,
                                         int readBudgetMs,
                                         int staleAfterMs)
    : m_batchSize(batchSize)
    , m_readBudget(readBudgetMs)
    , m_staleAfter(staleAfterMs)
{}

// -------------------------------------------------------------------------
// TryRead
// -------------------------------------------------------------------------

bool FarEntityListReader::TryRead(VmmAccessor&                     mem,
                                  const RuntimeSession&            session,
                                  std::vector<DayZFarEntityEntry>& entries)
{
    const uint32_t pid = session.processId;

    // 1. Far-entity list pointer
    uint64_t farListPtr = 0;
    if (!mem.TryReadPointer(pid,
                            session.worldAddress + RuntimeOffsets::World::FarEntityList,
                            farListPtr))
    {
        return false;
    }

    // 2. Table size
    int32_t rawCount = 0;
    if (!mem.TryReadValue<int32_t>(pid,
                                   session.worldAddress + RuntimeOffsets::World::FarTableSize,
                                   rawCount))
    {
        return false;
    }

    const int count = std::clamp(rawCount, 0, kMaxEntities);

    // 3. If cursor has reached the end of the table, complete the current scan
    if (m_cursorOffset >= count * 8) {
        auto now = std::chrono::steady_clock::now();
        CompleteScan(now);
        PruneExpiredEntities(now);
        entries = BuildResult();
        return true;
    }

    // 4. Calculate chunk: how many pointer-sized slots remain vs batch limit
    const int remaining    = count - (m_cursorOffset / 8);
    const int chunkEntries = std::min(remaining, m_batchSize);

    // 5. Bulk-read chunk of pointers
    const size_t chunkBytes = static_cast<size_t>(chunkEntries) * 8;
    std::vector<uint8_t> tableBytes;
    tableBytes.reserve(chunkBytes);
    if (!mem.ReadBytes(pid, farListPtr + static_cast<uint64_t>(m_cursorOffset),
                       chunkBytes, tableBytes))
    {
        // On failure advance cursor to avoid getting stuck
        m_cursorOffset += static_cast<int>(chunkBytes);
        entries = BuildResult();
        return true;
    }

    // 6. Process entities within the read budget
    const auto batchStart = std::chrono::steady_clock::now();

    for (int i = 0; i < chunkEntries; ++i) {
        uint64_t entityAddr = 0;
        std::memcpy(&entityAddr,
                    tableBytes.data() + static_cast<size_t>(i) * 8,
                    sizeof(uint64_t));

        if (MemoryValidation::IsValidUserAddress(entityAddr)) {
            auto now = std::chrono::steady_clock::now();
            ProcessEntity(mem, session, entityAddr, now);

            // Check time budget after each entity
            auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now() - batchStart);
            if (elapsed >= m_readBudget) {
                // Advance cursor to account for entities processed so far
                m_cursorOffset += (i + 1) * 8;
                PruneExpiredEntities(std::chrono::steady_clock::now());
                entries = BuildResult();
                return true;
            }
        }
    }

    // 7. Advance cursor by the full chunk
    m_cursorOffset += chunkEntries * 8;

    // 8. Prune stale entries
    auto now = std::chrono::steady_clock::now();
    PruneExpiredEntities(now);

    // 9. If cursor reached end, complete scan
    if (m_cursorOffset >= count * 8) {
        CompleteScan(now);
    }

    entries = BuildResult();
    return true;
}

// -------------------------------------------------------------------------
// Reset
// -------------------------------------------------------------------------

void FarEntityListReader::Reset()
{
    m_entities.clear();
    m_typeCache.Clear();
    m_cachedResult.clear();
    m_sortedAddresses.clear();
    m_cursorOffset     = 0;
    m_scanGeneration   = 1;
    m_hasCompletedScan = false;
    m_sortDirty        = true;
    m_resultDirty      = true;
}

// -------------------------------------------------------------------------
// ProcessEntity
// -------------------------------------------------------------------------

void FarEntityListReader::ProcessEntity(VmmAccessor&                          mem,
                                        const RuntimeSession&                 session,
                                        uint64_t                              entityAddr,
                                        std::chrono::steady_clock::time_point now)
{
    const uint32_t pid = session.processId;

    // Read entity type pointer
    uint64_t typeAddr = 0;
    if (!mem.TryReadPointer(pid, entityAddr + Offsets::Common::Type, typeAddr)) {
        MarkTypeReadFailure(entityAddr, nullptr, now);
        return;
    }

    // Look up or read type metadata
    EntityTypeMetadata typeMeta;
    if (!m_typeCache.GetOrRead(typeAddr, typeMeta)) {
        typeMeta = ReadEntityTypeMetadata(mem, pid, typeAddr);
        if (typeMeta.typeName.empty() && typeMeta.configName.empty() &&
            typeMeta.cleanName.empty() && typeMeta.modelName.empty())
        {
            // Type read produced nothing useful — track failure on cached entity
            auto it = m_entities.find(entityAddr);
            MarkTypeReadFailure(entityAddr,
                                (it != m_entities.end()) ? &it->second : nullptr,
                                now);
            return;
        }
        m_typeCache.Set(typeAddr, typeMeta);
    }

    // Network ID
    uint32_t networkId = 0;
    mem.TryReadValue<uint32_t>(pid, entityAddr + Offsets::Network::EntityNetworkId, networkId);

    // Transform — use cached position/heading if the read fails.
    // Also cache the VS pointer for use in RefreshPositions.
    std::optional<Vector3> position;
    std::optional<float>   heading;

    uint64_t vsAddr = 0;
    mem.TryReadValue<uint64_t>(pid, entityAddr + Offsets::Common::VisualState, vsAddr);

    bool transformOk = ReadEntityTransform(mem, session, entityAddr, position, heading);

    // Retrieve or create cached entry
    auto& cached = m_entities[entityAddr];

    if (vsAddr) cached.visualStateAddr = vsAddr;

    if (!transformOk && cached.entry.position.has_value()) {
        position = cached.entry.position;
        heading  = cached.entry.headingDegrees;
    }

    // Build display name
    const std::string entityName = GetBestEntityName(typeMeta.typeName,
                                                      typeMeta.configName,
                                                      typeMeta.cleanName,
                                                      typeMeta.modelName,
                                                      entityAddr);

    // Update cached entry
    cached.entry.address        = entityAddr;
    cached.entry.networkId      = networkId;
    cached.entry.position       = position;
    cached.entry.headingDegrees = heading;
    cached.entry.entityName     = entityName;
    cached.entry.typeName       = typeMeta.typeName;
    cached.entry.configName     = typeMeta.configName;
    cached.entry.modelName      = typeMeta.modelName;
    cached.lastSeenGeneration   = m_scanGeneration;
    cached.failedTypeReads      = 0;

    if (transformOk) {
        cached.lastSuccessfulRead = now;
    }

    m_resultDirty = true;
}

// -------------------------------------------------------------------------
// CompleteScan
// -------------------------------------------------------------------------

void FarEntityListReader::CompleteScan(std::chrono::steady_clock::time_point /*now*/)
{
    // Remove entities not seen in the current scan generation
    for (auto it = m_entities.begin(); it != m_entities.end(); ) {
        if (it->second.lastSeenGeneration != m_scanGeneration) {
            it = m_entities.erase(it);
            m_sortDirty   = true;
            m_resultDirty = true;
        } else {
            ++it;
        }
    }

    m_cursorOffset     = 0;
    m_hasCompletedScan = true;
    ++m_scanGeneration;
}

// -------------------------------------------------------------------------
// PruneExpiredEntities
// -------------------------------------------------------------------------

void FarEntityListReader::PruneExpiredEntities(std::chrono::steady_clock::time_point now)
{
    for (auto it = m_entities.begin(); it != m_entities.end(); ) {
        const auto age = std::chrono::duration_cast<std::chrono::milliseconds>(
            now - it->second.lastSuccessfulRead);
        if (age > m_staleAfter) {
            it = m_entities.erase(it);
            m_sortDirty   = true;
            m_resultDirty = true;
        } else {
            ++it;
        }
    }
}

// -------------------------------------------------------------------------
// BuildResult
// -------------------------------------------------------------------------

std::vector<DayZFarEntityEntry>& FarEntityListReader::BuildResult()
{
    if (!m_resultDirty) {
        return m_cachedResult;
    }

    // Rebuild sorted address list when topology changes
    if (m_sortDirty) {
        m_sortedAddresses.clear();
        m_sortedAddresses.reserve(m_entities.size());
        for (const auto& kv : m_entities) {
            m_sortedAddresses.push_back(kv.first);
        }
        std::sort(m_sortedAddresses.begin(), m_sortedAddresses.end(),
                  [this](uint64_t a, uint64_t b) {
                      const auto& ea = m_entities.at(a).entry;
                      const auto& eb = m_entities.at(b).entry;
                      if (ea.entityName != eb.entityName)
                          return ea.entityName < eb.entityName;
                      return a < b;
                  });
        m_sortDirty = false;
    }

    m_cachedResult.clear();
    m_cachedResult.reserve(m_sortedAddresses.size());
    for (uint64_t addr : m_sortedAddresses) {
        auto it = m_entities.find(addr);
        if (it != m_entities.end()) {
            m_cachedResult.push_back(it->second.entry);
        }
    }

    m_resultDirty = false;
    return m_cachedResult;
}

// -------------------------------------------------------------------------
// RefreshPositions
// Two-pass scatter read: all VS pointers in pass A, all positions in pass B.
// Entities with a stale/unknown VS pointer fall back to a sequential re-read.
// -------------------------------------------------------------------------

void FarEntityListReader::RefreshPositions(VmmAccessor&          mem,
                                            const RuntimeSession& session)
{
    if (m_entities.empty()) return;

    const uint32_t pid = session.processId;

    // Collect entities that have a cached VS pointer and those that don't.
    struct PosSlot {
        uint64_t entityAddr;
        uint64_t vsAddr;
        Vector3  posOut;
        bool     ok = false;
    };

    std::vector<PosSlot> slots;
    slots.reserve(m_entities.size());

    // Pass A: scatter-read VS pointers for entities whose cached value is 0.
    {
        std::vector<VmmAccessor::ScatterEntry> vsScatter;
        // We need a stable backing store for the scatter buffers.
        // Use a parallel vector of uint64_t.
        std::vector<uint64_t> vsBufs(m_entities.size(), 0);
        int idx = 0;
        for (auto& [addr, cached] : m_entities) {
            PosSlot s;
            s.entityAddr = addr;
            s.vsAddr     = cached.visualStateAddr;
            slots.push_back(s);

            if (cached.visualStateAddr == 0) {
                vsScatter.push_back({ addr + Offsets::Common::VisualState,
                                      &vsBufs[idx], sizeof(uint64_t) });
            }
            ++idx;
        }

        if (!vsScatter.empty()) {
            mem.ScatterRead(pid, vsScatter);
        }

        // Fill cached VS pointers from scatter result.
        idx = 0;
        for (auto& [addr, cached] : m_entities) {
            if (cached.visualStateAddr == 0 && vsBufs[idx] != 0) {
                cached.visualStateAddr = vsBufs[idx];
                slots[idx].vsAddr      = vsBufs[idx];
            }
            ++idx;
        }
    }

    // Pass B: scatter-read positions for all slots that have a VS pointer.
    {
        std::vector<VmmAccessor::ScatterEntry> posScatter;
        posScatter.reserve(slots.size());

        for (auto& slot : slots) {
            if (slot.vsAddr == 0) continue;
            posScatter.push_back({ slot.vsAddr + Offsets::Common::Position,
                                   &slot.posOut, sizeof(Vector3) });
        }

        if (!posScatter.empty()) {
            mem.ScatterRead(pid, posScatter);
        }
    }

    // Apply results: update position on cached entities.
    auto now = std::chrono::steady_clock::now();
    for (auto& slot : slots) {
        if (slot.vsAddr == 0) continue;
        if (!MemoryValidation::IsValidVector(slot.posOut)) {
            // Possibly stale VS pointer — invalidate so pass A re-reads it next time.
            auto it = m_entities.find(slot.entityAddr);
            if (it != m_entities.end())
                it->second.visualStateAddr = 0;
            continue;
        }

        auto it = m_entities.find(slot.entityAddr);
        if (it != m_entities.end()) {
            it->second.entry.position  = slot.posOut;
            it->second.lastSuccessfulRead = now;
            m_resultDirty = true;
        }
    }
}

// -------------------------------------------------------------------------
// MarkTypeReadFailure
// -------------------------------------------------------------------------

void FarEntityListReader::MarkTypeReadFailure(uint64_t                              entityAddr,
                                               CachedEntity*                         cached,
                                               std::chrono::steady_clock::time_point now)
{
    if (cached == nullptr) {
        // Entity not yet in cache — nothing more to do
        return;
    }
    ++cached->failedTypeReads;

    // Still mark it seen this generation so it isn't pruned by CompleteScan
    cached->lastSeenGeneration = m_scanGeneration;
}
