#include "ntptun/Bytes.hpp"
#include "ntptun/ChaCha20Poly1305.hpp"
#include "TestUtil.hpp"

#include <string>

using namespace ntptun;

namespace {

AeadNonce nonce_from_hex(const std::string& hex) {
    const auto v = from_hex(hex);
    AeadNonce n{};
    if (v) {
        for (std::size_t i = 0; i < n.size() && i < v->size(); ++i) {
            n[i] = (*v)[i];
        }
    }
    return n;
}

}  // namespace

void run_tests() {
    // RFC 8439 section 2.8.2 AEAD known-answer test. Validates the vendored
    // ChaCha20-Poly1305 implementation against the specification test vector.
    const auto key = from_hex(
        "808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f");
    const AeadNonce nonce = nonce_from_hex("070000004041424344454647");
    const auto aad = from_hex("50515253c0c1c2c3c4c5c6c7");
    const std::string pt_str =
        "Ladies and Gentlemen of the class of '99: If I could offer you only "
        "one tip for the future, sunscreen would be it.";
    const ByteVector plaintext(pt_str.begin(), pt_str.end());
    const auto expected = from_hex(
        "d31a8d34648e60db7b86afbc53ef7ec2a4aded51296e08fea9e2b5a736ee62d6"
        "3dbea45e8ca9671282fafb69da92728b1a71de0a9e060b2905d6a5b67ecd3b36"
        "92ddbd7f2d778b8c9803aee328091b58fab324e4fad675945585808b4831d7bc"
        "3ff4def08e4b7a9de576d26586cec64b6116"
        "1ae10b594f09e26a7e902ecbd0600691");

    CHECK(key.has_value());
    CHECK(aad.has_value());
    CHECK(expected.has_value());

    const ByteVector sealed =
        aead_seal(key->data(), nonce, ByteSpan(*aad), ByteSpan(plaintext));
    CHECK_EQ(sealed.size(), expected->size());
    CHECK(sealed == *expected);

    // Round trip: open recovers the original plaintext.
    const auto opened =
        aead_open(key->data(), nonce, ByteSpan(*aad), ByteSpan(sealed));
    CHECK(opened.has_value());
    CHECK(opened && *opened == plaintext);

    // A single flipped ciphertext bit fails authentication.
    {
        ByteVector bad = sealed;
        bad[0] ^= 0x01;
        CHECK(!aead_open(key->data(), nonce, ByteSpan(*aad), ByteSpan(bad)));
    }
    // A flipped tag bit fails authentication.
    {
        ByteVector bad = sealed;
        bad.back() ^= 0x80;
        CHECK(!aead_open(key->data(), nonce, ByteSpan(*aad), ByteSpan(bad)));
    }
    // Tampered associated data fails authentication.
    {
        ByteVector bad_aad = *aad;
        bad_aad[0] ^= 0xFF;
        CHECK(!aead_open(key->data(), nonce, ByteSpan(bad_aad), ByteSpan(sealed)));
    }
    // A different nonce fails authentication.
    {
        AeadNonce n2 = nonce;
        n2[11] ^= 0x01;
        CHECK(!aead_open(key->data(), n2, ByteSpan(*aad), ByteSpan(sealed)));
    }

    // Empty associated data still round-trips.
    {
        const ByteVector sealed2 =
            aead_seal(key->data(), nonce, ByteSpan{}, ByteSpan(plaintext));
        const auto opened2 =
            aead_open(key->data(), nonce, ByteSpan{}, ByteSpan(sealed2));
        CHECK(opened2 && *opened2 == plaintext);
    }
}

TEST_MAIN()