#pragma once

#include "MessageRegistry.hpp"

#include "message_protocol/MessageConnection.hpp"
#include "message_protocol/MessageEndpoint.hpp"

#include <spdlog/spdlog.h>

#include <functional>
#include <span>
#include <utility>
#include <vector>

namespace tw {

/**
 * Sends and receives protobuf messages over an endpoint.
 *
 * A view rather than an owner: several of these may share one endpoint, and
 * messages encoded some other way can travel over it at the same time.
 */
class ProtobufMessages {
    msg::MessageEndpoint* m_endpoint;

    // Reused between sends so a steady stream of messages does not allocate.
    std::vector<std::byte> m_buffer;

    template<typename T>
    bool serialize(const T& message) {
        m_buffer.resize(message.ByteSizeLong());
        return message.SerializeToArray(m_buffer.data(), static_cast<int>(m_buffer.size()));
    }

public:
    explicit ProtobufMessages(msg::MessageEndpoint* endpoint) :
        m_endpoint(endpoint) {
    }

    /** Calls `handler` for every T that arrives from any peer. */
    template<typename T>
    void set_handler(std::function<void(msg::PeerId, const T&)> handler) {
        m_endpoint->set_handler(
            Message<T>::value,
            [handler = std::move(handler), message = T{}](msg::PeerId                  peer,
                                                          std::span<const std::byte> body) mutable {
                if(!message.ParseFromArray(body.data(), static_cast<int>(body.size()))) {
                    spdlog::warn("Failed to parse a {} of {} bytes", Message<T>::value, body.size());
                    return;
                }

                handler(peer, message);
            });
    }

    template<typename T>
    tl::expected<void, msg::MessageError> send(msg::MessageConnection* peer,
                                               const T&                message,
                                               bool                    reliable = true) {
        if(!serialize(message)) {
            return tl::make_unexpected(
                msg::MessageError(msg::MessageErrorType::SendFailed, "failed to serialize the message"));
        }

        return peer->send(Message<T>::value, m_buffer, reliable);
    }

    template<typename T>
    tl::expected<void, msg::MessageError> send_to(msg::PeerId id, const T& message, bool reliable = true) {
        auto* peer = m_endpoint->peer(id);
        if(peer == nullptr) {
            return tl::make_unexpected(msg::MessageError(msg::MessageErrorType::NotConnected));
        }

        return send(peer, message, reliable);
    }

    template<typename T>
    void broadcast(const T& message, bool reliable = false) {
        if(!serialize(message)) {
            spdlog::error("Failed to serialize a {} for broadcast", Message<T>::value);
            return;
        }

        for(auto* peer : m_endpoint->peers()) {
            auto send_r = peer->send(Message<T>::value, m_buffer, reliable);
            if(!send_r) {
                spdlog::error("Failed to send to peer {}: {}", peer->peer_id(), send_r.error().message());
            }
        }
    }
};

}