/*
 * 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 14-byte AAD (voice frame header).
static std::vector<uint8_t> make_aad(uint16_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 (uint16_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 14-byte AAD (voice frame header) with a given ssrc + seq.
static std::vector<uint8_t> make_aad_ssrc(uint32_t ssrc, uint16_t seq) {
    VoiceFrame f;
    f.ssrc = ssrc;
    f.seq  = seq;
    std::vector<uint8_t> aad(kVoiceHeaderSize);
    serialize_header(f, aad.data());
    return aad;
}

// 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, static_cast<uint16_t>(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) {
        uint64_t ctr = r_send.peek_send_counter();
        out_aad[8] = static_cast<uint8_t>((ctr >> 8) & 0xFF);
        out_aad[9] = static_cast<uint8_t>(ctr & 0xFF);
    }
    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)
    }
}

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();

    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;
}