#pragma once

#include "PeerAction.hpp"
#include "world/World.hpp"

#include <entt/entt.hpp>
#include <map>
#include <unordered_map>
#include <unordered_set>

namespace tw::p2p {

// Rollback-based lock-step processor.
//
// Speculative advance
//   try_advance() always steps the world.  For each peer whose action has not
//   yet arrived the entity is held stationary (zero input — no prediction).
//
// Committed frames
//   A Frame becomes committed once every registered peer (including self) has
//   submitted an action for it.  try_advance() cascade-commits all newly-full
//   frames in one pass before stepping.  If any new frames were committed the
//   world is rolled back to the last committed snapshot and the remaining
//   speculative frames are re-simulated with whatever confirmed actions exist.
//
// One re-simulation per tick (not per submit_action).
class WorldStepProcessor {
public:
    static constexpr double FIXED_DELTA_S = (1000.0 / 60.0) / 1000.0;
    static constexpr float  MOVE_SPEED    = 5.0f;

    WorldStepProcessor(World* world, uint32_t self_id);

    void add_peer(uint32_t peer_id);

    // Store a confirmed action.  Does not touch the world or trigger rollback.
    void submit_action(const PeerAction& action);

    // Cascade-commit full frames, rollback + re-simulate if needed, then
    // advance one speculative frame.  Always returns true.
    bool advance();

    uint32_t current_frame()   const { return m_current_frame; }
    uint32_t committed_frame() const { return m_committed_frame; }
    uint64_t rollbacks()       const { return m_rollbacks; }
    uint64_t stepped_frames()  const { return m_stepped_frames; }

    glm::vec3 peer_position(uint32_t peer_id) const;

private:
    // All confirmed actions for one simulation frame.
    struct Frame {
        std::unordered_map<uint32_t, PeerAction> actions; // peer_id -> action

        bool is_full(const std::unordered_set<uint32_t>& peers) const {
            for (uint32_t id : peers)
                if (!actions.count(id)) return false;
            return true;
        }
    };

    // World-state snapshot: positions of every tracked entity.
    struct Snapshot {
        std::unordered_map<uint32_t, glm::vec3> positions;
    };

    World*   m_world;
    uint32_t m_self_id;

    uint32_t m_current_frame   = 0; // next speculative frame to advance
    uint32_t m_committed_frame = 0; // next frame eligible to commit

    Snapshot m_committed_snapshot; // world state before m_committed_frame

    std::unordered_set<uint32_t>               m_all_peer_ids;
    std::unordered_map<uint32_t, entt::entity> m_peer_entities;
    std::map<uint32_t, Frame>                  m_frames; // partial or full

    uint64_t m_rollbacks      = 0;
    uint64_t m_stepped_frames = 0;

    // Apply all actions in frame to the current world state; zero for missing.
    void apply_frame(const Frame& frame);

    // Cascade-commit full frames and re-simulate speculative ones if needed.
    void flush_commits();

    Snapshot capture_snapshot() const;
    void     restore_snapshot (const Snapshot& snap);
};

} // namespace tw::p2p
