#include <gtest/gtest.h>

#include <cstdint>
#include <cstring>

#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)
{
    simde__m128i x;
    std::uint64_t lane[2] = {lo, hi};
    std::memcpy(&x, lane, sizeof(x));
    return x;
}

}  // namespace

TEST(LowmcPrg, EncryptDecryptRoundtrip)
{
    lowmc::LowMC cipher;
    lowmc::block plain = 0;
    plain[0] = 1;
    plain[127] = 1;
    plain[64] = 1;
    lowmc::block ciphered = cipher.encrypt(plain);
    EXPECT_TRUE(cipher.decrypt(ciphered) == plain);
    EXPECT_TRUE(ciphered != plain);
}

TEST(LowmcPrg, EvalAgreesWithBulk)
{
    using prg = dpf::prg::lowmc128;
    simde__m128i seed = block_from_lanes(0x0123456789abcdefull, 0xfedcba9876543210ull);

    auto kids = prg::eval01(seed);
    EXPECT_TRUE(blocks_equal(kids[0], prg::eval(seed, 0)));
    EXPECT_TRUE(blocks_equal(kids[1], prg::eval(seed, 1)));
    EXPECT_FALSE(blocks_equal(kids[0], kids[1]));
    EXPECT_FALSE(blocks_equal(kids[0], seed));

    alignas(16) simde__m128i bulk[4];
    prg::eval(seed, bulk, 4, 0);
    EXPECT_TRUE(blocks_equal(bulk[0], kids[0]));
    EXPECT_TRUE(blocks_equal(bulk[1], kids[1]));
    EXPECT_TRUE(blocks_equal(bulk[2], prg::eval(seed, 2)));
    EXPECT_TRUE(blocks_equal(bulk[3], prg::eval(seed, 3)));

    alignas(16) simde__m128i seeds[8];
    alignas(16) simde__m128i out4[4];
    alignas(16) simde__m128i out8[8];
    alignas(16) simde__m128i left[4];
    alignas(16) simde__m128i right[4];
    for (int i = 0; i < 8; ++i)
    {
        seeds[i] = block_from_lanes(0x1000u + static_cast<unsigned>(i), 0x2000u);
    }
    prg::eval_x4(seeds, out4, 3);
    prg::eval_x8(seeds, out8, 5);
    prg::eval01_x4(seeds, left, right);
    for (int i = 0; i < 4; ++i)
    {
        EXPECT_TRUE(blocks_equal(out4[i], prg::eval(seeds[i], 3)));
        EXPECT_TRUE(blocks_equal(left[i], prg::eval(seeds[i], 0)));
        EXPECT_TRUE(blocks_equal(right[i], prg::eval(seeds[i], 1)));
    }
    for (int i = 0; i < 8; ++i)
    {
        EXPECT_TRUE(blocks_equal(out8[i], prg::eval(seeds[i], 5)));
    }
}

TEST(LowmcPrg, DpfPoint)
{
    using prg = dpf::prg::lowmc128;
    const std::uint8_t x = 0x2a;
    const std::uint32_t y = 0x01020304;
    auto [k0, k1] = dpf::make_dpf<prg>(x, y);

    for (int i = 0; i < 256; ++i)
    {
        auto y0 = *dpf::eval_point(k0, static_cast<std::uint8_t>(i));
        auto y1 = *dpf::eval_point(k1, static_cast<std::uint8_t>(i));
        auto sum = dpf::reconstruct(y0, y1);
        if (static_cast<std::uint8_t>(i) == x)
        {
            EXPECT_EQ(sum, y);
        }
        else
        {
            EXPECT_EQ(sum, 0u);
        }
    }
}

TEST(LowmcPrg, CounterWrapperCountsEval01)
{
    using prg = dpf::prg::counter_wrapper<dpf::prg::lowmc128>;
    const auto before = prg::count();
    simde__m128i seed = block_from_lanes(0x1111ull, 0x2222ull);
    auto kids = prg::eval01(seed);
    EXPECT_FALSE(blocks_equal(kids[0], kids[1]));
    EXPECT_EQ(prg::count(), before + 2u);

    alignas(16) simde__m128i bulk[4];
    prg::eval(seed, bulk, 4, 0);
    EXPECT_EQ(prg::count(), before + 2u + 4u);
}

TEST(LowmcPrg, CounterWrapperDpfGenAndPoint)
{
    using prg = dpf::prg::counter_wrapper<dpf::prg::lowmc128>;
    const auto before = prg::count();
    const std::uint8_t x = 0x91;
    const std::uint32_t y = 0xdeadbeef;
    auto [k0, k1] = dpf::make_dpf<prg>(x, y);
    EXPECT_GT(prg::count(), before);
    auto y0 = *dpf::eval_point(k0, x);
    auto y1 = *dpf::eval_point(k1, x);
    EXPECT_EQ(dpf::reconstruct(y0, y1), y);
}