#pragma once

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

#include <tl/expected.hpp>

#include <chrono>
#include <cstddef>
#include <functional>
#include <span>
#include <unordered_map>
#include <vector>

namespace tw::net::quicr {
class QuicrConnection;
}

namespace tw::msg {

class MessageDispatcher;

/**
 * Sends messages to one peer and routes the ones it sends back.
 *
 * Owned by the endpoint that created it and valid until the peer disconnects.
 * Handlers for message addresses are registered on the endpoint and shared by
 * every peer; a connection only holds the reply handlers for requests it made
 * itself.
 */
class MessageConnection {
public:
    using ReplyHandler = std::function<void(std::span<const std::byte>)>;

private:
    struct PendingRequest {
        ReplyHandler                          on_reply;
        std::function<void()>                 on_timeout;
        std::chrono::steady_clock::time_point expires_at;
    };

    net::quicr::QuicrConnection* m_connection;
    MessageDispatcher*           m_dispatcher;
    PeerId                       m_peer_id;

    std::vector<std::byte>                       m_send_buffer;
    std::unordered_map<uint32_t, PendingRequest> m_pending;
    uint32_t                                     m_next_seq = 1;
    uint64_t                                     m_bytes_sent = 0;
    uint64_t                                     m_messages_sent = 0;
    uint64_t                                     m_messages_received = 0;

    uint32_t next_seq();

    tl::expected<void, MessageError> send_impl(MessageType                type,
                                               std::span<const std::byte> body,
                                               uint32_t                   seq,
                                               bool                       reliable);

public:
    MessageConnection(PeerId peer_id,
                      net::quicr::QuicrConnection* connection,
                      MessageDispatcher* dispatcher);

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

    PeerId peer_id() const {
        return m_peer_id;
    }

    uint64_t bytes_sent() const {
        return m_bytes_sent;
    }

    /** Messages handed over for sending since the connection was created. */
    uint64_t messages_sent() const {
        return m_messages_sent;
    }

    /** Messages routed from the peer since the connection was created. */
    uint64_t messages_received() const {
        return m_messages_received;
    }

    bool is_established() const;

    tl::expected<void, MessageError> send(MessageType                type,
                                          std::span<const std::byte> body,
                                          bool                       reliable = false);

    /**
     * Sends a message the caller has already written a header into, for
     * callers that build the whole message in a buffer of their own.
     */
    tl::expected<void, MessageError> send_framed(std::span<const std::byte> message,
                                                 bool                       reliable = false);

    /**
     * Sends `body` and calls `on_reply` with the reply carrying the same
     * sequence number, or `on_timeout` if no reply arrives in time.
     */
    tl::expected<void, MessageError> request(
        MessageType                type,
        std::span<const std::byte> body,
        ReplyHandler               on_reply,
        std::chrono::milliseconds  timeout    = std::chrono::seconds(5),
        std::function<void()>      on_timeout = nullptr,
        bool                       reliable   = true);

    /**
     * Reads everything the peer has sent, routing each message, and returns
     * how many bytes were read. `scratch` is used to hold one message at a
     * time and may be reused between peers.
     */
    size_t receive(std::span<std::byte> scratch);

    /** Routes one received message to its reply handler, or to the dispatcher. */
    void on_message(std::span<const std::byte> message);

    /** Fails every request whose reply did not arrive before `now`. */
    void expire_requests(std::chrono::steady_clock::time_point now);
};

}
