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

#include <cstdint>
#include <cstring>
#include <iterator>
#include <stdexcept>

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

simde__m128i block_from_bytes(const std::uint8_t * p)
{
    simde__m128i x;
    std::memcpy(&x, p, sizeof(x));
    return x;
}

// RFC 8439 §2.3.2. Counter = 1, nonce = 00:00:00:09:00:00:00:4a:00:00:00:00.
constexpr std::uint8_t k_rfc_keystream[64] = {
    0x10, 0xf1, 0xe7, 0xe4, 0xd1, 0x3b, 0x59, 0x15,
    0x50, 0x0f, 0xdd, 0x1f, 0xa3, 0x20, 0x71, 0xc4,
    0xc7, 0xd1, 0xf4, 0xc7, 0x33, 0xc0, 0x68, 0x03,
    0x04, 0x22, 0xaa, 0x9a, 0xc3, 0xd4, 0x6c, 0x4e,
    0xd2, 0x82, 0x64, 0x46, 0x07, 0x9f, 0xaa, 0x09,
    0x14, 0xc2, 0xd7, 0x05, 0xd9, 0x8b, 0x02, 0xa2,
    0xb5, 0x12, 0x9c, 0xd1, 0xde, 0x16, 0x4e, 0xb9,
    0xcb, 0xd0, 0x83, 0xe8, 0xa2, 0x50, 0x3c, 0x4e
};

}  // namespace

TEST(ChachaPrg, Rfc8439Block)
{
    const std::uint32_t key[8] = {
        0x03020100u, 0x07060504u, 0x0b0a0908u, 0x0f0e0d0cu,
        0x13121110u, 0x17161514u, 0x1b1a1918u, 0x1f1e1d1cu
    };
    const std::uint32_t nonce[3] = {0x09000000u, 0x4a000000u, 0x00000000u};
    std::uint8_t out[64];
    dpf::prg::chacha_detail::block<20>(key, 1, nonce, out);
    EXPECT_EQ(0, std::memcmp(out, k_rfc_keystream, sizeof(out)));

    std::uint8_t fewer[64];
    dpf::prg::chacha_detail::block<8>(key, 1, nonce, fewer);
    EXPECT_NE(0, std::memcmp(fewer, k_rfc_keystream, sizeof(fewer)));
}

TEST(ChachaPrg, WideBlockMatchesScalar)
{
    std::uint32_t keys[4][8];
    std::uint32_t counters[4] = {0u, 1u, 5u, 0x00fffff0u};
    for (int lane = 0; lane < 4; ++lane)
    {
        for (int w = 0; w < 8; ++w)
        {
            keys[lane][w] = 0x9e3779b9u * static_cast<std::uint32_t>(lane + 1)
                + static_cast<std::uint32_t>(w) * 0x01000193u;
        }
    }
    std::uint8_t wide[4][64];
    dpf::prg::chacha_detail::block4<20>(keys, counters, wide);
    for (int lane = 0; lane < 4; ++lane)
    {
        std::uint8_t scalar[64];
        dpf::prg::chacha_detail::block<20>(keys[lane], counters[lane],
            dpf::prg::chacha_detail::zero_nonce, scalar);
        EXPECT_EQ(0, std::memcmp(wide[lane], scalar, 64)) << "lane=" << lane;
    }
}

TEST(ChachaPrg, EvalIsKeystreamChunk)
{
    using prg = dpf::prg::chacha20;
    simde__m128i seed = block_from_lanes(0x0123456789abcdefull, 0xfedcba9876543210ull);
    std::uint32_t key[8];
    dpf::prg::chacha_detail::seed_key(seed, key);

    for (psnip_uint32_t pos = 0; pos < 8; ++pos)
    {
        std::uint8_t buf[64];
        dpf::prg::chacha_detail::block<20>(key, pos >> 2,
            dpf::prg::chacha_detail::zero_nonce, buf);
        EXPECT_TRUE(blocks_equal(prg::eval(seed, pos),
            block_from_bytes(buf + 16 * (pos & 3u)))) << "pos=" << pos;
    }

    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));
    EXPECT_FALSE(blocks_equal(dpf::prg::chacha12::eval(seed, 0), kids[0]));
    EXPECT_FALSE(blocks_equal(dpf::prg::chacha8::eval(seed, 0), kids[0]));
}

TEST(ChachaPrg, BulkAndWideAgree)
{
    using prg = dpf::prg::chacha20;
    simde__m128i seed = block_from_lanes(0x0123456789abcdefull, 0xfedcba9876543210ull);

    auto check_bulk = [&](psnip_uint32_t pos, psnip_uint32_t count)
    {
        alignas(16) simde__m128i bulk[32];
        prg::eval(seed, bulk, count, pos);
        for (psnip_uint32_t i = 0; i < count; ++i)
        {
            EXPECT_TRUE(blocks_equal(bulk[i], prg::eval(seed, pos + i)))
                << "pos=" << pos << " i=" << i;
        }
    };
    check_bulk(0, 1);
    check_bulk(0, 2);
    check_bulk(0, 4);
    check_bulk(0, 16);
    check_bulk(1, 20);
    check_bulk(3, 6);
    check_bulk(4, 7);
    check_bulk(0xfffffffeu, 2);

    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, 5);
    prg::eval_x8(seeds, out8, 0);
    prg::eval01_x4(seeds, left, right);
    for (int i = 0; i < 4; ++i)
    {
        EXPECT_TRUE(blocks_equal(out4[i], prg::eval(seeds[i], 5)));
        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], 0)));
    }

    alignas(16) simde__m128i one[1];
    EXPECT_NO_THROW(prg::eval(seed, one, 1, 0xffffffffu));
    EXPECT_TRUE(blocks_equal(one[0], prg::eval(seed, 0xffffffffu)));
    EXPECT_THROW(prg::eval(seed, out4, 2, 0xffffffffu), std::invalid_argument);
    prg::eval(seed, out4, 0, 0xffffffffu);
}

TEST(ChachaPrg, DpfPointAndFull)
{
    using prg = dpf::prg::chacha20;
    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);
        EXPECT_EQ(sum, static_cast<std::uint8_t>(i) == x ? y : 0u) << "i=" << i;
    }

    auto [buf0, iter0] = dpf::eval_full(k0);
    auto [buf1, iter1] = dpf::eval_full(k1);
    (void)buf0;
    (void)buf1;
    std::size_t i = 0;
    auto it0 = std::begin(iter0);
    auto it1 = std::begin(iter1);
    for (; it0 != std::end(iter0); ++it0, ++it1, ++i)
    {
        auto sum = dpf::reconstruct(*it0, *it1);
        EXPECT_EQ(sum, static_cast<std::uint8_t>(i) == x ? y : 0u) << "i=" << i;
    }
    EXPECT_EQ(i, std::size_t{256});
}

TEST(ChachaPrg, ReducedRoundsStillCorrect)
{
    using prg = dpf::prg::chacha8;
    const std::uint8_t x = 0x11;
    const std::uint32_t y = 0xabcdu;
    auto [k0, k1] = dpf::make_dpf<prg, dpf::prg::chacha12>(x, y);
    for (int i = 0; i < 256; ++i)
    {
        auto sum = dpf::reconstruct(
            *dpf::eval_point(k0, static_cast<std::uint8_t>(i)),
            *dpf::eval_point(k1, static_cast<std::uint8_t>(i)));
        EXPECT_EQ(sum, static_cast<std::uint8_t>(i) == x ? y : 0u) << "i=" << i;
    }
}

TEST(ChachaPrg, CounterWrapperCountsEval01)
{
    using prg = dpf::prg::counter_wrapper<dpf::prg::chacha20>;
    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(ChachaPrg, ExpandAndBufferedReplay)
{
    using prg = dpf::prg::chacha20;
    simde__m128i seed = block_from_lanes(0x1111ull, 0x2222ull);
    auto t0 = prg::expand<std::uint32_t, 0>(seed, 3);
    auto t1 = prg::expand<std::uint32_t, 1>(seed, 3);
    EXPECT_EQ(dpf::reconstruct(t0, t1), 0u);

    auto u0 = prg::expand<std::uint64_t, 0>(seed, 0);
    auto u1 = prg::expand<std::uint64_t, 1>(seed, 1);
    EXPECT_NE(u0.raw(), u1.raw());

    dpf::randomness::buffered_prg<prg, std::uint64_t, std::uint64_t> streamed(seed, 8);
    auto s0 = streamed.get<0>();
    auto s1 = streamed.get<1>();
    dpf::randomness::buffered_prg<prg, std::uint64_t, std::uint64_t> replay(seed, 8);
    EXPECT_EQ(replay.at<0>(0), s0);
    EXPECT_EQ(replay.at<1>(0), s1);
    EXPECT_EQ(replay.get<0>(), s0);
    EXPECT_NE(s0, streamed.get<0>());
}

template <typename PRG>
void expect_chacha_comparison_and_interval()
{
    const std::uint8_t alpha = 0x3C;
    auto [c0, c1] = dpf::make_dpf<PRG>(alpha, dpf::lt(std::uint32_t{7}, std::uint32_t{1}));
    for (int i = 0; i < 256; ++i)
    {
        const auto x = static_cast<std::uint8_t>(i);
        const auto got = dpf::reconstruct(
            dpf::eval_point(dpf::cmp, c0, x),
            dpf::eval_point(dpf::cmp, c1, x));
        EXPECT_EQ(static_cast<std::uint64_t>(got) & 0xffffffffu,
            x < alpha ? 7u : 1u)
            << "x=" << i;
    }

    const std::uint8_t r = 9, p = 2, q = 5;
    auto keys = dpf::make_dpf<PRG>(r, dpf::ic(p, q, std::uint32_t{4}, std::uint32_t{1}));
    for (int i = 0; i < 256; ++i)
    {
        const auto x = static_cast<std::uint8_t>(i);
        const auto got = dpf::reconstruct(
            dpf::eval_point(dpf::ic, keys.first, x),
            dpf::eval_point(dpf::ic, keys.second, x));
        const auto w = static_cast<std::uint8_t>(x - r);
        EXPECT_EQ(static_cast<std::uint64_t>(got) & 0xffffffffu,
            (w >= p && w <= q) ? 4u : 1u)
            << "x=" << i;
    }
}

TEST(ChachaPrg, ComparisonAndIntervalKeys)
{
    expect_chacha_comparison_and_interval<dpf::prg::chacha20>();
    expect_chacha_comparison_and_interval<dpf::prg::chacha12>();
    expect_chacha_comparison_and_interval<dpf::prg::chacha8>();
}