#pragma once

#include "Entity.pb.h"
#include "MessageRegistry.hpp"
#include "WorldState.pb.h"
#include "network/NetworkReceiver.hpp"
#include "network/PlayerSessionRegistry.hpp"
#include "systems/Interest.hpp"
#include "systems/InterestSystem.hpp"
#include "world/Transform.hpp"

#include <tw/serial/Serial.hpp>

#include <entt/entt.hpp>
#include <spdlog/spdlog.h>
#include <tracy/Tracy.hpp>
#include <cassert>

namespace tw::net {

/**
 * Decides what to replicate and to whom.  Does not mutate world state.
 */
template<im::SpatialBackend Backend>
class StateReplicator {
    PlayerSessionRegistry*                m_client_registry;
    NetworkReceiver*                      m_network;

    // Per-client backing buffers reused every frame.
    std::vector<tw::serial::BinaryBuffer> m_frames;

    // Header(16) + spawn_hdr(4) + despawn_hdr(4) + 512 entities × 16 bytes
    static constexpr std::size_t kHeaderCapacity  = 24;
    static constexpr std::size_t kInitialCapacity = kHeaderCapacity + 512 * 16;

public:
    StateReplicator(
        PlayerSessionRegistry* client_registry,
        NetworkReceiver*       network
    ) :
        m_client_registry(client_registry),
        m_network(network)
    {}

    /**
     * Replicates the current world state for one zone to its connected clients.
     */
    void replicate(
            const entt::registry& registry,
            const im::InterestSystem<Backend>* interest_manager
    ) {
        ZoneScopedN("Replicator");

        const auto& sessions = m_client_registry->sessions();
        const std::size_t session_count = sessions.size();

        if (session_count == 0) return;

        if (m_frames.size() < session_count) {
            m_frames.resize(session_count);
            for (auto& buf : m_frames)
                buf.reserve(kInitialCapacity);
        }

        std::vector<tw::serial::WorldStateWriter> writers;
        writers.reserve(session_count);
        for (std::size_t i = 0; i < session_count; ++i)
            writers.emplace_back(m_frames[i]);

        std::vector<const im::Interest*> client_states(session_count);

        {
            ZoneScopedN("Preparing messages");

            for (std::size_t i = 0; i < session_count; ++i) {
                const auto* session = sessions[i];
                const im::Interest* state =
                    interest_manager->get_interest(session->session_id);

                client_states[i] = state;
                if (!state) continue;

                const std::size_t needed =
                    kHeaderCapacity
                    + state->spawn().size()    * 4
                    + state->despawn().size()  * 4
                    + state->interest().size() * 16;

                m_frames[i].reserve(needed);
                writers[i].reset();
                writers[i].begin(session->last_received_frame, Message<mmo::WorldStateMessage>::value);
                writers[i].write_spawns(state->spawn());
                writers[i].write_despawns(state->despawn());
                
                mmo::EntitySpawnMessage spawn_message = {};

                for(auto& spawn : state->spawn()) {
                    auto spawn_item = spawn_message.add_spawns();
                    spawn_item->set_entity_id((uint32_t)spawn);
                    spawn_item->set_is_player(false);
                    spawn_item->set_name("Test");
                }

                m_network->send_mesg(sessions[i]->session_id, spawn_message);
            }
        }

        {
            ZoneScopedN("Putting transforms into messages");

            auto view = registry.view<Transform>();
            view.each([&](entt::entity e, const Transform& t) {
                for (std::size_t i = 0; i < session_count; ++i) {
                    if (client_states[i] && client_states[i]->is_interested_in_entity(e)) {
                        writers[i].write_entity(e, t.position());
                        // spdlog::info("Entity {} is at {} {} {}", (uint32_t)e, t.position().x, t.position().y, t.position().z);
                    }
                }
            });
        }

        {
            ZoneScopedN("Sending messages");

            for (std::size_t i = 0; i < session_count; ++i) {
                if (!client_states[i]) continue;
                writers[i].end();
                m_network->send_framed(sessions[i]->session_id, writers[i].view());
            }
        }
    }
};

} // namespace tw::net
