#pragma once

#include "message_protocol/MessageConnection.hpp"
#include "message_protocol/MessageDispatcher.hpp"
#include "message_protocol/MessageError.hpp"
#include "message_protocol/MessageType.hpp"
#include "message_protocol/PeerId.hpp"

#include <tl/expected.hpp>

#include <cstddef>
#include <functional>
#include <memory>
#include <span>
#include <string>
#include <unordered_map>
#include <vector>

namespace tw::net::quicr {
class QuicrEndpoint;
class QuicrConnectionListener;
}

namespace tw::msg {

/**
 * Owns the connections to every peer and the handlers shared between them.
 *
 * An endpoint created with bind() accepts incoming peers; either kind may
 * connect() outwards, so one endpoint can serve peers and reach out to others
 * at the same time.
 *
 * update() must be called regularly. Nothing is received and no request ever
 * times out between calls.
 */
class MessageEndpoint {
    std::unique_ptr<net::quicr::QuicrEndpoint>           m_endpoint;
    std::unique_ptr<net::quicr::QuicrConnectionListener> m_listener;

    MessageDispatcher m_dispatcher;

    std::unordered_map<PeerId, std::unique_ptr<MessageConnection>> m_peers;
    std::function<void(PeerId)>                                    m_on_peer_connected;
    PeerId                                                         m_next_peer_id = 1;

    std::vector<std::byte> m_receive_buffer;
    uint64_t               m_bytes_received = 0;

    explicit MessageEndpoint(std::unique_ptr<net::quicr::QuicrEndpoint> endpoint);

    MessageConnection* add_peer(net::quicr::QuicrConnection* connection);
    void               accept_peers();
    void               receive();

public:
    ~MessageEndpoint();

    MessageEndpoint(const MessageEndpoint&)            = delete;
    MessageEndpoint& operator=(const MessageEndpoint&) = delete;

    /** Creates an endpoint that only connects outwards. */
    static tl::expected<std::unique_ptr<MessageEndpoint>, MessageError> create();

    /** Creates an endpoint that also accepts peers on `port`. */
    static tl::expected<std::unique_ptr<MessageEndpoint>, MessageError> bind(int port);

    tl::expected<MessageConnection*, MessageError> connect(const std::string& host, int port);

    /** Registers `handler` for every peer. */
    void set_handler(MessageType type, MessageDispatcher::Handler handler) {
        m_dispatcher.set_handler(type, std::move(handler));
    }

    MessageDispatcher& dispatcher() {
        return m_dispatcher;
    }

    /** Receives pending messages, accepts new peers and times out requests. */
    void update();

    MessageConnection* peer(PeerId id);

    /**
     * Calls `handler` for each peer that connects or is accepted, before any
     * of that peer's messages are dispatched.
     */
    void set_on_peer_connected(std::function<void(PeerId)> handler) {
        m_on_peer_connected = std::move(handler);
    }

    std::vector<MessageConnection*> peers() const;

    void broadcast(MessageType type, std::span<const std::byte> body, bool reliable = false);

    /** Bytes received since the endpoint was created. */
    uint64_t bytes_received() const {
        return m_bytes_received;
    }

    /** Bytes handed to every peer for sending since the endpoint was created. */
    uint64_t bytes_sent() const;

    /** Messages handed to every peer for sending since the endpoint was created. */
    uint64_t messages_sent() const;

    /** Messages routed from every peer since the endpoint was created. */
    uint64_t messages_received() const;
};

}
