#include <gtest/gtest.h>

#include <array>
#include <cstdint>
#include <cstring>
#include <type_traits>
#include <utility>

#include "dpf.hpp"

namespace
{

using block_t = dpf::prg::aes128::block_type;
constexpr std::size_t kBits = 8;

static block_t g_roots[2];
static int g_ri = 0;
block_t take_root()
{
    return g_roots[g_ri++];
}

void reset_roots(block_t r0, block_t r1)
{
    g_roots[0] = r0;
    g_roots[1] = r1;
    g_ri = 0;
}

std::uint8_t sibling_at(std::uint8_t alpha, std::size_t level)
{
    return static_cast<std::uint8_t>(alpha ^ (1u << (kBits - 1 - level)));
}

template <std::size_t I, std::size_t N, typename Keys>
auto grow_rest(Keys keys, std::uint8_t alpha,
    const std::array<std::uint64_t, N> & betas)
{
    if constexpr (I >= N)
        return keys;
    else
    {
        auto next = dpf::extend(keys.first, keys.second, alpha,
            dpf::at<I + 1>(betas[I]));
        return grow_rest<I + 1, N>(std::move(next), alpha, betas);
    }
}

template <std::size_t N>
auto grow_incrementally(std::uint8_t alpha,
    const std::array<std::uint64_t, N> & betas, block_t r0, block_t r1)
{
    reset_roots(r0, r1);
    auto keys = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<1>(betas[0]));
    return grow_rest<1, N>(std::move(keys), alpha, betas);
}

}  // namespace

TEST(GrowingIdpf, PrefixesOpenToOwnPayload)
{
    const std::uint8_t alpha = 0xB2;
    const std::array<std::uint64_t, kBits> betas = {
        0x11ull, 0x2222ull, 0x333333ull, 0x44444444ull,
        0x5555555555ull, 0x666666666666ull, 0x77777777777777ull,
        0x8888888888888888ull};

    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();
    auto [k0, k1] = grow_incrementally(alpha, betas, r0, r1);

    using KT = std::decay_t<decltype(k0)>;
    EXPECT_EQ(KT::depth, kBits - 1);
    EXPECT_EQ(KT::num_outputs, kBits);

    auto open = [&](auto out_ic, std::uint8_t x) {
        return dpf::reconstruct(
            *dpf::eval_point(out_ic, k0, x),
            *dpf::eval_point(out_ic, k1, x));
    };

    [&]<std::size_t... L>(std::index_sequence<L...>) {
        (([&] {
            EXPECT_EQ(open(dpf::out<L>, alpha), betas[L]) << "prefix " << L;
            EXPECT_EQ(open(dpf::out<L>, sibling_at(alpha, L)), 0ull)
                << "sibling at prefix " << L;
        }()), ...);
    }(std::make_index_sequence<kBits>{});
}

TEST(GrowingIdpf, ExtendMatchesFullySpecified)
{
    const std::uint8_t alpha = 0x6D;
    const std::array<std::uint64_t, kBits> betas = {
        1001, 2002, 3003, 4004, 5005, 6006, 7007, 8008};

    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();

    auto [k0, k1] = grow_incrementally(alpha, betas, r0, r1);

    reset_roots(r0, r1);
    auto [ref0, ref1] = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<1>(betas[0]), dpf::at<2>(betas[1]), dpf::at<3>(betas[2]),
        dpf::at<4>(betas[3]), dpf::at<5>(betas[4]), dpf::at<6>(betas[5]),
        dpf::at<7>(betas[6]), dpf::at<8>(betas[7]));

    using KT = std::decay_t<decltype(k0)>;
    static_assert(std::is_same_v<KT, std::decay_t<decltype(ref0)>>);
    EXPECT_EQ(std::memcmp(k0.correction_words().data(),
        ref0.correction_words().data(),
        sizeof(typename KT::correction_words_array)), 0);
    EXPECT_EQ(std::memcmp(k0.correction_advice().data(),
        ref0.correction_advice().data(),
        sizeof(typename KT::correction_advice_array)), 0);
    EXPECT_EQ(std::memcmp(&k0.root(), &ref0.root(), sizeof(block_t)), 0);
    EXPECT_EQ(std::memcmp(&k1.root(), &ref1.root(), sizeof(block_t)), 0);

    [&]<std::size_t... L>(std::index_sequence<L...>) {
        (([&] {
            EXPECT_EQ(std::memcmp(&k0.template leaf<L>(),
                &ref0.template leaf<L>(), sizeof(k0.template leaf<L>())), 0)
                << "leaf0 " << L;
            EXPECT_EQ(std::memcmp(&k1.template leaf<L>(),
                &ref1.template leaf<L>(), sizeof(k1.template leaf<L>())), 0)
                << "leaf1 " << L;
            EXPECT_EQ(dpf::reconstruct(
                    *dpf::eval_point(dpf::out<L>, k0, alpha),
                    *dpf::eval_point(dpf::out<L>, k1, alpha)),
                betas[L])
                << "open " << L;
        }()), ...);
    }(std::make_index_sequence<kBits>{});
}

TEST(GrowingIdpf, AddOutputWildcardOnExistingLevel)
{
    const std::uint8_t alpha = 0xA5;
    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();
    reset_roots(r0, r1);
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<2>(std::uint64_t{42}));
    // at<2>(u64) → level 1. at<3>(u32) → level 1 (lg_opl=2), new group.
    auto [a0, a1] = dpf::add_output(k0, k1, alpha,
        dpf::at<3>(dpf::wildcard_value<std::uint32_t>{}));
    using AT = std::decay_t<decltype(a0)>;
    EXPECT_EQ(AT::depth, 1u);
    EXPECT_EQ(AT::num_outputs, 2u);
    EXPECT_TRUE(AT::wildcard_bits[1]);

    auto o0 = dpf::eval_point(dpf::out<0>, a0, alpha);
    auto o1 = dpf::eval_point(dpf::out<0>, a1, alpha);
    EXPECT_EQ(dpf::reconstruct(*o0, *o1), 42ull);
}

TEST(GrowingIdpf, DeepestPrefixMatchesClassicDpf)
{
    const std::uint8_t alpha = 0xC3;
    const std::array<std::uint64_t, kBits> betas = {
        1, 2, 4, 8, 16, 32, 64, 128};
    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();
    auto [k0, k1] = grow_incrementally(alpha, betas, r0, r1);

    auto recon = [&](std::uint8_t x) {
        return dpf::reconstruct(
            *dpf::eval_point(dpf::out<kBits - 1>, k0, x),
            *dpf::eval_point(dpf::out<kBits - 1>, k1, x));
    };

    auto [c0, c1] = dpf::make_dpf(alpha, betas.back());
    for (std::uint32_t x = 0; x < 256; ++x)
    {
        const auto want = dpf::reconstruct(
            *dpf::eval_point(c0, static_cast<std::uint8_t>(x)),
            *dpf::eval_point(c1, static_cast<std::uint8_t>(x)));
        EXPECT_EQ(recon(static_cast<std::uint8_t>(x)),
            static_cast<std::uint64_t>(want))
            << "x=" << x;
    }
}

TEST(GrowingIdpf, ChaChaPrgOpensThePrefix)
{
    using prg = dpf::prg::chacha20;
    using block = prg::block_type;
    static block roots[2];
    static int ri = 0;
    auto take = +[]() -> block { return roots[ri++]; };
    const std::uint8_t alpha = 0x3C;
    const std::uint64_t beta = 0x1234;
    roots[0] = dpf::uniform_sample<block>();
    roots[1] = dpf::uniform_sample<block>();
    ri = 0;
    auto [k0, k1] = dpf::make_dpf<prg, prg>(alpha,
        dpf::root_sampler_t<prg>{take}, dpf::at<1>(beta));
    const auto opened = dpf::reconstruct(
        *dpf::eval_point(dpf::out<0>, k0, alpha),
        *dpf::eval_point(dpf::out<0>, k1, alpha));
    EXPECT_EQ(opened, beta);
}

namespace
{

struct GrowPad
{
    simde__m128i block() { return dpf::uniform_sample<simde__m128i>(); }
    std::uint8_t bit()
    {
        return static_cast<std::uint8_t>(
            dpf::uniform_sample<unsigned char>() & 1u);
    }
};

}  // namespace

TEST(GrowingIdpf, MemoizerExtendMatchesRewalk)
{
    const std::uint8_t alpha = 0xB2;
    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();
    reset_roots(r0, r1);
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<1>(std::uint64_t{7}));

    auto [d0, d1] = dpf::extend(k0, k1, alpha, dpf::at<2>(std::uint64_t{9}));

    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    (void)dpf::eval_point(dpf::out<0>, k0, alpha, m0);
    (void)dpf::eval_point(dpf::out<0>, k1, alpha, m1);
    auto [m_k0, m_k1] =
        dpf::extend(k0, k1, m0, m1, alpha, dpf::at<2>(std::uint64_t{9}));

    EXPECT_EQ(0, std::memcmp(d0.correction_words().data(),
                     m_k0.correction_words().data(),
                     sizeof(d0.correction_words())));
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<0>, m_k0, alpha),
                  *dpf::eval_point(dpf::out<0>, m_k1, alpha)),
        7ull);
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<1>, m_k0, alpha),
                  *dpf::eval_point(dpf::out<1>, m_k1, alpha)),
        9ull);
    (void)d1;
}

TEST(GrowingIdpf, ExtendDsMatchesDealer)
{
    const std::uint8_t alpha = 0xB2;
    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();
    reset_roots(r0, r1);
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<1>(std::uint64_t{7}));
    auto [d0, d1] = dpf::extend(k0, k1, alpha, dpf::at<2>(std::uint64_t{9}));

    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    (void)dpf::eval_point(dpf::out<0>, k0, alpha, m0);
    (void)dpf::eval_point(dpf::out<0>, k1, alpha, m1);
    GrowPad pads{};
    dpf::local_cw_protocol<GrowPad> proto{pads};
    auto [s0, s1] = dpf::extend_ds(k0, k1, m0, m1, alpha, std::uint8_t{0}, proto,
        dpf::at<2>(std::uint64_t{9}));

    EXPECT_EQ(0, std::memcmp(d0.correction_words().data(),
                     s0.correction_words().data(),
                     sizeof(d0.correction_words())));
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<1>, s0, alpha),
                  *dpf::eval_point(dpf::out<1>, s1, alpha)),
        9ull);
    (void)d1;
}

TEST(GrowingIdpf, AddOutputDsWildcard)
{
    const std::uint8_t alpha = 0xA5;
    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();
    reset_roots(r0, r1);
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<2>(std::uint64_t{42}));
    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    (void)dpf::eval_point(dpf::out<0>, k0, alpha, m0);
    (void)dpf::eval_point(dpf::out<0>, k1, alpha, m1);
    GrowPad pads{};
    dpf::local_cw_protocol<GrowPad> proto{pads};
    auto [a0, a1] = dpf::add_output_ds(k0, k1, m0, m1, alpha, std::uint8_t{0},
        proto, dpf::at<3>(dpf::wildcard_value<std::uint32_t>{}));
    EXPECT_EQ(std::decay_t<decltype(a0)>::num_outputs, 2u);
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<0>, a0, alpha),
                  *dpf::eval_point(dpf::out<0>, a1, alpha)),
        42ull);
}

TEST(GrowingIdpf, ShallowMemoizerThrows)
{
    const std::uint8_t alpha = 0x10;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<2>(std::uint64_t{1}));
    using bare = dpf::unwrap_party_key_t<std::decay_t<decltype(k0)>>;
    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    const bare & b0 = k0;
    const bare & b1 = k1;
    dpf::detail::ensure_level(b0, alpha, m0, 0);
    dpf::detail::ensure_level(b1, alpha, m1, 0);
    EXPECT_THROW(
        (void)dpf::extend(k0, k1, m0, m1, alpha, dpf::at<3>(std::uint64_t{2})),
        std::invalid_argument);
}

TEST(GrowingIdpfDeath, MismatchedKeyPairThrows)
{
    const std::uint8_t alpha = 0x55;
    // depth >= 1 so correction-word memcmp runs
    auto [k0, _] = dpf::make_dpf(alpha, dpf::at<2>(std::uint64_t{1}));
    auto [__, k1] = dpf::make_dpf(alpha, dpf::at<2>(std::uint64_t{1}));
    EXPECT_THROW(
        (void)dpf::extend(k0, k1, alpha, dpf::at<3>(std::uint64_t{2})),
        std::invalid_argument);
}

TEST(GrowingIdpfDeath, WrongPathBitThrows)
{
    const std::uint8_t alpha = 0x80;  // MSB = 1
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<1>(std::uint64_t{1}));
    // depth 0; programmed bit is 1; flip it
    EXPECT_THROW(
        (void)dpf::extend(k0, k1, /*bit=*/false, alpha,
            dpf::at<2>(std::uint64_t{2})),
        std::invalid_argument);
}

TEST(GrowingIdpfDeath, SpecNotOnNewDepthThrows)
{
    const std::uint8_t alpha = 0x11;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<1>(std::uint64_t{1}));
    // Deepens by one via at<2>, but also plants at<1> on the old level.
    EXPECT_THROW(
        (void)dpf::extend(k0, k1, alpha, dpf::at<2>(std::uint64_t{9}),
            dpf::at<1>(std::uint64_t{8})),
        std::invalid_argument);
}

TEST(GrowingIdpf, ExtendSaturatesEightBitDomain)
{
    const std::uint8_t alpha = 0x01;
    const std::array<std::uint64_t, 8> betas = {1, 2, 3, 4, 5, 6, 7, 8};
    block_t r0 = dpf::uniform_sample<block_t>();
    block_t r1 = dpf::uniform_sample<block_t>();
    auto [k0, k1] = grow_incrementally(alpha, betas, r0, r1);
    using KT = std::decay_t<decltype(k0)>;
    EXPECT_EQ(KT::depth, 7u);
    EXPECT_EQ(KT::depth + 1, dpf::utils::bitlength_of_v<std::uint8_t>);
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<7>, k0, alpha),
                  *dpf::eval_point(dpf::out<7>, k1, alpha)),
        8ull);
}

TEST(GrowingIdpfDeath, MemoizerOnSiblingCorruptsNewSlot)
{
    const std::uint8_t alpha = 0xB2;
    const std::uint8_t sibling = static_cast<std::uint8_t>(alpha ^ 0x80);
    // depth 1 so the frontier is past the root
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<2>(std::uint64_t{7}));
    auto [good0, good1] =
        dpf::extend(k0, k1, alpha, dpf::at<3>(std::uint64_t{9}));

    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    (void)dpf::eval_point(dpf::out<0>, k0, sibling, m0);
    (void)dpf::eval_point(dpf::out<0>, k1, sibling, m1);
    auto [bad0, bad1] =
        dpf::extend(k0, k1, m0, m1, alpha, dpf::at<3>(std::uint64_t{9}));

    EXPECT_NE(0, std::memcmp(good0.correction_words().data(),
                     bad0.correction_words().data(),
                     sizeof(good0.correction_words())));
    const auto got = dpf::reconstruct(*dpf::eval_point(dpf::out<1>, bad0, alpha),
        *dpf::eval_point(dpf::out<1>, bad1, alpha));
    EXPECT_NE(got, 9ull);
    (void)good1;
}

TEST(GrowingIdpfDeath, DsWrongSharesDisagreeWithDealer)
{
    const std::uint8_t alpha = 0xB2;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<2>(std::uint64_t{7}));
    auto [good0, good1] =
        dpf::extend(k0, k1, alpha, dpf::at<3>(std::uint64_t{9}));

    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    (void)dpf::eval_point(dpf::out<0>, k0, alpha, m0);
    (void)dpf::eval_point(dpf::out<0>, k1, alpha, m1);
    GrowPad pads{};
    dpf::local_cw_protocol<GrowPad> proto{pads};
    // Flip the bit used at the extend level (depth 1 → second MSB).
    const std::uint8_t wrong = static_cast<std::uint8_t>(alpha ^ 0x40);
    auto [bad0, bad1] = dpf::extend_ds(k0, k1, m0, m1, wrong, std::uint8_t{0},
        proto, dpf::at<3>(std::uint64_t{9}));

    EXPECT_NE(0, std::memcmp(good0.correction_words().data(),
                     bad0.correction_words().data(),
                     sizeof(good0.correction_words())));
    (void)good1;
    (void)bad1;
}

TEST(GrowingIdpf, ExhaustiveDomainAfterMemoGrow)
{
    const std::uint8_t alpha = 0x5A;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<1>(std::uint64_t{3}));
    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    (void)dpf::eval_point(dpf::out<0>, k0, alpha, m0);
    (void)dpf::eval_point(dpf::out<0>, k1, alpha, m1);
    auto [g0, g1] =
        dpf::extend(k0, k1, m0, m1, alpha, dpf::at<2>(std::uint64_t{11}));
    auto [d0, d1] =
        dpf::extend(k0, k1, alpha, dpf::at<2>(std::uint64_t{11}));

    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xa = static_cast<std::uint8_t>(x);
        EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<0>, g0, xa),
                      *dpf::eval_point(dpf::out<0>, g1, xa)),
            dpf::reconstruct(*dpf::eval_point(dpf::out<0>, d0, xa),
                *dpf::eval_point(dpf::out<0>, d1, xa)))
            << "out0 x=" << x;
        EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<1>, g0, xa),
                      *dpf::eval_point(dpf::out<1>, g1, xa)),
            dpf::reconstruct(*dpf::eval_point(dpf::out<1>, d0, xa),
                *dpf::eval_point(dpf::out<1>, d1, xa)))
            << "out1 x=" << x;
    }
}

TEST(GrowingIdpf, DsAndDealerAgreeFullDomain)
{
    const std::uint8_t alpha = 0x3C;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<1>(std::uint64_t{4}));
    auto [d0, d1] =
        dpf::extend(k0, k1, alpha, dpf::at<2>(std::uint64_t{8}));
    auto m0 = dpf::make_basic_path_memoizer(k0);
    auto m1 = dpf::make_basic_path_memoizer(k1);
    (void)dpf::eval_point(dpf::out<0>, k0, alpha, m0);
    (void)dpf::eval_point(dpf::out<0>, k1, alpha, m1);
    GrowPad pads{};
    dpf::local_cw_protocol<GrowPad> proto{pads};
    auto [s0, s1] = dpf::extend_ds(k0, k1, m0, m1, alpha, std::uint8_t{0}, proto,
        dpf::at<2>(std::uint64_t{8}));

    EXPECT_EQ(0, std::memcmp(d0.correction_words().data(),
                     s0.correction_words().data(),
                     sizeof(d0.correction_words())));
    EXPECT_EQ(0, std::memcmp(d0.correction_advice().data(),
                     s0.correction_advice().data(),
                     sizeof(d0.correction_advice())));
    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xa = static_cast<std::uint8_t>(x);
        EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(dpf::out<1>, s0, xa),
                      *dpf::eval_point(dpf::out<1>, s1, xa)),
            dpf::reconstruct(*dpf::eval_point(dpf::out<1>, d0, xa),
                *dpf::eval_point(dpf::out<1>, d1, xa)))
            << "x=" << x;
    }
}

TEST(GrowingIdpf, EmptyAddOutputThrows)
{
    const std::uint8_t alpha = 0x01;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::at<1>(std::uint64_t{1}));
    EXPECT_THROW((void)dpf::add_output(k0, k1, alpha), std::invalid_argument);
}