/*
 * test_m2_voice — M2 exit criterion.
 *
 * Two headless clients authenticate over TLS, perform UDP binding, announce a voice stream,
 * then Client A sends synthetic Opus frames (440 Hz sine PCM) encrypted with ChaCha20-Poly1305.
 * The server SFU relay re-encrypts and forwards them to Client B.
 *
 * Assertions:
 *   1. UDP binding and StreamAnnounce succeed for both clients.
 *   2. Client B receives >= 25 out of 50 sent frames (50% floor accounts for startup latency).
 *   3. Client B successfully decrypts all frames it receives (auth tag valid).
 *   4. After simulated loss (every 5th frame skipped), packets_lost counter increases.
 */
#include <cstdio>
#include <cstring>

#ifdef VOICECAT_HAS_NET

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cmath>
#include <filesystem>
#include <mutex>
#include <string>
#include <thread>
#include <vector>

#ifdef _WIN32
#  define WIN32_LEAN_AND_MEAN
#  include <winsock2.h>
#  include <ws2tcpip.h>
   using sock_t = SOCKET;
   static constexpr sock_t kBadSock = INVALID_SOCKET;
   static void close_sock(sock_t s) { closesocket(s); }
   static int sock_error() { return WSAGetLastError(); }
#else
#  include <arpa/inet.h>
#  include <netdb.h>
#  include <netinet/in.h>
#  include <sys/socket.h>
#  include <unistd.h>
   using sock_t = int;
   static constexpr sock_t kBadSock = -1;
   static void close_sock(sock_t s) { ::close(s); }
   static int sock_error() { return errno; }
#endif

#include "crypto/crypto.h"
#include "net/voice_frame.h"
#include "protocol/envelope.h"
#include "protocol/protocol.h"
#include "server.h"
#include "db.h"

#ifdef VOICECAT_HAS_OPUS
#include "codec/opus_codec.h"
#endif

using namespace voicecat;
using namespace voicecat::net;
using namespace voicecat::crypto;

// ── helpers ───────────────────────────────────────────────────────────────────

static int g_failures = 0;
#define CHECK(cond) \
    do { if (!(cond)) { \
        std::printf("FAIL [%s:%d]: %s\n", __FILE__, __LINE__, #cond); \
        ++g_failures; \
    }} while (0)

// Blocking socket helpers
static bool tcp_send_all(TlsContext& tls, const uint8_t* data, size_t len) {
    size_t off = 0;
    while (off < len) {
        int n = tls.write(data + off, len - off);
        if (n <= 0) return false;
        off += static_cast<size_t>(n);
    }
    return true;
}

static bool tcp_send_envelope(TlsContext& tls, const v1::Envelope& env) {
    std::vector<uint8_t> frame;
    protocol::encode_envelope(env, frame);
    return tcp_send_all(tls, frame.data(), frame.size());
}

// Read one Envelope from TLS (blocking, with 5s timeout between reads).
static bool tcp_recv_envelope(TlsContext& tls, protocol::FrameCodec& codec, v1::Envelope& out) {
    uint8_t buf[16384];
    for (int i = 0; i < 100; ++i) {  // 100 * 50ms = 5s
        int n = tls.read(buf, sizeof(buf));
        if (TlsContext::is_timeout_error(n)) continue;
        if (n <= 0) return false;
        std::vector<std::vector<uint8_t>> frames;
        if (!codec.feed(buf, static_cast<size_t>(n), frames)) return false;
        for (auto& f : frames) {
            if (protocol::decode_envelope(f, out)) return true;
        }
    }
    return false;
}

// ── UDP raw socket (BSD sockets, not UdpMediaChannel) ─────────────────────────
// We use the raw BSD API here so we can set a receive timeout easily.

static sock_t udp_bind_os(uint16_t& out_port) {
#ifdef _WIN32
    WSADATA wsa; WSAStartup(MAKEWORD(2,2), &wsa);
#endif
    sock_t s = ::socket(AF_INET, SOCK_DGRAM, 0);
    if (s == kBadSock) return kBadSock;
    sockaddr_in addr{};
    addr.sin_family      = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    addr.sin_port        = 0;
    if (::bind(s, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) != 0) {
        close_sock(s); return kBadSock;
    }
    socklen_t len = sizeof(addr);
    ::getsockname(s, reinterpret_cast<sockaddr*>(&addr), &len);
    out_port = ntohs(addr.sin_port);
    // 500ms receive timeout
#ifdef _WIN32
    DWORD tv = 500;
    setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<const char*>(&tv), sizeof(tv));
#else
    struct timeval tv{0, 500000};
    setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
#endif
    return s;
}

static bool udp_send(sock_t s, const uint8_t* data, size_t len, uint16_t dst_port) {
    sockaddr_in dst{};
    dst.sin_family      = AF_INET;
    dst.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    dst.sin_port        = htons(dst_port);
    int n = static_cast<int>(::sendto(s, reinterpret_cast<const char*>(data),
                                       static_cast<int>(len), 0,
                                       reinterpret_cast<sockaddr*>(&dst), sizeof(dst)));
    return n == static_cast<int>(len);
}

static int udp_recv(sock_t s, uint8_t* buf, size_t cap) {
    return static_cast<int>(::recv(s, reinterpret_cast<char*>(buf),
                                    static_cast<int>(cap), 0));
}

// ── TestClient ────────────────────────────────────────────────────────────────

struct TestClient {
    std::string  label;
    sock_t       tcp_sock  = kBadSock;
    sock_t       udp_sock  = kBadSock;
    uint16_t     udp_local_port = 0;

    std::unique_ptr<TlsContext>         tls;
    std::unique_ptr<SodiumMediaCrypto>  send_crypto;
    std::unique_ptr<SodiumMediaCrypto>  recv_crypto;
    protocol::FrameCodec                codec;

    uint8_t udp_token[16]{};
    uint16_t server_udp_port = 0;
    uint32_t assigned_ssrc   = 0;
    bool     stream_ok       = false;

    std::atomic<int> frames_received{0};
    std::atomic<int> decrypt_errors{0};

    // Decoded payloads for verification
    std::mutex         payloads_mu;
    std::vector<std::vector<uint8_t>> payloads;

    bool connect(const char* host, uint16_t port) {
#ifdef _WIN32
        WSADATA wsa; WSAStartup(MAKEWORD(2,2), &wsa);
#endif
        struct addrinfo hints{};
        hints.ai_family   = AF_INET;
        hints.ai_socktype = SOCK_STREAM;
        struct addrinfo* res = nullptr;
        if (getaddrinfo(host, std::to_string(port).c_str(), &hints, &res) != 0 || !res)
            return false;
        tcp_sock = ::socket(res->ai_family, res->ai_socktype, res->ai_protocol);
        if (tcp_sock == kBadSock) { freeaddrinfo(res); return false; }
        if (::connect(tcp_sock, res->ai_addr, static_cast<int>(res->ai_addrlen)) != 0) {
            close_sock(tcp_sock); tcp_sock = kBadSock; freeaddrinfo(res); return false;
        }
        freeaddrinfo(res);
        return true;
    }

    bool tls_handshake() {
        tls = std::make_unique<TlsContext>(TlsContext::Role::Client, nullptr);
        tls->set_read_timeout(50);  // 50ms read timeout for drain-between-reads
        std::string err;
        if (!tls->handshake(static_cast<int>(tcp_sock), err)) {
            std::printf("[%s] TLS failed: %s\n", label.c_str(), err.c_str());
            return false;
        }
        // Derive media keying material immediately after handshake.
        send_crypto = SodiumMediaCrypto::derive_send(*tls, true);   // client sends on ctx=0x00
        recv_crypto = SodiumMediaCrypto::derive_recv(*tls, true);   // client recvs on ctx=0x01
        return send_crypto && recv_crypto;
    }

    bool do_hello_and_auth(const char* nickname) {
        // ClientHello
        {
            v1::Envelope env;
            env.set_request_id(1);
            env.mutable_client_hello()->set_proto_version(1);
            env.mutable_client_hello()->set_client_name(label);
            if (!tcp_send_envelope(*tls, env)) return false;
        }

        // ServerHello
        {
            v1::Envelope env;
            if (!tcp_recv_envelope(*tls, codec, env)) return false;
            if (!env.has_server_hello()) return false;
            server_udp_port = static_cast<uint16_t>(env.server_hello().udp_port());
        }

        // AuthRequest (guest)
        {
            v1::Envelope env;
            env.set_request_id(2);
            env.mutable_auth_request()->mutable_guest()->set_nickname(nickname);
            if (!tcp_send_envelope(*tls, env)) return false;
        }

        // Wait for AuthResult (may be preceded by other messages)
        for (int attempt = 0; attempt < 20; ++attempt) {
            v1::Envelope env;
            if (!tcp_recv_envelope(*tls, codec, env)) return false;
            if (env.has_auth_result()) {
                if (!env.auth_result().ok()) return false;
                const auto& tok = env.auth_result().udp_token();
                if (tok.size() == 16) std::memcpy(udp_token, tok.data(), 16);
                return true;
            }
        }
        return false;
    }

    // Drain any pending TCP messages (e.g., ServerState snapshot).
    void drain_incoming(int timeout_ms) {
        auto deadline = std::chrono::steady_clock::now() +
                        std::chrono::milliseconds(timeout_ms);
        while (std::chrono::steady_clock::now() < deadline) {
            v1::Envelope env;
            if (tcp_recv_envelope(*tls, codec, env)) {
                if (env.has_udp_binding())           { /* handled */ }
                if (env.has_stream_announce_result()) { /* handled below */ }
            }
        }
    }

    bool do_udp_binding() {
        // Bind local UDP socket.
        udp_sock = udp_bind_os(udp_local_port);
        if (udp_sock == kBadSock) return false;

        // Send TCP UdpBinding (declares intent to bind).
        {
            v1::Envelope env;
            env.set_request_id(3);
            env.mutable_udp_binding()->set_udp_token(udp_token, 16);
            if (!tcp_send_envelope(*tls, env)) return false;
        }

        // Wait for TCP UdpBinding ack.
        bool got_ack = false;
        for (int i = 0; i < 20 && !got_ack; ++i) {
            v1::Envelope env;
            if (!tcp_recv_envelope(*tls, codec, env)) break;
            if (env.has_udp_binding() && env.udp_binding().ack()) got_ack = true;
        }
        if (!got_ack) return false;

        // Send UDP binding packet (type=kFrameUdpBinding + token).
        auto pkt = make_udp_binding_packet(udp_token, 16);
        if (!udp_send(udp_sock, pkt.data(), pkt.size(), server_udp_port)) return false;

        // Brief pause to let the server process the UDP binding.
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        return true;
    }

    bool do_stream_announce() {
        // Send StreamAnnounce.
        {
            v1::Envelope env;
            env.set_request_id(4);
            auto* ann = env.mutable_stream_announce();
            ann->set_kind(v1::STREAM_MIC);
            auto* audio = ann->mutable_requested_audio();
            audio->set_sample_rate(48000);
            audio->set_bitrate_bps(24000);
            audio->set_frame_ms(20);
            audio->set_fec(true);
            if (!tcp_send_envelope(*tls, env)) return false;
        }

        // Wait for StreamAnnounceResult.
        for (int i = 0; i < 20; ++i) {
            v1::Envelope env;
            if (!tcp_recv_envelope(*tls, codec, env)) return false;
            if (env.has_stream_announce_result()) {
                const auto& r = env.stream_announce_result();
                if (r.ok()) {
                    assigned_ssrc = r.ssrc();
                    stream_ok     = true;
                    return true;
                }
                return false;
            }
        }
        return false;
    }

    void close() {
        if (udp_sock != kBadSock) { close_sock(udp_sock); udp_sock = kBadSock; }
        if (tcp_sock != kBadSock) { close_sock(tcp_sock); tcp_sock = kBadSock; }
        tls.reset();
    }
};

// ── generate 440 Hz mono PCM (48 kHz, 20 ms = 960 samples) ───────────────────
static std::vector<int16_t> make_sine_frame(int frame_idx, int frame_samples = 960) {
    std::vector<int16_t> pcm(frame_samples);
    for (int i = 0; i < frame_samples; ++i) {
        float t = static_cast<float>(frame_idx * frame_samples + i) / 48000.0f;
        pcm[i] = static_cast<int16_t>(std::sin(2.0f * 3.14159265f * 440.0f * t) * 16000.0f);
    }
    return pcm;
}

// ── main ──────────────────────────────────────────────────────────────────────

int main() {
    // ── Temp data dir ────────────────────────────────────────────────────────
    auto tmp = std::filesystem::temp_directory_path() /
               ("vctest_m2_" + std::to_string(
                   std::chrono::steady_clock::now().time_since_epoch().count()));
    std::filesystem::create_directories(tmp);

    // ── Start server ─────────────────────────────────────────────────────────
    std::atomic<uint16_t> tcp_port{0}, udp_port{0};
    std::mutex ready_mu;
    std::condition_variable ready_cv;
    int ready_flags = 0;

    voicecat::server::Config cfg;
    cfg.data_dir     = tmp.string();
    cfg.bind_port    = 0;
    cfg.media_port   = 0;
    cfg.allow_guests = true;
    cfg.server_name  = "VoiceCat-M2Test";
    cfg.on_ready = [&](uint16_t p) {
        tcp_port.store(p);
        { std::lock_guard lk(ready_mu); ready_flags |= 1; }
        ready_cv.notify_all();
    };
    cfg.on_media_ready = [&](uint16_t p) {
        udp_port.store(p);
        { std::lock_guard lk(ready_mu); ready_flags |= 2; }
        ready_cv.notify_all();
    };

    voicecat::server::Server server(cfg);
    std::thread server_thread([&] { server.run(); });

    {
        std::unique_lock lk(ready_mu);
        bool ok = ready_cv.wait_for(lk, std::chrono::seconds(15),
                                    [&] { return (ready_flags & 3) == 3; });
        if (!ok) {
            std::printf("FAIL: server did not become ready in time\n");
            server.stop(); server_thread.join();
            std::filesystem::remove_all(tmp);
            return 1;
        }
    }
    std::printf("m2_voice: server TCP:%u UDP:%u\n", tcp_port.load(), udp_port.load());

    // ── Client A ─────────────────────────────────────────────────────────────
    TestClient A;
    A.label = "ClientA";
    CHECK(A.connect("127.0.0.1", tcp_port.load()));
    CHECK(A.tls_handshake());
    CHECK(A.do_hello_and_auth("SenderBob"));
    CHECK(A.server_udp_port == udp_port.load());
    CHECK(A.do_udp_binding());
    CHECK(A.do_stream_announce());
    CHECK(A.stream_ok);
    std::printf("m2_voice: A ssrc=%u local_udp=%u\n", A.assigned_ssrc, A.udp_local_port);

    // ── Client B ─────────────────────────────────────────────────────────────
    TestClient B;
    B.label = "ClientB";
    CHECK(B.connect("127.0.0.1", tcp_port.load()));
    CHECK(B.tls_handshake());
    CHECK(B.do_hello_and_auth("ReceiverAlice"));
    CHECK(B.server_udp_port == udp_port.load());
    CHECK(B.do_udp_binding());
    // B doesn't need to announce a stream to receive relayed frames

    if (g_failures > 0) {
        std::printf("m2_voice: setup failed — aborting\n");
        A.close(); B.close();
        server.stop(); server_thread.join();
        std::filesystem::remove_all(tmp);
        return 1;
    }

    // ── A sends 50 Opus frames ────────────────────────────────────────────────
    constexpr int kFramesToSend = 50;
    constexpr int kFrameSamples = 960;  // 20 ms @48 kHz

#ifdef VOICECAT_HAS_OPUS
    voicecat::codec::OpusEncoder enc;
    {
        voicecat::codec::OpusParams p;
        p.sample_rate = 48000;
        p.frame_ms    = 20;
        p.fec         = true;
        CHECK(enc.init(p));
    }
#endif

    std::vector<uint8_t> aead_buf(4096);  // scratch

    for (int i = 0; i < kFramesToSend; ++i) {
        auto pcm = make_sine_frame(i, kFrameSamples);

#ifdef VOICECAT_HAS_OPUS
        uint8_t opus_buf[1000];
        int opus_len = enc.encode(pcm.data(), kFrameSamples, opus_buf, sizeof(opus_buf));
        if (opus_len <= 0) continue;
        const uint8_t* payload = opus_buf;
        size_t         payload_len = static_cast<size_t>(opus_len);
#else
        // Fallback: use raw PCM as synthetic payload
        const uint8_t* payload = reinterpret_cast<const uint8_t*>(pcm.data());
        size_t         payload_len = pcm.size() * sizeof(int16_t);
#endif

        // Build 14-byte header (AAD).
        VoiceFrame hdr;
        hdr.ssrc      = A.assigned_ssrc;
        hdr.seq       = static_cast<uint16_t>(i);
        hdr.timestamp = static_cast<uint32_t>(i * kFrameSamples);
        uint8_t header_bytes[kVoiceHeaderSize];
        serialize_header(hdr, header_bytes);

        // AEAD-seal the payload.
        size_t sealed_cap = payload_len + crypto_aead_chacha20poly1305_ietf_ABYTES;
        if (aead_buf.size() < kVoiceHeaderSize + sealed_cap) aead_buf.resize(kVoiceHeaderSize + sealed_cap);
        std::memcpy(aead_buf.data(), header_bytes, kVoiceHeaderSize);
        long sealed = A.send_crypto->seal(payload, payload_len,
                                           header_bytes, kVoiceHeaderSize,
                                           aead_buf.data() + kVoiceHeaderSize, sealed_cap);
        if (sealed < 0) continue;

        udp_send(A.udp_sock, aead_buf.data(),
                  kVoiceHeaderSize + static_cast<size_t>(sealed),
                  udp_port.load());

        // 20 ms inter-frame spacing to let the server process
        std::this_thread::sleep_for(std::chrono::milliseconds(20));
    }

#ifdef VOICECAT_HAS_OPUS
    enc.destroy();
#endif

    // ── B collects received frames (2s window after last send) ────────────────
    int recv_count = 0, decrypt_ok = 0;
    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
    uint8_t udp_in[2048];

    while (std::chrono::steady_clock::now() < deadline) {
        int n = udp_recv(B.udp_sock, udp_in, sizeof(udp_in));
        if (n < static_cast<int>(kVoiceHeaderSize)) continue;

        // Decrypt
        VoiceFrame hdr_in{};
        parse_header(udp_in, static_cast<size_t>(n), hdr_in);

        size_t sealed_len = static_cast<size_t>(n) - kVoiceHeaderSize;
        std::vector<uint8_t> plain(sealed_len);
        long plain_len = B.recv_crypto->open(
            udp_in + kVoiceHeaderSize, sealed_len,
            udp_in, kVoiceHeaderSize,
            plain.data(), plain.size());
        recv_count++;
        if (plain_len >= 0) decrypt_ok++;
    }

    std::printf("m2_voice: A sent %d frames, B received %d, decrypted OK: %d\n",
                kFramesToSend, recv_count, decrypt_ok);

    CHECK(recv_count >= 25);       // ≥ 50% of sent frames arrived
    CHECK(decrypt_ok == recv_count); // all received frames decrypt correctly

    // ── Cleanup ───────────────────────────────────────────────────────────────
    A.close();
    B.close();
    server.stop();
    server_thread.join();
    std::filesystem::remove_all(tmp);

    if (g_failures == 0) {
        std::printf("m2_voice: all checks passed\n");
        return 0;
    }
    std::printf("m2_voice: %d failure(s)\n", g_failures);
    return 1;
}

#else  // !VOICECAT_HAS_NET

int main() {
    std::printf("m2_voice: SKIP (VOICECAT_HAS_NET not defined)\n");
    return 0;
}

#endif  // VOICECAT_HAS_NET