#include "QuicrPeerLink.hpp"
#include "quicr/QuicrConnectionListener.hpp"
#include "quicr/QuicrEndpoint.hpp"

#include <cstring>
#include <spdlog/spdlog.h>

namespace tw::p2p {
using net::quicr::QuicrEndpoint;

QuicrPeerLink::QuicrPeerLink(uint32_t self_id, uint16_t port)
    : m_self_id(self_id),
      m_endpoint(QuicrEndpoint::create_and_bind(port).value()),
      m_listener(net::quicr::QuicrConnectionListener::listen(m_endpoint.get()).value())
{}

void QuicrPeerLink::connect_to(uint32_t peer_id, const tw::net::Address& addr) {
    // The address is rebuilt from its text, which for a mapped or IPv6 peer is
    // more than the literal constructor can parse, so it goes back through the
    // resolver — into the endpoint's family, since that is what will send it.
    auto address_r = net::quicr::QuicrAddress::resolve(addr.ip_string(), addr.port(),
                                                      m_endpoint->family());
    if (!address_r) {
        spdlog::warn("QuicrPeerLink[{}]: address of peer {} failed to resolve: {}",
                     m_self_id, peer_id, address_r.error().message());
        return;
    }

    auto r = m_endpoint->connect(address_r.value());
    if (!r) {
        spdlog::warn("QuicrPeerLink[{}]: connect to peer {} failed", m_self_id, peer_id);
        return;
    }
    auto* conn = *r;
    mmo::peer::PeerHello hello;
    hello.set_peer_id(m_self_id);
    send_raw(conn, MsgType::Hello, hello, /*reliable=*/true);
    m_conns.push_back(Conn{conn, peer_id});
}

void QuicrPeerLink::send_batch(uint32_t peer_id, const mmo::peer::PeerActionBatch& batch) {
    auto it = m_by_id.find(peer_id);
    if (it == m_by_id.end()) return;
    send_raw(it->second, MsgType::ActionBatch, batch, /*reliable=*/false);
}

void QuicrPeerLink::poll() {
    m_endpoint->poll();
    drain_listener();
    poll_conns();
}

void QuicrPeerLink::drain_listener() {
    while (auto* conn = m_listener->listen()) {
        spdlog::error("PRDIDKI");
        mmo::peer::PeerHello hello;
        hello.set_peer_id(m_self_id);
        send_raw(conn, MsgType::Hello, hello, /*reliable=*/true);
        m_conns.push_back(Conn{conn, std::nullopt});
    }
}

void QuicrPeerLink::poll_conns() {
    for (auto& c : m_conns) {
        auto r = c.raw->read_into(std::span(c.recv_buf));
        if (r && *r > 0)
            dispatch(c, std::span<const std::byte>(c.recv_buf.data(), *r));
    }
}

void QuicrPeerLink::dispatch(Conn& c, std::span<const std::byte> frame) {
    if (frame.size() < sizeof(uint32_t)) return;

    uint32_t tag;
    std::memcpy(&tag, frame.data(), sizeof(tag));
    auto payload = frame.subspan(sizeof(tag));

    switch (static_cast<MsgType>(tag)) {
        case MsgType::Hello: {
            mmo::peer::PeerHello msg;
            if (!msg.ParseFromArray(payload.data(), static_cast<int>(payload.size()))) break;
            uint32_t remote_id = msg.peer_id();
            c.peer_id = remote_id;
            m_by_id[remote_id] = c.raw;
            if (m_connected_handler) m_connected_handler(remote_id);
            break;
        }
        case MsgType::ActionBatch: {
            if (!c.peer_id) break;
            mmo::peer::PeerActionBatch batch;
            if (!batch.ParseFromArray(payload.data(), static_cast<int>(payload.size()))) break;
            if (m_action_handler) {
                for (const auto& pa : batch.actions()) {
                    m_action_handler(*c.peer_id, PeerAction{
                        .peer_id   = batch.peer_id(),
                        .frame_idx = pa.frame_idx(),
                        .input     = {pa.input().x(), pa.input().y(), pa.input().z()},
                        .ack_frame = batch.ack_frame(),
                    });
                }
            }
            break;
        }
        case MsgType::Bye:
            spdlog::debug("QuicrPeerLink[{}]: bye from peer {}", m_self_id, c.peer_id.value_or(0));
            break;
        default:
            spdlog::warn("QuicrPeerLink[{}]: unknown msg type {}", m_self_id, tag);
    }
}

void QuicrPeerLink::send_raw(tw::net::quicr::QuicrConnection* conn,
                              MsgType type,
                              const google::protobuf::MessageLite& msg,
                              bool reliable)
{
    std::string payload = msg.SerializeAsString();
    std::vector<std::byte> frame(sizeof(uint32_t) + payload.size());
    uint32_t tag = static_cast<uint32_t>(type);
    std::memcpy(frame.data(), &tag, sizeof(tag));
    std::memcpy(frame.data() + sizeof(tag), payload.data(), payload.size());
    if (!conn->send_message(std::span(frame), reliable))
        spdlog::warn("QuicrPeerLink[{}]: send failed", m_self_id);
}

} // namespace tw::p2p
