#include <gtest/gtest.h>

#include <algorithm>
#include <cstdint>
#include <iterator>
#include <limits>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

#include "dpf.hpp"

namespace
{

template <typename A, typename B>
auto recon(const A & a, const B & b)
{
    if constexpr (dpf::is_secret_share_v<std::decay_t<A>>
        && dpf::is_secret_share_v<std::decay_t<B>>)
        return dpf::reconstruct(a, b);
    else
    {
        using T = std::common_type_t<std::decay_t<A>, std::decay_t<B>>;
        if constexpr (std::is_integral_v<T> && std::is_unsigned_v<T>)
            return static_cast<T>(a - b);
        else
            return a - b;
    }
}

template <typename Key0, typename Key1, typename InputT>
void assign_input_local(Key0 & k0, Key1 & k1, InputT alpha)
{
    const InputT a0 = static_cast<InputT>(0x12);
    const InputT a1 = static_cast<InputT>(alpha - a0);
    const auto sh0 = k0.offset_x.compute_and_get_share(a0);
    const auto sh1 = k1.offset_x.compute_and_get_share(a1);
    k0.offset_x.reconstruct(sh1);
    k1.offset_x.reconstruct(sh0);
}

template <typename InputT>
std::size_t inclusive_span(InputT from, InputT to)
{
    constexpr auto bits = dpf::utils::bitlength_of_v<InputT>;
    constexpr auto to_int = dpf::utils::to_integral_type<InputT>{};
    auto span = to_int(to) - to_int(from);
    if constexpr (bits < dpf::utils::bitlength_of_v<decltype(span)>)
        span &= (decltype(span){1} << bits) - 1;
    return static_cast<std::size_t>(span) + 1;
}

template <typename InputT, typename Fn>
void for_inclusive_wrap(InputT from, InputT to, Fn && fn)
{
    constexpr auto to_int = dpf::utils::to_integral_type<InputT>{};
    const auto n = inclusive_span(from, to);
    InputT x = from;
    for (std::size_t i = 0; i < n; ++i)
    {
        fn(x);
        x = static_cast<InputT>(to_int(x) + 1);
    }
}

}  // namespace

// ---------------------------------------------------------------------------
// utils: interval_wraps / split_leaf_nodes / get_leafnodes
// ---------------------------------------------------------------------------

TEST(WrapUtils, IntervalWrapsUnsigned)
{
    EXPECT_FALSE(dpf::utils::interval_wraps(std::uint8_t{10}, std::uint8_t{20}, 8));
    EXPECT_TRUE(dpf::utils::interval_wraps(std::uint8_t{200}, std::uint8_t{10}, 8));
    EXPECT_FALSE(dpf::utils::interval_wraps(std::uint8_t{0}, std::uint8_t{255}, 8));
    EXPECT_TRUE(dpf::utils::interval_wraps(std::uint8_t{255}, std::uint8_t{0}, 8));
    EXPECT_FALSE(dpf::utils::interval_wraps(0, 0, 0));
}

TEST(WrapUtils, IntervalWrapsMaskedLowBits)
{
    // Only low 4 bits participate; 0x1F vs 0x02 wraps in nibble space.
    EXPECT_TRUE(dpf::utils::interval_wraps(0x1Fu, 0x02u, 4));
    EXPECT_FALSE(dpf::utils::interval_wraps(0x12u, 0x1Eu, 4));
}

TEST(WrapUtils, SplitNonWrapSingleSegment)
{
    auto segs = dpf::utils::split_leaf_nodes(std::uint8_t{3}, std::uint8_t{10}, 8, false);
    ASSERT_EQ(segs.n, 1u);
    EXPECT_EQ(segs.seg[0].from_node, 3);
    EXPECT_EQ(segs.seg[0].to_node, 10);
    EXPECT_EQ(segs.seg[0].count, 7u);
    EXPECT_EQ(segs.total, 7u);
}

TEST(WrapUtils, SplitWrapTwoSegments)
{
    // depth < bitwidth so domain_end is representable.
    auto segs = dpf::utils::split_leaf_nodes(
        std::uint16_t{250}, std::uint16_t{4}, 8, true);
    ASSERT_EQ(segs.n, 2u);
    EXPECT_EQ(segs.seg[0].from_node, 250);
    EXPECT_EQ(segs.seg[0].to_node, 256);
    EXPECT_EQ(segs.seg[0].count, 6u);
    EXPECT_EQ(segs.seg[1].from_node, 0);
    EXPECT_EQ(segs.seg[1].to_node, 4);
    EXPECT_EQ(segs.seg[1].count, 4u);
    EXPECT_EQ(segs.total, 10u);
}

TEST(WrapUtils, SplitWrapShallowDepth)
{
    // depth 6 => domain_end 64
    auto segs = dpf::utils::split_leaf_nodes(std::uint8_t{60}, std::uint8_t{3}, 6, true);
    ASSERT_EQ(segs.n, 2u);
    EXPECT_EQ(segs.seg[0].from_node, 60);
    EXPECT_EQ(segs.seg[0].to_node, 64);
    EXPECT_EQ(segs.seg[0].count, 4u);
    EXPECT_EQ(segs.seg[1].from_node, 0);
    EXPECT_EQ(segs.seg[1].to_node, 3);
    EXPECT_EQ(segs.seg[1].count, 3u);
    EXPECT_EQ(segs.total, 7u);
}

TEST(WrapUtils, InputWrapWithFromNodeLeToNodeStillSplits)
{
    // Packing can make from_node <= to_node while the input still wraps.
    // Collapsing would omit the duplicated shared leaf.
    auto segs = dpf::utils::split_leaf_nodes(std::uint8_t{2}, std::uint8_t{5}, 6, true);
    ASSERT_EQ(segs.n, 2u);
    EXPECT_EQ(segs.seg[0].from_node, 2);
    EXPECT_EQ(segs.seg[0].to_node, 64);
    EXPECT_EQ(segs.seg[1].from_node, 0);
    EXPECT_EQ(segs.seg[1].to_node, 5);
    EXPECT_GT(segs.total, 5u - 2u);
}

TEST(WrapUtils, LeafCountAgreesWithSplitTotal)
{
    using In = std::uint8_t;
    auto [k0, k1] = dpf::make_dpf(In{1}, std::uint32_t{1});
    using key_t = std::decay_t<decltype(k0)>;
    auto check = [](In from, In to)
    {
        In ff = from, ft = to;
        dpf::utils::flip_msb_if_signed_integral(ff);
        dpf::utils::flip_msb_if_signed_integral(ft);
        constexpr auto to_int = dpf::utils::to_integral_type<In>{};
        using integral = typename key_t::integral_type;
        const auto from_i = static_cast<integral>(to_int(ff));
        const auto to_i = static_cast<integral>(to_int(ft));
        const bool wraps = dpf::utils::interval_wraps(from_i, to_i, 8);
        const auto segs = dpf::utils::split_leaf_nodes(
            dpf::utils::get_from_node<key_t>(ff),
            dpf::utils::get_to_node<key_t>(ft),
            static_cast<std::size_t>(key_t::depth), wraps);
        EXPECT_EQ((dpf::utils::get_leafnodes_in_output_interval<key_t>(from, to)),
            segs.total);
    };
    check(In{10}, In{20});
    check(In{200}, In{10});
    check(In{10}, In{9});
    check(In{0}, In{255});
    (void)k1;
}

TEST(WrapUtils, LeafCountVariesWithAlignmentWhenPacked)
{
    using In = std::int8_t;
    auto [k0, k1] = dpf::make_dpf(In{0}, std::uint32_t{7});
    using key_t = std::decay_t<decltype(k0)>;
    ASSERT_GT(key_t::outputs_per_leaf, 1u);
    const In from{-40}, to{10};
    const auto n0 = dpf::utils::get_leafnodes_in_output_interval<key_t>(from, to);
    bool saw_larger = false;
    for (int d = 0; d < 64; ++d)
    {
        const In tfrom = static_cast<In>(from + d);
        const In tto = static_cast<In>(to + d);
        const auto n = dpf::utils::get_leafnodes_in_output_interval<key_t>(tfrom, tto);
        if (n > n0)
            saw_larger = true;
    }
    EXPECT_TRUE(saw_larger);
    (void)k1;
}

// ---------------------------------------------------------------------------
// rotation_iterable: indexing past the wrap (the deferred-view bug class)
// ---------------------------------------------------------------------------

TEST(RotationIterable, OperatorBracketMatchesIteratorOrder)
{
    std::vector<int> values{10, 20, 30, 40, 50};
    dpf::rotation_iterable rot(values.begin(), values.end(), 3);
    std::vector<int> by_it;
    for (auto it = rot.begin(); it != rot.end(); ++it)
        by_it.push_back(*it);
    EXPECT_EQ(by_it, (std::vector<int>{40, 50, 10, 20, 30}));

    std::vector<int> by_idx;
    for (std::ptrdiff_t i = 0; i < static_cast<std::ptrdiff_t>(values.size()); ++i)
        by_idx.push_back(rot[i]);
    EXPECT_EQ(by_idx, by_it);
}

TEST(RotationIterable, NegativeAndFullCycleDistanceNormalize)
{
    std::vector<int> values{1, 2, 3, 4};
    dpf::rotation_iterable neg(values.begin(), values.end(), -1);
    EXPECT_EQ(neg.distance(), 3);
    EXPECT_EQ(neg[0], 4);
    EXPECT_EQ(neg[1], 1);

    dpf::rotation_iterable full(values.begin(), values.end(), 8);
    EXPECT_EQ(full.distance(), 0);
    EXPECT_EQ(full[0], 1);
    EXPECT_EQ(full[3], 4);
}

TEST(RotationIterable, WrappingSubrangeViaIndexDoesNotWalkPastEnd)
{
    // Contiguous std::next from rot.begin() past the physical end is wrong
    // for a wrapping logical slice; operator[] with modulo is the safe path.
    std::vector<int> values{0, 1, 2, 3, 4, 5, 6, 7};
    dpf::rotation_iterable rot(values.begin(), values.end(), 5);
    // Logical slice starting at index 6 of the unrotated domain, length 4:
    // rotated indices (6+5)%8 ... => values 3,4,5,6 in rot order from start 6.
    const std::size_t start = 6;
    const std::size_t count = 4;
    const std::size_t n = values.size();
    std::vector<int> got;
    for (std::size_t i = 0; i < count; ++i)
        got.push_back(rot[static_cast<std::ptrdiff_t>((start + i) % n)]);
    EXPECT_EQ(got, (std::vector<int>{3, 4, 5, 6}));
}

TEST(RotationIterable, BidirectionalRoundTrip)
{
    std::vector<int> values{1, 2, 3, 4, 5};
    dpf::rotation_iterable rot(values.begin(), values.end(), 2);
    auto it = rot.begin();
    ++it;
    ++it;
    EXPECT_EQ(*it, 5);
    --it;
    EXPECT_EQ(*it, 4);
    --it;
    EXPECT_EQ(it, rot.begin());
    EXPECT_EQ(*it, 3);
}

// ---------------------------------------------------------------------------
// Eager eval_interval wrap vs pointwise
// ---------------------------------------------------------------------------

TEST(EagerWrap, Uint8MatchesPointwise)
{
    using In = std::uint8_t;
    using Out = std::uint32_t;
    constexpr In from{200}, to{10}, alpha{250};
    constexpr Out beta{0xABCDEF01u};
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    auto [buf0, it0] = dpf::eval_interval(k0, from, to);
    auto [buf1, it1] = dpf::eval_interval(k1, from, to);
    auto a = std::begin(it0);
    auto b = std::begin(it1);
    std::size_t n = 0;
    for_inclusive_wrap(from, to, [&](In x)
    {
        ASSERT_NE(a, std::end(it0));
        ASSERT_NE(b, std::end(it1));
        EXPECT_EQ(recon(*a, *b),
            recon(*dpf::eval_point(k0, x), *dpf::eval_point(k1, x)))
            << "x=" << +x;
        ++a;
        ++b;
        ++n;
    });
    EXPECT_EQ(a, std::end(it0));
    EXPECT_EQ(b, std::end(it1));
    EXPECT_EQ(n, inclusive_span(from, to));
    (void)buf0;
    (void)buf1;
}

TEST(EagerWrap, SameLeafWrapMatchesPointwise)
{
    using In = std::uint8_t;
    using Out = std::uint32_t;
    // from=10, to=9 wraps almost the full domain; opl may share a leaf.
    auto [k0, k1] = dpf::make_dpf(In{40}, Out{0x11111111u});
    auto [buf0, it0] = dpf::eval_interval(k0, In{10}, In{9});
    auto [buf1, it1] = dpf::eval_interval(k1, In{10}, In{9});
    auto a = std::begin(it0);
    auto b = std::begin(it1);
    for_inclusive_wrap(In{10}, In{9}, [&](In x)
    {
        ASSERT_NE(a, std::end(it0));
        ASSERT_NE(b, std::end(it1));
        EXPECT_EQ(recon(*a, *b),
            recon(*dpf::eval_point(k0, x), *dpf::eval_point(k1, x)));
        ++a;
        ++b;
    });
    EXPECT_EQ(a, std::end(it0));
    EXPECT_EQ(b, std::end(it1));
    (void)buf0;
    (void)buf1;
}

// ---------------------------------------------------------------------------
// Wildcard offset: memoizer / buffer must size on tree coordinates
// ---------------------------------------------------------------------------

TEST(OffsetSizing, LogicalMemoizerCanUndersizeAfterAssign)
{
    using In = std::int8_t;
    using Out = std::uint32_t;
    constexpr In from{-40}, to{10};
    using key_t = std::decay_t<decltype(
        std::get<0>(dpf::make_dpf(dpf::wildcard_value<In>{}, Out{7})))>;
    ASSERT_GT(key_t::outputs_per_leaf, 1u);

    const auto n_logical =
        dpf::utils::get_leafnodes_in_output_interval<key_t>(from, to);
    bool found = false;
    for (int trial = 0; trial < 256; ++trial)
    {
        auto [k0, k1] = dpf::make_dpf(dpf::wildcard_value<In>{}, Out{7});
        const In alpha = static_cast<In>(trial - 128);
        assign_input_local(k0, k1, alpha);
        const auto tfrom = k0.offset_x(from);
        const auto tto = k0.offset_x(to);
        const auto n_tree =
            dpf::utils::get_leafnodes_in_output_interval<key_t>(tfrom, tto);

        In ff = tfrom, ft = tto;
        dpf::utils::flip_msb_if_signed_integral(ff);
        dpf::utils::flip_msb_if_signed_integral(ft);
        constexpr auto to_int = dpf::utils::to_integral_type<In>{};
        using integral = typename key_t::integral_type;
        const bool wraps = dpf::utils::interval_wraps(
            static_cast<integral>(to_int(ff)),
            static_cast<integral>(to_int(ft)), 8);
        // Single-segment undersize: each wrap half may still fit in n_logical.
        if (wraps || n_tree <= n_logical)
            continue;

        found = true;
        auto small = dpf::make_basic_interval_memoizer<key_t>(from, to);
        auto buf = dpf::make_output_buffer_for_interval(k0, from, to);
        EXPECT_THROW(
            (void)dpf::eval_interval(k0, from, to, buf, small),
            std::exception);

        auto memo = dpf::make_basic_interval_memoizer(k0, from, to);
        auto buf2 = dpf::make_output_buffer_for_interval(k0, from, to);
        auto it = dpf::eval_interval(k0, from, to, buf2, memo);
        auto it1_buf = dpf::make_output_buffer_for_interval(k1, from, to);
        auto memo1 = dpf::make_basic_interval_memoizer(k1, from, to);
        auto it1 = dpf::eval_interval(k1, from, to, it1_buf, memo1);
        auto a = std::begin(it);
        auto b = std::begin(it1);
        for_inclusive_wrap(from, to, [&](In x)
        {
            ASSERT_NE(a, std::end(it));
            ASSERT_NE(b, std::end(it1));
            EXPECT_EQ(recon(*a, *b),
                recon(*dpf::eval_point(k0, x), *dpf::eval_point(k1, x)));
            ++a;
            ++b;
        });
        break;
    }
    ASSERT_TRUE(found) << "no non-wrapping misaligned offset found";
}

TEST(OffsetSizing, ConvenienceEvalIntervalSurvivesSignedWildcardSweep)
{
    using In = std::int8_t;
    using Out = std::uint32_t;
    constexpr In from{-40}, to{10};
    constexpr Out beta{9};
    int ok = 0;
    for (int trial = 0; trial < 64; ++trial)
    {
        auto [d0, d1] = dpf::make_dpf(dpf::wildcard_value<In>{}, beta);
        assign_input_local(d0, d1, static_cast<In>(trial * 3 - 96));
        auto buf0 = dpf::make_output_buffer_for_interval(d0, from, to);
        auto buf1 = dpf::make_output_buffer_for_interval(d1, from, to);
        // Default memoizer path must not throw after the sizing fix.
        auto it0 = dpf::eval_interval(d0, from, to, buf0);
        auto it1 = dpf::eval_interval(d1, from, to, buf1);
        auto a = std::begin(it0);
        auto b = std::begin(it1);
        for_inclusive_wrap(from, to, [&](In x)
        {
            ASSERT_NE(a, std::end(it0));
            ASSERT_NE(b, std::end(it1));
            EXPECT_EQ(recon(*a, *b),
                recon(*dpf::eval_point(d0, x), *dpf::eval_point(d1, x)))
                << "x=" << +x << " trial=" << trial;
            ++a;
            ++b;
        });
        ++ok;
    }
    EXPECT_EQ(ok, 64);
}

// ---------------------------------------------------------------------------
// Deferred wrap: index-based view vs eager (regression for contiguous-next bug)
// ---------------------------------------------------------------------------

TEST(DeferWrap, WrappingUint8MatchesEagerAndPointwise)
{
    using In = std::uint8_t;
    using Out = std::uint32_t;
    constexpr In from{200}, to{10};
    constexpr Out beta{42};
    for (unsigned alpha_i = 0; alpha_i < 256; alpha_i += 37)
    {
        const In alpha = static_cast<In>(alpha_i);
        auto [d0, d1] = dpf::make_dpf(dpf::wildcard_value<In>{}, beta);
        auto buf0 = dpf::make_output_buffer_for_full(d0);
        auto buf1 = dpf::make_output_buffer_for_full(d1);
        auto def0 = dpf::defer_eval_interval(d0, from, to, buf0);
        auto def1 = dpf::defer_eval_interval(d1, from, to, buf1);
        assign_input_local(d0, d1, alpha);

        auto eager_buf0 = dpf::make_output_buffer_for_interval(d0, from, to);
        auto eager_buf1 = dpf::make_output_buffer_for_interval(d1, from, to);
        auto eager0 = dpf::eval_interval(d0, from, to, eager_buf0);
        auto eager1 = dpf::eval_interval(d1, from, to, eager_buf1);

        auto v0 = def0.get();
        auto v1 = def1.get();
        auto it_a = std::begin(v0);
        auto it_b = std::begin(v1);
        auto it_c = std::begin(eager0);
        auto it_d = std::begin(eager1);
        for_inclusive_wrap(from, to, [&](In x)
        {
            ASSERT_NE(it_a, std::end(v0));
            ASSERT_NE(it_b, std::end(v1));
            ASSERT_NE(it_c, std::end(eager0));
            ASSERT_NE(it_d, std::end(eager1));
            const auto y_def = recon(*it_a, *it_b);
            const auto y_eag = recon(*it_c, *it_d);
            const auto y_pt = recon(*dpf::eval_point(d0, x),
                *dpf::eval_point(d1, x));
            EXPECT_EQ(y_def, y_eag) << "x=" << +x << " alpha=" << +alpha;
            EXPECT_EQ(y_def, y_pt) << "x=" << +x << " alpha=" << +alpha;
            ++it_a;
            ++it_b;
            ++it_c;
            ++it_d;
        });
        EXPECT_EQ(it_a, std::end(v0));
        EXPECT_EQ(it_b, std::end(v1));
    }
}

TEST(DeferWrap, SignedSpanAcrossZeroMatchesEager)
{
    using In = std::int8_t;
    using Out = std::uint32_t;
    constexpr In from{-5}, to{5};
    constexpr Out beta{3};
    for (int trial = 0; trial < 32; ++trial)
    {
        auto [d0, d1] = dpf::make_dpf(dpf::wildcard_value<In>{}, beta);
        auto buf0 = dpf::make_output_buffer_for_full(d0);
        auto buf1 = dpf::make_output_buffer_for_full(d1);
        auto def0 = dpf::defer_eval_interval(d0, from, to, buf0);
        auto def1 = dpf::defer_eval_interval(d1, from, to, buf1);
        assign_input_local(d0, d1, static_cast<In>(trial * 7 - 112));

        auto eager_buf0 = dpf::make_output_buffer_for_interval(d0, from, to);
        auto eager_buf1 = dpf::make_output_buffer_for_interval(d1, from, to);
        auto eager0 = dpf::eval_interval(d0, from, to, eager_buf0);
        auto eager1 = dpf::eval_interval(d1, from, to, eager_buf1);

        auto v0 = def0.get();
        auto v1 = def1.get();
        auto a = std::begin(v0);
        auto b = std::begin(v1);
        auto c = std::begin(eager0);
        auto d = std::begin(eager1);
        while (a != std::end(v0))
        {
            ASSERT_NE(b, std::end(v1));
            ASSERT_NE(c, std::end(eager0));
            ASSERT_NE(d, std::end(eager1));
            EXPECT_EQ(recon(*a, *b), recon(*c, *d));
            ++a;
            ++b;
            ++c;
            ++d;
        }
        EXPECT_EQ(b, std::end(v1));
        EXPECT_EQ(c, std::end(eager0));
        EXPECT_EQ(d, std::end(eager1));
    }
}

TEST(DeferWrap, MultiOplUnalignedWrappingMatchesEager)
{
    using In = std::uint8_t;
    using Out = std::uint64_t;
    constexpr In from{0xF1}, to{0x0E};
    constexpr Out beta{0x55};
    auto [d0, d1] = dpf::make_dpf(dpf::wildcard_value<In>{}, beta);
    ASSERT_GT(decltype(d0)::outputs_per_leaf, std::size_t{1});
    auto buf0 = dpf::make_output_buffer_for_full(d0);
    auto buf1 = dpf::make_output_buffer_for_full(d1);
    auto def0 = dpf::defer_eval_interval(d0, from, to, buf0);
    auto def1 = dpf::defer_eval_interval(d1, from, to, buf1);
    assign_input_local(d0, d1, In{0xA3});

    auto eager_buf0 = dpf::make_output_buffer_for_interval(d0, from, to);
    auto eager_buf1 = dpf::make_output_buffer_for_interval(d1, from, to);
    auto eager0 = dpf::eval_interval(d0, from, to, eager_buf0);
    auto eager1 = dpf::eval_interval(d1, from, to, eager_buf1);

    auto v0 = def0.get();
    auto v1 = def1.get();
    auto a = std::begin(v0);
    auto b = std::begin(v1);
    auto c = std::begin(eager0);
    auto d = std::begin(eager1);
    for_inclusive_wrap(from, to, [&](In x)
    {
        ASSERT_NE(a, std::end(v0));
        ASSERT_NE(b, std::end(v1));
        ASSERT_NE(c, std::end(eager0));
        ASSERT_NE(d, std::end(eager1));
        EXPECT_EQ(recon(*a, *b), recon(*c, *d)) << "x=" << +x;
        EXPECT_EQ(recon(*a, *b),
            recon(*dpf::eval_point(d0, x), *dpf::eval_point(d1, x)));
        ++a;
        ++b;
        ++c;
        ++d;
    });
}

// ---------------------------------------------------------------------------
// Inner-product wrap already covered; keep a wildcard-offset variant
// ---------------------------------------------------------------------------

TEST(InnerProductWrap, WildcardOffsetWrappingMatchesPoints)
{
    using In = std::uint8_t;
    constexpr In from{250}, to{4};
    auto [k0, k1] = dpf::make_dpf(dpf::wildcard_value<In>{}, std::uint32_t{9});
    assign_input_local(k0, k1, In{1});

    std::vector<std::uint32_t> w;
    std::uint64_t expect = 0;
    for_inclusive_wrap(from, to, [&](In x)
    {
        w.push_back(static_cast<std::uint32_t>(w.size() + 1));
        expect += static_cast<std::uint64_t>(
            recon(*dpf::eval_point(k0, x), *dpf::eval_point(k1, x)))
            * w.back();
    });
    EXPECT_EQ(recon(
            dpf::eval_inner_product(dpf::paired, k0, from, to, w),
            dpf::eval_inner_product(dpf::paired, k1, from, to, w)),
        expect);
}