#include "TcpPeerLink.hpp"

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

namespace tw::p2p {

TcpPeerLink::TcpPeerLink(uint32_t self_id, const tw::net::Address& listen_addr)
    : m_self_id(self_id)
    , m_listener(tw::net::TcpListener::listen(listen_addr, listen_addr.port()).value())
{
    (void)m_listener.set_non_blocking();
}

void TcpPeerLink::connect_to(uint32_t peer_id, const tw::net::Address& addr) {
    auto r = tw::net::TcpStream::connect(addr);
    if (!r) {
        spdlog::warn("TcpPeerLink[{}]: connect to peer {} failed", m_self_id, peer_id);
        return;
    }
    (void)r->set_non_blocking();

    mmo::peer::PeerHello hello;
    hello.set_peer_id(m_self_id);
    send_raw(*r, MsgType::Hello, hello);

    m_conns.push_back(Conn{std::move(*r), peer_id});
}

void TcpPeerLink::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);
}

void TcpPeerLink::poll() {
    drain_listener();
    poll_conns();
}

void TcpPeerLink::drain_listener() {
    while (true) {
        auto r = m_listener.listen();
        if (!r) break;
        (void)r->set_non_blocking();

        mmo::peer::PeerHello hello;
        hello.set_peer_id(m_self_id);
        send_raw(*r, MsgType::Hello, hello);

        m_conns.push_back(Conn{std::move(*r), std::nullopt});
    }
}

void TcpPeerLink::poll_conns() {
    for (auto& c : m_conns) {
        auto space = std::span(c.recv_buf.data() + c.recv_filled,
                               c.recv_buf.size() - c.recv_filled);
        auto r = c.stream.read_into(space);
        if (r && *r > 0) {
            c.recv_filled += *r;
            dispatch(c);
        }
    }
}

void TcpPeerLink::dispatch(Conn& c) {
    size_t pos = 0;
    while (c.recv_filled - pos >= sizeof(uint32_t)) {
        uint32_t total_len;
        std::memcpy(&total_len, c.recv_buf.data() + pos, sizeof(total_len));

        if (c.recv_filled - pos < sizeof(uint32_t) + total_len) break;

        pos += sizeof(uint32_t);
        auto msg_span = std::span(c.recv_buf.data() + pos, total_len);

        if (total_len >= sizeof(uint32_t)) {
            uint32_t tag;
            std::memcpy(&tag, msg_span.data(), sizeof(tag));
            auto payload = msg_span.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.stream;
                    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("TcpPeerLink[{}]: bye from peer {}",
                                  m_self_id, c.peer_id.value_or(0));
                    break;
                default:
                    spdlog::warn("TcpPeerLink[{}]: unknown msg type {}", m_self_id, tag);
            }
        }
        pos += total_len;
    }

    if (pos > 0) {
        std::memmove(c.recv_buf.data(), c.recv_buf.data() + pos, c.recv_filled - pos);
        c.recv_filled -= pos;
    }
}

void TcpPeerLink::send_raw(tw::net::TcpStream& stream,
                            MsgType type,
                            const google::protobuf::MessageLite& msg)
{
    std::string payload  = msg.SerializeAsString();
    uint32_t total_len   = static_cast<uint32_t>(sizeof(uint32_t) + payload.size());
    uint32_t tag         = static_cast<uint32_t>(type);

    std::vector<std::byte> frame(sizeof(uint32_t) * 2 + payload.size());
    std::memcpy(frame.data(),                      &total_len, sizeof(total_len));
    std::memcpy(frame.data() + sizeof(uint32_t),   &tag,       sizeof(tag));
    std::memcpy(frame.data() + sizeof(uint32_t)*2, payload.data(), payload.size());

    if (!stream.write(std::span(frame)))
        spdlog::warn("TcpPeerLink[{}]: send failed", m_self_id);
}

} // namespace tw::p2p
