#include "media_relay.h"

#ifdef VOICECAT_HAS_NET

#include <array>
#include <cstdio>

#include "conn_session.h"
#include "crypto/crypto.h"
#include "net/voice_frame.h"
#include "session_registry.h"

namespace voicecat::server {

MediaRelay::MediaRelay(asio::io_context& io, std::shared_ptr<SessionRegistry> registry)
    : io_(io), registry_(std::move(registry)) {}

MediaRelay::~MediaRelay() { stop(); }

bool MediaRelay::bind(uint16_t port) {
    return udp_.bind(io_, port);
}

void MediaRelay::start() {
    udp_.start_recv([this](const uint8_t* data, size_t len, asio::ip::udp::endpoint sender) {
        on_udp_frame(data, len, sender);
    });
}

void MediaRelay::stop() {
    udp_.close();
}

uint16_t MediaRelay::media_port() const {
    return static_cast<uint16_t>(udp_.local_endpoint().port());
}

void MediaRelay::on_udp_frame(const uint8_t* data, size_t len,
                                asio::ip::udp::endpoint sender) {
    if (len < 1) return;

    const uint8_t frame_type = data[0];

    if (frame_type == voicecat::net::kFrameUdpBinding) {
        // Payload = 16-byte token after the 14-byte header.
        if (len < voicecat::net::kVoiceHeaderSize + 16) return;
        std::array<uint8_t, 16> token{};
        std::memcpy(token.data(), data + voicecat::net::kVoiceHeaderSize, 16);

        auto session = registry_->find_by_udp_token(token);
        if (!session) return;

        session->set_udp_endpoint(sender);
        return;
    }

    if (frame_type == voicecat::net::kFrameVoice) {
        if (len < voicecat::net::kVoiceHeaderSize + crypto_aead_chacha20poly1305_ietf_ABYTES) return;

        // Resolve sender session.
        auto sender_session = registry_->find_by_udp_endpoint(sender);
        if (!sender_session) return;

        auto* recv_crypto = sender_session->recv_crypto();
        if (!recv_crypto) return;

        // AAD = 14-byte header (authenticated, not encrypted).
        const uint8_t* aad     = data;
        const uint8_t* sealed  = data + voicecat::net::kVoiceHeaderSize;
        size_t          sealed_len = len - voicecat::net::kVoiceHeaderSize;

        if (plain_buf_.size() < sealed_len) plain_buf_.resize(sealed_len);

        long plain_len = recv_crypto->open(sealed, sealed_len, aad, voicecat::net::kVoiceHeaderSize,
                                            plain_buf_.data(), plain_buf_.size());
        if (plain_len < 0) return;  // auth failure or replay

        // Parse the voice frame header to find the source ssrc/channel.
        voicecat::net::VoiceFrame hdr{};
        if (!voicecat::net::parse_header(data, len, hdr)) return;

        // Find the channel and get all other members.
        uint32_t uid     = sender_session->user_id();
        uint32_t channel = registry_->user_channel(uid);
        if (channel == 0) return;

        auto members = registry_->find_channel_sessions(channel, sender_session->session_id());

        // Re-encrypt and relay to each member.
        for (auto& member : members) {
            if (!member->has_udp_endpoint()) continue;

            auto* send_crypto = member->send_crypto();
            if (!send_crypto) continue;

            // Build an outgoing frame with the same 14-byte header.
            if (seal_buf_.size() < voicecat::net::kVoiceHeaderSize +
                                    static_cast<size_t>(plain_len) +
                                    crypto_aead_chacha20poly1305_ietf_ABYTES) {
                seal_buf_.resize(voicecat::net::kVoiceHeaderSize +
                                  static_cast<size_t>(plain_len) +
                                  crypto_aead_chacha20poly1305_ietf_ABYTES);
            }

            // Copy header (ssrc, ts, flags pass through for demux/playout), then rewrite the
            // seq field to THIS recipient's next send counter. The media AEAD nonce is an
            // implicit per-direction monotonic counter; open() reconstructs it from the seq in
            // the header (the AAD). Since we re-seal with the recipient's send_crypto (its own
            // counter), the verbatim sender seq would no longer match the nonce seal() uses and
            // every relayed frame would fail auth. Set seq = peek_send_counter() BEFORE sealing
            // so the header (which is the authenticated AAD) carries the matching counter.
            std::memcpy(seal_buf_.data(), data, voicecat::net::kVoiceHeaderSize);
            const uint64_t send_ctr = send_crypto->peek_send_counter();
            seal_buf_[8] = static_cast<uint8_t>((send_ctr >> 8) & 0xFF);
            seal_buf_[9] = static_cast<uint8_t>(send_ctr & 0xFF);

            uint8_t* out_payload = seal_buf_.data() + voicecat::net::kVoiceHeaderSize;
            long sealed_out = send_crypto->seal(
                plain_buf_.data(), static_cast<size_t>(plain_len),
                seal_buf_.data(), voicecat::net::kVoiceHeaderSize,
                out_payload,
                static_cast<size_t>(plain_len) + crypto_aead_chacha20poly1305_ietf_ABYTES);

            if (sealed_out < 0) continue;

            udp_.send_to(seal_buf_.data(),
                          voicecat::net::kVoiceHeaderSize + static_cast<size_t>(sealed_out),
                          member->udp_endpoint());
        }
        return;
    }

    if (frame_type == voicecat::net::kFrameKeepalive) {
        // Plaintext KEEPALIVE (docs/voice.md §6): identify the sender by its verified UDP
        // endpoint, bump last_seen so the reaper doesn't drop a client whose TCP control
        // channel is idle but whose media path is alive, and echo the keepalive back so the
        // client can measure media-path RTT/loss independently of the TCP ping.
        auto session = registry_->find_by_udp_endpoint(sender);
        if (!session) return;
        session->touch_last_seen();
        // Echo back to the sender (same plaintext header).
        udp_.send_to(data, len, sender);
        return;
    }

    // Unknown frame type: silently discard.
}

}  // namespace voicecat::server

#endif  // VOICECAT_HAS_NET