#include "PeerWorldController.hpp"

#include "PeerRegistry.hpp"
#include "WorldStepProcessor.hpp"

#include <cmath>
#include <limits>
#include <spdlog/spdlog.h>
#include <thread>

namespace tw::p2p {

PeerWorldController::PeerWorldController(uint32_t self_id, uint16_t port,
                                         std::unique_ptr<PeerLink> link,
                                         const std::string& registry_path,
                                         World* world)
    : m_self_id(self_id)
    , m_port(port)
    , m_registry_path(registry_path)
    , m_link(std::move(link))
    , m_processor(world, self_id)
{
    m_link->on_connected([this](uint32_t peer_id) {
        m_peer_ids.push_back(peer_id);
        m_peer_acks[peer_id] = 0;
        m_processor.add_peer(peer_id);
        spdlog::info("PeerWorldController[{}]: peer {} connected", m_self_id, peer_id);
    });

    m_link->on_action([this](uint32_t /*peer_id*/, const PeerAction& action) {
        m_processor.submit_action(action);
        // Take max: unreliable transport can deliver packets out of order.
        auto& ack = m_peer_acks[action.peer_id];
        ack = std::max(action.ack_frame, ack);
        spdlog::info("PeerWorldController[{}]: peer {} ack={}", m_self_id, action.peer_id, ack);
    });
}

void PeerWorldController::register_self() {
    PeerRegistry reg(m_registry_path);
    reg.register_self(m_self_id, "127.0.0.1:" + std::to_string(m_port));
}

bool PeerWorldController::wait_until_peers_registered(uint32_t expected_count,
                                                       std::chrono::milliseconds timeout)
{
    PeerRegistry reg(m_registry_path);
    auto deadline = std::chrono::steady_clock::now() + timeout;
    while (std::chrono::steady_clock::now() < deadline) {
        if (reg.read_all().size() >= expected_count) return true;
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    spdlog::warn("PeerWorldController[{}]: timed out waiting for {} peers in registry",
                 m_self_id, expected_count);
    return false;
}

void PeerWorldController::discover_and_connect() {
    PeerRegistry reg(m_registry_path);
    for (const auto& info : reg.read_all()) {
        if (info.id == m_self_id || info.id < m_self_id) continue;

        auto colon = info.address.rfind(':');
        if (colon == std::string::npos) continue;

        std::string host = info.address.substr(0, colon);
        int         port = std::stoi(info.address.substr(colon + 1));
        m_link->connect_to(info.id, tw::net::Address{host, port});
    }
}

void PeerWorldController::wait_for_connections(uint32_t expected_count,
                                               std::chrono::milliseconds timeout)
{
    auto deadline = std::chrono::steady_clock::now() + timeout;
    while (m_peer_ids.size() < expected_count &&
           std::chrono::steady_clock::now() < deadline)
    {
        m_link->poll();
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }
    if (m_peer_ids.size() < expected_count)
        spdlog::warn("PeerWorldController[{}]: timed out ({}/{} connections)",
                     m_self_id, m_peer_ids.size(), expected_count);
}

bool PeerWorldController::tick() {
    m_link->poll();

    if(m_processor.committed_frame() >= std::max((int32_t)m_processor.current_frame() - 50, 0)) {
        uint32_t frame = m_processor.current_frame();
        mmo::peer::PeerAction local = build_local_action(frame);
        m_self_history[frame] = local;

        m_processor.submit_action(PeerAction{
            .peer_id   = m_self_id,
            .frame_idx = frame,
            .input     = glm::vec3(local.input().x(), local.input().y(), local.input().z()),
        });
    }

    for (uint32_t peer_id : m_peer_ids)
        send_history_to(peer_id);

    m_link->poll(); // flush outbound, pick up immediate responses

    // Prune history that all peers have acked.
    if (!m_peer_ids.empty()) {
        uint32_t min_ack = std::numeric_limits<uint32_t>::max();
        for (auto& [pid, acked] : m_peer_acks)
            min_ack = std::min(min_ack, acked);
        while (!m_self_history.empty() && m_self_history.begin()->first < min_ack)
            m_self_history.erase(m_self_history.begin());
    }

    bool has_advanced = m_processor.advance();

    return has_advanced;
}

mmo::peer::PeerAction PeerWorldController::build_local_action(uint32_t frame) const {
    float t = static_cast<float>(frame) *
              static_cast<float>(WorldStepProcessor::FIXED_DELTA_S);

    mmo::peer::PeerAction a;
    a.set_frame_idx(frame);
    a.mutable_input()->set_x(std::cos(t));
    a.mutable_input()->set_y(0.0f);
    a.mutable_input()->set_z(std::sin(t));
    return a;
}

void PeerWorldController::send_history_to(uint32_t peer_id) {
    uint32_t acked = m_peer_acks.at(peer_id);

    mmo::peer::PeerActionBatch batch;
    batch.set_peer_id(m_self_id);
    batch.set_ack_frame(m_processor.committed_frame());

    for (auto& [frame, action] : m_self_history) {
        if (frame < acked) continue;
        *batch.add_actions() = action;
    }

    if (batch.actions_size() > 0)
        m_link->send_batch(peer_id, batch);
}

} // namespace tw::p2p
