#include "Readers/SlowEntityListReader.h"

#include <algorithm>
#include <cstring>

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

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

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

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

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

    // 1. Slow-entity list pointer
    uint64_t slowListPtr = 0;
    if (!mem.TryReadPointer(pid,
                            session.worldAddress + Offsets::World::SlowEntityList,
                            slowListPtr))
    {
        return false;
    }

    // 2. Alloc count (total slots) and valid count (live entities)
    int32_t rawAllocCount = 0;
    int32_t rawValidCount = 0;
    if (!mem.TryReadValue<int32_t>(pid,
                                   session.worldAddress + Offsets::World::SlowTableAllocCount,
                                   rawAllocCount))
    {
        return false;
    }
    mem.TryReadValue<int32_t>(pid,
                              session.worldAddress + Offsets::World::SlowTableValidCount,
                              rawValidCount);

    const int allocCount     = std::clamp(rawAllocCount, 0, kMaxEntities);
    const int safeValidCount = std::clamp(rawValidCount, 0, allocCount);

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

    // 4. Chunk size: max(batchSize, kMinTableChunkEntries) entries, clamped to remaining
    const int remaining    = allocCount - m_cursorOffset;
    const int chunkEntries = std::min(remaining,
                                      std::max(m_batchSize, kMinTableChunkEntries));

    // 5. Bulk-read chunk — each slot is Offsets::SlowTable::EntrySize (0x18) bytes
    const size_t entrySize  = static_cast<size_t>(Offsets::SlowTable::EntrySize);
    const size_t chunkBytes = static_cast<size_t>(chunkEntries) * entrySize;
    const uint64_t chunkAddr = slowListPtr +
                               static_cast<uint64_t>(m_cursorOffset) * entrySize;

    std::vector<uint8_t> tableBytes;
    tableBytes.reserve(chunkBytes);
    if (!mem.ReadBytes(pid, chunkAddr, chunkBytes, tableBytes)) {
        m_cursorOffset += chunkEntries;
        entries = BuildResult();
        return true;
    }

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

    for (int i = 0; i < chunkEntries; ++i) {
        const uint8_t* slot = tableBytes.data() + static_cast<size_t>(i) * entrySize;

        // Flag at offset 0 — uint16_t, valid == 1
        uint16_t flag = 0;
        std::memcpy(&flag, slot, sizeof(uint16_t));
        if (flag != 1) continue;

        // Entity pointer at offset 0x8
        uint64_t entityAddr = 0;
        std::memcpy(&entityAddr,
                    slot + static_cast<size_t>(Offsets::SlowTable::EntryPointerOffset),
                    sizeof(uint64_t));

        if (!MemoryValidation::IsValidUserAddress(entityAddr)) continue;

        ++m_validEntriesSeenInScan;

        auto now = std::chrono::steady_clock::now();
        ProcessEntity(mem, session, entityAddr, now);

        // Stop early when we have seen all valid entries
        if (safeValidCount > 0 && m_validEntriesSeenInScan >= safeValidCount) {
            m_cursorOffset += (i + 1);
            CompleteScan(now);
            PruneExpiredEntities(now);
            entries = BuildResult();
            return true;
        }

        // Check time budget
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now() - batchStart);
        if (elapsed >= m_readBudget) {
            m_cursorOffset += (i + 1);
            budgetExceeded  = true;
            break;
        }
    }

    if (!budgetExceeded) {
        m_cursorOffset += chunkEntries;
    }

    auto now = std::chrono::steady_clock::now();
    PruneExpiredEntities(now);

    if (m_cursorOffset >= allocCount) {
        CompleteScan(now);
    }

    entries = BuildResult();
    return true;
}

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

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

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

void SlowEntityListReader::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)) {
        auto it = m_entities.find(entityAddr);
        MarkTypeReadFailure(entityAddr,
                            (it != m_entities.end()) ? &it->second : 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())
        {
            auto it = m_entities.find(entityAddr);
            MarkTypeReadFailure(entityAddr,
                                (it != m_entities.end()) ? &it->second : nullptr,
                                now);
            return;
        }
        m_typeCache.Set(typeAddr, typeMeta);
    }

    // Skip entities with no model — slow list contains many invisible objects
    if (typeMeta.modelName.empty()) {
        return;
    }

    // Transform
    std::optional<Vector3> position;
    std::optional<float>   heading;
    bool transformOk = ReadEntityTransform(mem, session, entityAddr, position, heading);

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

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

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

    // Update cached entry
    cached.entry.address        = entityAddr;
    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 SlowEntityListReader::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_validEntriesSeenInScan = 0;
    m_hasCompletedScan       = true;
    ++m_scanGeneration;
}

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

void SlowEntityListReader::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<DayZSlowEntityEntry>& SlowEntityListReader::BuildResult()
{
    if (!m_resultDirty) {
        return m_cachedResult;
    }

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

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

void SlowEntityListReader::MarkTypeReadFailure(uint64_t                              /*entityAddr*/,
                                                CachedEntity*                         cached,
                                                std::chrono::steady_clock::time_point /*now*/)
{
    if (cached == nullptr) return;

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