#include <gtest/gtest.h>
#include <tuple>

#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <type_traits>

#include "dpf.hpp"
#include "simde/simde/x86/avx2.h"

namespace
{

bool blocks_equal(simde__m128i a, simde__m128i b)
{
    return simde_mm_movemask_epi8(simde_mm_cmpeq_epi8(a, b)) == 0xFFFF;
}

simde__m128i block_from_lanes(std::uint64_t lo, std::uint64_t hi)
{
    return simde_mm_set_epi64x(static_cast<long long>(hi),
        static_cast<long long>(lo));
}

}  // namespace

TEST(AesCcrPrg, SigmaRoundTrip)
{
    using prg = dpf::prg::aes128_ccr;
    const simde__m128i samples[] = {
        block_from_lanes(0, 0),
        block_from_lanes(1, 0),
        block_from_lanes(0, 1),
        block_from_lanes(0x0123456789abcdefull, 0xfedcba9876543210ull),
        block_from_lanes(~0ull, ~0ull),
    };
    for (simde__m128i x : samples)
    {
        EXPECT_TRUE(blocks_equal(prg::sigma_inv(prg::sigma(x)), x));
        EXPECT_TRUE(blocks_equal(prg::sigma(prg::sigma_inv(x)), x));
        EXPECT_TRUE(blocks_equal(prg::sigma_prime(x),
            simde_mm_xor_si128(prg::sigma(x), x)));
    }
}

TEST(AesCcrPrg, HashIsAesMmoOfSigma)
{
    using prg = dpf::prg::aes128_ccr;
    simde__m128i x = block_from_lanes(0x1111222233334444ull, 0x5555666677778888ull);
    EXPECT_TRUE(blocks_equal(prg::hash(x),
        dpf::prg::aes128::eval(prg::sigma(x), 0)));
    EXPECT_TRUE(blocks_equal(prg::H(x), prg::hash(x)));
    EXPECT_TRUE(blocks_equal(prg::eval(x, 0), prg::hash(x)));
    EXPECT_TRUE(blocks_equal(prg::eval(x, 3),
        dpf::prg::aes128::eval(prg::sigma(x), 3)));
}

TEST(AesCcrPrg, Eval01IsHalfTreeChildren)
{
    using prg = dpf::prg::aes128_ccr;
    simde__m128i seed = block_from_lanes(0x0123456789abcdefull, 0xfedcba9876543210ull);
    auto kids = prg::eval01(seed);
    const simde__m128i h = prg::hash(seed);
    EXPECT_TRUE(blocks_equal(kids[0], h));
    EXPECT_TRUE(blocks_equal(kids[1], simde_mm_xor_si128(h, seed)));
    EXPECT_TRUE(blocks_equal(simde_mm_xor_si128(kids[0], kids[1]), seed));
}

TEST(AesCcrPrg, TwoTweakDiffersFromHalfExpand)
{
    using prg = dpf::prg::aes128_ccr;
    simde__m128i seed = block_from_lanes(7, 9);
    auto half = prg::eval01(seed);
    auto tw = prg::eval01_twotweak(seed);
    const simde__m128i base = dpf::unset_lo_bit(seed);
    EXPECT_TRUE(blocks_equal(tw[0], prg::hash(base)));
    EXPECT_TRUE(blocks_equal(tw[1], prg::hash(dpf::set_lo_bit(base))));
    EXPECT_FALSE(blocks_equal(half[0], tw[0]));  // half hashes full seed
    EXPECT_FALSE(blocks_equal(half[1], tw[1]));
}

TEST(AesCcrPrg, HashX4MatchesScalar)
{
    using prg = dpf::prg::aes128_ccr;
    alignas(16) simde__m128i in[4] = {
        block_from_lanes(1, 2),
        block_from_lanes(3, 4),
        block_from_lanes(5, 6),
        block_from_lanes(7, 8),
    };
    alignas(16) simde__m128i out[4];
    prg::hash_x4(in, out);
    for (int i = 0; i < 4; ++i)
        EXPECT_TRUE(blocks_equal(out[i], prg::hash(in[i]))) << i;
}

TEST(AesCcrPrg, Eval01X4MatchesScalar)
{
    using prg = dpf::prg::aes128_ccr;
    alignas(16) simde__m128i seeds[4] = {
        block_from_lanes(10, 20),
        block_from_lanes(30, 40),
        block_from_lanes(50, 60),
        block_from_lanes(70, 80),
    };
    alignas(16) simde__m128i left[4], right[4];
    prg::eval01_x4(seeds, left, right);
    for (int i = 0; i < 4; ++i)
    {
        auto kids = prg::eval01(seeds[i]);
        EXPECT_TRUE(blocks_equal(left[i], kids[0])) << i;
        EXPECT_TRUE(blocks_equal(right[i], kids[1])) << i;
    }
}

TEST(AesCcrPrg, BulkEvalMatchesScalarAndRejectsAWrappingSpan)
{
    using prg = dpf::prg::aes128_ccr;
    simde__m128i seed = block_from_lanes(0xabcdu, 0x1234u);
    alignas(16) simde__m128i out[3];
    out[0] = block_from_lanes(1, 1);
    prg::eval(seed, out, 0);
    EXPECT_TRUE(blocks_equal(out[0], block_from_lanes(1, 1)));
    prg::eval(seed, out, 3, 5);
    for (int i = 0; i < 3; ++i)
        EXPECT_TRUE(blocks_equal(out[i], prg::eval(seed, static_cast<std::uint32_t>(5 + i))))
            << i;
    EXPECT_THROW(prg::eval(seed, out, 2, 0xffffffffu), std::invalid_argument);
}

TEST(AesCcrPrg, HalfTreeTagIsPresent)
{
    static_assert(std::is_void_v<dpf::prg::aes128_ccr::half_tree_tag>);
    SUCCEED();
}