#include "WorldStepProcessor.hpp"
#include "world/Transform.hpp"

namespace tw::p2p {

WorldStepProcessor::WorldStepProcessor(World* world, uint32_t self_id)
    : m_world(world)
    , m_self_id(self_id)
{
    m_all_peer_ids.insert(self_id);
    auto entity = m_world->registry().create();
    m_world->registry().emplace<tw::Transform>(entity, glm::vec3(0.0f));
    m_peer_entities[self_id] = entity;
    m_committed_snapshot.positions[self_id] = glm::vec3(0.0f);
}

void WorldStepProcessor::add_peer(uint32_t peer_id) {
    if (m_all_peer_ids.count(peer_id)) return;
    m_all_peer_ids.insert(peer_id);
    auto entity = m_world->registry().create();
    m_world->registry().emplace<tw::Transform>(entity, glm::vec3(0.0f));
    m_peer_entities[peer_id] = entity;
    m_committed_snapshot.positions[peer_id] = glm::vec3(0.0f);
}

// ── submit ────────────────────────────────────────────────────────────────────

void WorldStepProcessor::submit_action(const PeerAction& action) {
    m_frames[action.frame_idx].actions[action.peer_id] = action;
}

// ── advance ───────────────────────────────────────────────────────────────────

bool WorldStepProcessor::advance() {
    flush_commits();

    Frame empty;
    const Frame& frame = m_frames.count(m_current_frame)
                             ? m_frames.at(m_current_frame)
                             : empty;
    apply_frame(frame);
    ++m_current_frame;
    ++m_stepped_frames;
    return true;
}

// ── flush_commits ─────────────────────────────────────────────────────────────

void WorldStepProcessor::flush_commits() {
    const uint32_t old_committed = m_committed_frame;

    while (m_frames.count(m_committed_frame) &&
           m_frames.at(m_committed_frame).is_full(m_all_peer_ids))
    {
        ++m_committed_frame;
    }

    if (m_committed_frame == old_committed) {
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
        return;
    }

    // Advance the committed snapshot through the newly committed frames.
    restore_snapshot(m_committed_snapshot);
    for (uint32_t f = old_committed; f < m_committed_frame; ++f)
        apply_frame(m_frames.at(f));
    m_committed_snapshot = capture_snapshot();

    // Re-simulate any speculative frames on top of the new committed state.
    for (uint32_t f = m_committed_frame; f < m_current_frame; ++f) {
        Frame empty;
        const Frame& fr = m_frames.count(f) ? m_frames.at(f) : empty;
        apply_frame(fr);
    }

    // Prune frames that are now part of the committed snapshot.
    while (!m_frames.empty() && m_frames.begin()->first < m_committed_frame)
        m_frames.erase(m_frames.begin());

    ++m_rollbacks;
}

// ── helpers ───────────────────────────────────────────────────────────────────

void WorldStepProcessor::apply_frame(const Frame& frame) {
    for (uint32_t peer_id : m_all_peer_ids) {
        glm::vec3 input(0.0f);
        auto it = frame.actions.find(peer_id);
        if (it != frame.actions.end())
            input = it->second.input;

        auto* t = m_world->registry().try_get<tw::Transform>(m_peer_entities.at(peer_id));
        if (t) t->translate(input * MOVE_SPEED * static_cast<float>(FIXED_DELTA_S));
    }
    m_world->step(FIXED_DELTA_S);
}

WorldStepProcessor::Snapshot WorldStepProcessor::capture_snapshot() const {
    Snapshot snap;
    for (auto& [peer_id, entity] : m_peer_entities) {
        const auto* t = m_world->registry().try_get<tw::Transform>(entity);
        snap.positions[peer_id] = t ? t->position() : glm::vec3(0.0f);
    }
    return snap;
}

void WorldStepProcessor::restore_snapshot(const Snapshot& snap) {
    for (auto& [peer_id, pos] : snap.positions) {
        auto it = m_peer_entities.find(peer_id);
        if (it == m_peer_entities.end()) continue;
        auto* t = m_world->registry().try_get<tw::Transform>(it->second);
        if (t) t->set_position(pos);
    }
}

glm::vec3 WorldStepProcessor::peer_position(uint32_t peer_id) const {
    auto it = m_peer_entities.find(peer_id);
    if (it == m_peer_entities.end()) return glm::vec3(0.0f);
    const auto* t = m_world->registry().try_get<tw::Transform>(it->second);
    return t ? t->position() : glm::vec3(0.0f);
}

} // namespace tw::p2p
