#include "message_protocol/MessageConnection.hpp"

#include "message_protocol/MessageDispatcher.hpp"
#include "quicr/QuicrConnection.hpp"

#include <spdlog/spdlog.h>

#include <cstring>

namespace tw::msg {

namespace {
constexpr size_t INITIAL_SEND_BUFFER_SIZE = 64 * 1024;
}

MessageConnection::MessageConnection(PeerId                       peer_id,
                                     net::quicr::QuicrConnection* connection,
                                     MessageDispatcher*           dispatcher) :
    m_connection(connection),
    m_dispatcher(dispatcher),
    m_peer_id(peer_id),
    m_send_buffer(INITIAL_SEND_BUFFER_SIZE) {
}

bool MessageConnection::is_established() const {
    return m_connection->state() == net::quicr::QuicrConnectionState::Established;
}

uint32_t MessageConnection::next_seq() {
    uint32_t seq = m_next_seq++;
    if(m_next_seq == MessageHeader::SEQ_NONE) {
        m_next_seq = 1;
    }

    return seq;
}

tl::expected<void, MessageError> MessageConnection::send_impl(MessageType                type,
                                                              std::span<const std::byte> body,
                                                              uint32_t                   seq,
                                                              bool                       reliable) {
    const size_t size = MessageHeader::SIZE + body.size();

    if(m_send_buffer.size() < size) {
        m_send_buffer.resize(size);
    }

    MessageHeader{ type, seq }.encode(m_send_buffer);
    std::memcpy(m_send_buffer.data() + MessageHeader::SIZE, body.data(), body.size());

    auto send_r = m_connection->send_message(std::span(m_send_buffer).subspan(0, size), reliable);
    if(!send_r) {
        return tl::make_unexpected(MessageError(MessageErrorType::SendFailed, send_r.error().message()));
    }

    m_bytes_sent += size;
    m_messages_sent++;

    return {};
}

tl::expected<void, MessageError> MessageConnection::send(MessageType                type,
                                                         std::span<const std::byte> body,
                                                         bool                       reliable) {
    return send_impl(type, body, MessageHeader::SEQ_NONE, reliable);
}

tl::expected<void, MessageError> MessageConnection::send_framed(std::span<const std::byte> message,
                                                                bool                       reliable) {
    if(message.size() < MessageHeader::SIZE) {
        return tl::make_unexpected(
            MessageError(MessageErrorType::SendFailed, "the message is too short to hold a header"));
    }

    // send_message takes a writable span, so the bytes are staged in the send
    // buffer rather than sent straight from the caller's buffer.
    if(m_send_buffer.size() < message.size()) {
        m_send_buffer.resize(message.size());
    }

    std::memcpy(m_send_buffer.data(), message.data(), message.size());

    auto send_r = m_connection->send_message(std::span(m_send_buffer).subspan(0, message.size()), reliable);
    if(!send_r) {
        return tl::make_unexpected(MessageError(MessageErrorType::SendFailed, send_r.error().message()));
    }

    m_bytes_sent += message.size();
    m_messages_sent++;

    return {};
}

tl::expected<void, MessageError> MessageConnection::request(MessageType                type,
                                                            std::span<const std::byte> body,
                                                            ReplyHandler               on_reply,
                                                            std::chrono::milliseconds  timeout,
                                                            std::function<void()>      on_timeout,
                                                            bool                       reliable) {
    const uint32_t seq = next_seq();

    auto send_r = send_impl(type, body, seq, reliable);
    if(!send_r) {
        return send_r;
    }

    m_pending.emplace(seq,
                      PendingRequest{ std::move(on_reply),
                                      std::move(on_timeout),
                                      std::chrono::steady_clock::now() + timeout });

    return {};
}

size_t MessageConnection::receive(std::span<std::byte> scratch) {
    size_t total = 0;

    while(true) {
        auto read_r = m_connection->read_into(scratch);
        if(!read_r) {
            spdlog::error("Failed to read from peer {}: {}", m_peer_id, read_r.error().message());
            break;
        }

        if(*read_r == 0) {
            break;
        }

        total += *read_r;
        m_messages_received++;
        on_message(scratch.subspan(0, *read_r));
    }

    return total;
}

void MessageConnection::on_message(std::span<const std::byte> message) {
    auto header = MessageHeader::decode(message);
    if(!header) {
        spdlog::warn("Dropped a message of {} bytes, too short to hold a header", message.size());
        return;
    }

    auto body = message.subspan(MessageHeader::SIZE);

    if(header->seq != MessageHeader::SEQ_NONE) {
        auto pending = m_pending.find(header->seq);
        if(pending != m_pending.end()) {
            auto on_reply = std::move(pending->second.on_reply);
            m_pending.erase(pending);
            on_reply(body);
            return;
        }
    }

    if(!m_dispatcher->dispatch(m_peer_id, header->type, body)) {
        spdlog::warn("No handler for message type {}", header->type);
    }
}

void MessageConnection::expire_requests(std::chrono::steady_clock::time_point now) {
    std::erase_if(m_pending, [&](auto& entry) {
        if(entry.second.expires_at > now) {
            return false;
        }

        spdlog::warn("Request {} timed out", entry.first);
        if(entry.second.on_timeout) {
            entry.second.on_timeout();
        }

        return true;
    });
}

}
