/*
 * test_media_aead — ChaCha20-Poly1305 AEAD seal/open, anti-replay, tamper detection.
 *
 * Uses a synthetic 32-byte key directly (no TLS context needed for unit tests).
 */
#include <array>
#include <cstdio>
#include <cstring>
#include <vector>

#include <sodium.h>

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

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

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)

// Build a synthetic voice-frame-header AAD.
static std::vector<uint8_t> make_aad(uint64_t seq) {
    VoiceFrame f;
    f.ssrc = 0xCAFEBABE;
    f.seq  = seq;
    std::vector<uint8_t> aad(kVoiceHeaderSize);
    serialize_header(f, aad.data());
    return aad;
}

static void test_seal_open_round_trip() {
    uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES];
    crypto_generichash(key, sizeof(key),
                       reinterpret_cast<const uint8_t*>("test-key"), 8, nullptr, 0);

    SodiumMediaCrypto sender(key);
    SodiumMediaCrypto receiver(key);

    // Copy the receiver state so it starts with the same key but its own counter.

    std::vector<uint8_t> plain(100, 0xAB);
    auto aad = make_aad(0);

    // Seal
    std::vector<uint8_t> cipher(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES);
    long sealed_len = sender.seal(plain.data(), plain.size(),
                                   aad.data(), aad.size(),
                                   cipher.data(), cipher.size());
    CHECK(sealed_len == static_cast<long>(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES));

    // Open
    std::vector<uint8_t> recovered(plain.size());
    long plain_len = receiver.open(cipher.data(), static_cast<size_t>(sealed_len),
                                    aad.data(), aad.size(),
                                    recovered.data(), recovered.size());
    CHECK(plain_len == static_cast<long>(plain.size()));
    CHECK(std::memcmp(plain.data(), recovered.data(), plain.size()) == 0);
}

static void test_anti_replay() {
    uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES];
    crypto_generichash(key, sizeof(key),
                       reinterpret_cast<const uint8_t*>("replay-key"), 10, nullptr, 0);

    SodiumMediaCrypto sender(key);
    SodiumMediaCrypto receiver(key);

    std::vector<uint8_t> plain(50, 0x55);
    auto aad = make_aad(0);

    std::vector<uint8_t> cipher(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES);
    long sealed_len = sender.seal(plain.data(), plain.size(),
                                   aad.data(), aad.size(),
                                   cipher.data(), cipher.size());
    CHECK(sealed_len > 0);

    std::vector<uint8_t> recovered(plain.size());

    // First open succeeds.
    long r1 = receiver.open(cipher.data(), static_cast<size_t>(sealed_len),
                             aad.data(), aad.size(),
                             recovered.data(), recovered.size());
    CHECK(r1 == static_cast<long>(plain.size()));

    // Replay of the same ciphertext must fail.
    long r2 = receiver.open(cipher.data(), static_cast<size_t>(sealed_len),
                             aad.data(), aad.size(),
                             recovered.data(), recovered.size());
    CHECK(r2 < 0);
}

static void test_tamper_detection() {
    uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES];
    crypto_generichash(key, sizeof(key),
                       reinterpret_cast<const uint8_t*>("tamper-key"), 10, nullptr, 0);

    SodiumMediaCrypto sender(key);
    SodiumMediaCrypto receiver(key);

    std::vector<uint8_t> plain(40, 0x77);
    auto aad = make_aad(0);

    std::vector<uint8_t> cipher(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES);
    long sealed_len = sender.seal(plain.data(), plain.size(),
                                   aad.data(), aad.size(),
                                   cipher.data(), cipher.size());
    CHECK(sealed_len > 0);

    // Flip a byte in the ciphertext.
    cipher[5] ^= 0xFF;

    std::vector<uint8_t> recovered(plain.size());
    long r = receiver.open(cipher.data(), static_cast<size_t>(sealed_len),
                            aad.data(), aad.size(),
                            recovered.data(), recovered.size());
    CHECK(r < 0);
}

static void test_multiple_packets() {
    uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES];
    crypto_generichash(key, sizeof(key),
                       reinterpret_cast<const uint8_t*>("multi-key"), 9, nullptr, 0);

    SodiumMediaCrypto sender(key);
    SodiumMediaCrypto receiver(key);

    std::vector<uint8_t> plain(60, 0x99);

    for (uint64_t seq = 0; seq < 10; ++seq) {
        auto aad = make_aad(seq);
        std::vector<uint8_t> cipher(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES);
        long sealed_len = sender.seal(plain.data(), plain.size(),
                                       aad.data(), aad.size(),
                                       cipher.data(), cipher.size());
        CHECK(sealed_len > 0);

        std::vector<uint8_t> recovered(plain.size());
        long plain_len = receiver.open(cipher.data(), static_cast<size_t>(sealed_len),
                                        aad.data(), aad.size(),
                                        recovered.data(), recovered.size());
        CHECK(plain_len == static_cast<long>(plain.size()));
        CHECK(std::memcmp(plain.data(), recovered.data(), plain.size()) == 0);
    }
}

// Build a voice-frame-header AAD with a given ssrc + seq.
static std::vector<uint8_t> make_aad_ssrc(uint32_t ssrc, uint64_t seq) {
    VoiceFrame f;
    f.ssrc = ssrc;
    f.seq  = seq;
    std::vector<uint8_t> aad(kVoiceHeaderSize);
    serialize_header(f, aad.data());
    return aad;
}

// Overwrite the 8-byte big-endian seq field (header bytes [8..15]) in an AAD buffer.
static void set_aad_seq(std::vector<uint8_t>& aad, uint64_t seq) {
    for (int i = 0; i < 8; ++i)
        aad[8 + i] = static_cast<uint8_t>((seq >> (56 - 8 * i)) & 0xFF);
}

// Simulate one server relay hop for a single frame, sender → recipient R.
//   - sender seals with its send key, setting header seq = its own send counter (client contract).
//   - server opens with the sender's key, then re-seals with R's send key.
//   - if rewrite_seq, the re-sealed header's seq is set to R's send counter (the fix); otherwise
//     the sender's seq is forwarded verbatim (the bug).
// Returns true iff R successfully decrypts the relayed frame.
static bool relay_one(SodiumMediaCrypto& sender_send, SodiumMediaCrypto& server_recv,
                      SodiumMediaCrypto& r_send, SodiumMediaCrypto& r_recv,
                      uint32_t ssrc, const std::vector<uint8_t>& plain, bool rewrite_seq) {
    // Client A→server: seq carries the sender's send counter.
    auto in_aad = make_aad_ssrc(ssrc, sender_send.peek_send_counter());
    std::vector<uint8_t> cipher(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES);
    long sealed = sender_send.seal(plain.data(), plain.size(), in_aad.data(), in_aad.size(),
                                   cipher.data(), cipher.size());
    if (sealed < 0) return false;

    // Server decrypts the inbound frame.
    std::vector<uint8_t> recovered(plain.size());
    long opened = server_recv.open(cipher.data(), static_cast<size_t>(sealed),
                                   in_aad.data(), in_aad.size(),
                                   recovered.data(), recovered.size());
    if (opened < 0) return false;

    // Server re-seals to R. Header passes through except seq, which (when fixed) is set to R's
    // own send counter so R's open() reconstructs the matching nonce.
    std::vector<uint8_t> out_aad = in_aad;  // copy header verbatim
    if (rewrite_seq) set_aad_seq(out_aad, r_send.peek_send_counter());
    std::vector<uint8_t> relay_cipher(recovered.size() + crypto_aead_chacha20poly1305_ietf_ABYTES);
    long resealed = r_send.seal(recovered.data(), static_cast<size_t>(opened),
                                out_aad.data(), out_aad.size(),
                                relay_cipher.data(), relay_cipher.size());
    if (resealed < 0) return false;

    // R decrypts the relayed frame.
    std::vector<uint8_t> r_recovered(recovered.size());
    long r_opened = r_recv.open(relay_cipher.data(), static_cast<size_t>(resealed),
                                out_aad.data(), out_aad.size(),
                                r_recovered.data(), r_recovered.size());
    return r_opened == static_cast<long>(plain.size());
}

// Regression test for the relay nonce-desync bug: two senders (A, B) relayed into one recipient
// (R) interleaved. The media AEAD nonce is an implicit per-direction counter reconstructed from
// the header seq; if the relay forwards the sender's seq verbatim, it no longer matches R's send
// counter and frames fail to decrypt. The relay must rewrite seq = R's send counter.
static void test_relay_interleaved_reseal() {
    auto make_key = [](const char* label) {
        std::array<uint8_t, crypto_aead_chacha20poly1305_ietf_KEYBYTES> k{};
        crypto_generichash(k.data(), k.size(),
                           reinterpret_cast<const uint8_t*>(label),
                           std::strlen(label), nullptr, 0);
        return k;
    };
    auto kA = make_key("relay-A");  // A↔server direction
    auto kB = make_key("relay-B");  // B↔server direction
    auto kR = make_key("relay-R");  // server↔R direction

    std::vector<uint8_t> plain(80, 0x3C);

    // Fixed path: interleaved A/B frames all decrypt at R.
    {
        SodiumMediaCrypto a_send(kA.data()), srv_recv_a(kA.data());
        SodiumMediaCrypto b_send(kB.data()), srv_recv_b(kB.data());
        SodiumMediaCrypto r_send(kR.data()), r_recv(kR.data());

        bool all_ok = true;
        for (int i = 0; i < 8; ++i) {
            all_ok &= relay_one(a_send, srv_recv_a, r_send, r_recv, 0x1111, plain, /*rewrite=*/true);
            all_ok &= relay_one(b_send, srv_recv_b, r_send, r_recv, 0x2222, plain, /*rewrite=*/true);
        }
        CHECK(all_ok);  // with the fix, every interleaved relayed frame decrypts at R
    }

    // Control: forwarding seq verbatim (the bug) must drop frames once the counters diverge.
    {
        SodiumMediaCrypto a_send(kA.data()), srv_recv_a(kA.data());
        SodiumMediaCrypto b_send(kB.data()), srv_recv_b(kB.data());
        SodiumMediaCrypto r_send(kR.data()), r_recv(kR.data());

        int failures = 0;
        for (int i = 0; i < 8; ++i) {
            if (!relay_one(a_send, srv_recv_a, r_send, r_recv, 0x1111, plain, /*rewrite=*/false)) ++failures;
            if (!relay_one(b_send, srv_recv_b, r_send, r_recv, 0x2222, plain, /*rewrite=*/false)) ++failures;
        }
        CHECK(failures > 0);  // proves the verbatim-seq path is broken (locks in the regression)
    }
}

// Regression for the bad-wifi wedge: the anti-replay window must NOT be advanced by a
// packet that fails authentication. A single corrupted/forged frame carrying a huge seq
// used to shove recv_highest_ far ahead (before the AEAD tag was checked), after which
// every legitimate frame was rejected as "too old" — permanent silence. open() now
// advances the window only after a successful tag check (RFC 3711 §3.3).
static void test_corrupted_seq_does_not_poison_window() {
    uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES];
    crypto_generichash(key, sizeof(key),
                       reinterpret_cast<const uint8_t*>("poison-key"), 10, nullptr, 0);
    SodiumMediaCrypto sender(key);
    SodiumMediaCrypto receiver(key);

    std::vector<uint8_t> plain(64, 0x5A);
    std::vector<uint8_t> recovered(plain.size());

    auto seal_at_current = [&](std::vector<uint8_t>& aad_out, std::vector<uint8_t>& cipher_out) {
        aad_out = make_aad_ssrc(0xABCD, sender.peek_send_counter());
        cipher_out.assign(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES, 0);
        long s = sender.seal(plain.data(), plain.size(), aad_out.data(), aad_out.size(),
                             cipher_out.data(), cipher_out.size());
        CHECK(s > 0);
    };

    // 1. A normal frame (counter 0) decrypts. recv_highest_ = 0.
    std::vector<uint8_t> aad0, cipher0;
    seal_at_current(aad0, cipher0);  // sender counter 0 → 1
    CHECK(receiver.open(cipher0.data(), cipher0.size(), aad0.data(), aad0.size(),
                        recovered.data(), recovered.size()) == static_cast<long>(plain.size()));

    // 2. A frame whose header seq has been corrupted to a huge value: it fails auth
    //    (the AAD no longer matches what was sealed) and must NOT move the window.
    std::vector<uint8_t> aad1, cipher1;
    seal_at_current(aad1, cipher1);  // sender counter 1 → 2
    std::vector<uint8_t> forged_aad = aad1;
    set_aad_seq(forged_aad, 0x0000FFFFFFFFFFFFULL);  // bit-flip-style corruption
    CHECK(receiver.open(cipher1.data(), cipher1.size(), forged_aad.data(), forged_aad.size(),
                        recovered.data(), recovered.size()) < 0);

    // 3. The next legitimate frame (counter 2) must still decrypt. On the old code this
    //    returned "too old" because step 2 had poisoned recv_highest_.
    std::vector<uint8_t> aad2, cipher2;
    seal_at_current(aad2, cipher2);  // sender counter 2 → 3
    CHECK(receiver.open(cipher2.data(), cipher2.size(), aad2.data(), aad2.size(),
                        recovered.data(), recovered.size()) == static_cast<long>(plain.size()));
}

// Regression for the 16-bit seq wrap: with a full 64-bit wire counter, sealing/opening
// across the old u16 boundary (65,535 → 65,536) must keep decrypting. On the old code the
// nonce desynced at the wrap and every frame failed auth permanently.
static void test_seq_past_16bit_boundary() {
    uint8_t key[crypto_aead_chacha20poly1305_ietf_KEYBYTES];
    crypto_generichash(key, sizeof(key),
                       reinterpret_cast<const uint8_t*>("wrap-key"), 8, nullptr, 0);
    SodiumMediaCrypto sender(key);
    SodiumMediaCrypto receiver(key);

    std::vector<uint8_t> plain(48, 0x6B);
    std::vector<uint8_t> recovered(plain.size());
    std::vector<uint8_t> cipher(plain.size() + crypto_aead_chacha20poly1305_ietf_ABYTES);

    bool all_ok = true;
    for (uint64_t i = 0; i < 70000; ++i) {  // crosses 65,536
        auto aad = make_aad_ssrc(0x1234, sender.peek_send_counter());
        long s = sender.seal(plain.data(), plain.size(), aad.data(), aad.size(),
                             cipher.data(), cipher.size());
        if (s < 0) { all_ok = false; break; }
        long o = receiver.open(cipher.data(), static_cast<size_t>(s), aad.data(), aad.size(),
                               recovered.data(), recovered.size());
        if (o != static_cast<long>(plain.size())) { all_ok = false; break; }
    }
    CHECK(all_ok);
}

int main() {
    if (sodium_init() < 0) {
        std::printf("FAIL: sodium_init failed\n");
        return 1;
    }

    test_seal_open_round_trip();
    test_anti_replay();
    test_tamper_detection();
    test_multiple_packets();
    test_relay_interleaved_reseal();
    test_corrupted_seq_does_not_poison_window();
    test_seq_past_16bit_boundary();

    if (g_failures == 0) {
        std::printf("media_aead: all tests passed\n");
        return 0;
    }
    std::printf("media_aead: %d test(s) FAILED\n", g_failures);
    return 1;
}