/*
 * 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 <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);
    }
}

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

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