#pragma once

#include "message_protocol/MessageType.hpp"
#include "message_protocol/PeerId.hpp"

#include <functional>
#include <span>
#include <unordered_map>

namespace tw::msg {

/**
 * Routes a message body to the handler registered for its address.
 *
 * Never inspects the body, so how it is encoded is entirely the caller's
 * concern. At most one handler may be registered per address.
 */
class MessageDispatcher {
public:
    using Handler = std::function<void(PeerId, std::span<const std::byte>)>;

private:
    std::unordered_map<MessageType, Handler> m_handlers;

public:
    /** Registers `handler` for `type`, replacing any handler already there. */
    void set_handler(MessageType type, Handler handler) {
        m_handlers[type] = std::move(handler);
    }

    bool has_handler(MessageType type) const {
        return m_handlers.contains(type);
    }

    /** Invokes the handler for `type`. Returns false if there is none. */
    bool dispatch(PeerId peer, MessageType type, std::span<const std::byte> body) {
        auto handler = m_handlers.find(type);
        if(handler == m_handlers.end()) {
            return false;
        }

        handler->second(peer, body);
        return true;
    }
};

}
