/// @file gf2_adversarial_test.cpp
/// @brief Adversarial checks for GF(2^k) output types.
/// @details Full domains, lane boundaries, poisoned high bits, independent
///          polynomial inverses, leaf scaling, shares, multi-output keys,
///          comparisons, and proof tampering.

#include <gtest/gtest.h>

#include "dpf.hpp"

#include <cstdint>
#include <cstring>
#include <iterator>
#include <vector>

namespace
{

using u128 = unsigned __int128;

u128 poly_mul(u128 a, u128 b)
{
    unsigned __int128 p = 0;
    for (int i = 0; i < 128 && b != 0; ++i)
    {
        if ((b & 1) != 0)
            p ^= a;
        a <<= 1;
        b >>= 1;
    }
    return p;
}

u128 poly_quot_rem(u128 & rem, u128 den, int)
{
    unsigned __int128 q = 0;
    int den_bit = -1;
    for (int i = 127; i >= 0; --i)
    {
        if (((den >> i) & 1) != 0)
        {
            den_bit = i;
            break;
        }
    }
    if (den_bit < 0)
        return 0;
    for (;;)
    {
        int bit = -1;
        for (int i = 127; i >= 0; --i)
        {
            if (((rem >> i) & 1) != 0)
            {
                bit = i;
                break;
            }
        }
        if (bit < den_bit)
            break;
        const int sh = bit - den_bit;
        q ^= u128{1} << sh;
        rem ^= den << sh;
    }
    return q;
}

u128 poly_inv(u128 a, u128 mod)
{
    u128 r0 = mod;
    u128 r1 = a;
    u128 s0 = 0;
    u128 s1 = 1;
    while (r1 != 0)
    {
        const u128 q = poly_quot_rem(r0, r1, 0);
        const u128 nr = r0;
        u128 ns = s0 ^ poly_mul(q, s1);
        int mod_bit = -1;
        for (int i = 127; i >= 0; --i)
        {
            if (((mod >> i) & 1) != 0)
            {
                mod_bit = i;
                break;
            }
        }
        while (mod_bit >= 0)
        {
            int bit = -1;
            for (int i = 127; i >= 0; --i)
            {
                if (((ns >> i) & 1) != 0)
                {
                    bit = i;
                    break;
                }
            }
            if (bit < mod_bit)
                break;
            ns ^= mod << (bit - mod_bit);
        }
        r0 = r1;
        s0 = s1;
        r1 = nr;
        s1 = ns;
    }
    return r0 == 1 ? s0 : 0;
}

template <typename F>
u128 modulus_of()
{
    constexpr unsigned bits = F::bits;
    if constexpr (bits == 1)
        return 0x3;
    else if constexpr (bits == 2)
        return 0x7;
    else if constexpr (bits == 4)
        return 0x13;
    else if constexpr (bits == 8)
        return 0x11b;
    else if constexpr (bits == 16)
        return 0x1002d;
    else if constexpr (bits == 32)
        return 0x190200001ull;
    else
        return (u128{1} << 64) | (u128{1} << 63)
            | (u128{1} << 62) | (u128{1} << 53) | 1;
}

template <typename F>
F mask_elem(u128 v)
{
    using word = typename F::integral_type;
    if constexpr (F::bits >= sizeof(word) * 8u)
        return F{static_cast<word>(v)};
    else
    {
        const word m = static_cast<word>((word{1} << F::bits) - word{1});
        return F{static_cast<word>(static_cast<word>(v) & m)};
    }
}

std::uint64_t rng_state = 0x123456789abcdefull;

std::uint64_t next_rng()
{
    rng_state = rng_state * 6364136223846793005ull + 1ull;
    return rng_state;
}

template <typename F>
F random_elem()
{
    return mask_elem<F>(next_rng());
}

template <typename Key0, typename Key1, typename In, typename F>
F open_at(const Key0 & k0, const Key1 & k1, In x)
{
    const auto y0 = *dpf::eval_point(k0, x);
    const auto y1 = *dpf::eval_point(k1, x);
    const F fwd = dpf::reconstruct(y0, y1);
    const F rev = dpf::reconstruct(y1, y0);
    EXPECT_EQ(fwd, rev);
    return fwd;
}

template <typename F>
void expect_full_domain(F beta)
{
    using In = std::uint8_t;
    for (int alpha_i : {0, 1, 31, 32, 127, 128, 254, 255})
    {
        const In alpha = static_cast<In>(alpha_i);
        auto [k0, k1] = dpf::make_dpf(alpha, beta);
        for (int x = 0; x < 256; ++x)
        {
            const In q = static_cast<In>(x);
            const F got = open_at<decltype(k0), decltype(k1), In, F>(k0, k1, q);
            EXPECT_EQ(got, q == alpha ? beta : F{}) << +q << " alpha " << +alpha;
        }
    }
}

template <typename F, typename It0, typename It1>
F open_written(It0 it0, It1 it1)
{
    using V0 = typename std::iterator_traits<It0>::value_type;
    using V1 = typename std::iterator_traits<It1>::value_type;
    V0 a = *it0;
    V1 b = *it1;
    if constexpr (dpf::is_secret_share_v<V0>)
        return dpf::reconstruct(a, b);
    else
    {
        const auto lane = [](auto v) {
            return F{static_cast<typename F::integral_type>(static_cast<unsigned>(v))};
        };
        return lane(a) + lane(b);
    }
}

template <typename F>
void expect_interval_and_sequence(F beta)
{
    using In = std::uint8_t;
    const In alpha = 0x2a;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    auto [b0, i0] = dpf::eval_interval(k0, In{0}, In{255});
    auto [b1, i1] = dpf::eval_interval(k1, In{0}, In{255});
    auto p0 = std::begin(i0);
    auto p1 = std::begin(i1);
    for (int x = 0; x < 256; ++x, ++p0, ++p1)
    {
        EXPECT_EQ(open_written<F>(p0, p1),
            static_cast<In>(x) == alpha ? beta : F{}) << x;
    }

    const In pts[] = {0, 1, 41, 42, 43, 127, 128, 255};
    auto [s0, is0] = dpf::eval_sequence(k0, std::begin(pts), std::end(pts));
    auto [s1, is1] = dpf::eval_sequence(k1, std::begin(pts), std::end(pts));
    auto q0 = std::begin(is0);
    auto q1 = std::begin(is1);
    for (In q : pts)
    {
        EXPECT_EQ(open_written<F>(q0, q1),
            q == alpha ? beta : F{}) << +q;
        ++q0;
        ++q1;
    }
}

template <typename F>
void expect_comparison(F beta)
{
    using In = std::uint8_t;
    const In alpha = 10;
    auto [lt0, lt1] = dpf::make_dpf(alpha, dpf::lt(beta));
    auto [le0, le1] = dpf::make_dpf(alpha, dpf::leq(beta));
    for (int x = 0; x < 256; ++x)
    {
        const In q = static_cast<In>(x);
        const auto a0 = dpf::eval_point<F>(dpf::cmp, lt0, q);
        const auto a1 = dpf::eval_point<F>(dpf::cmp, lt1, q);
        const auto b0 = dpf::eval_point<F>(dpf::cmp, le0, q);
        const auto b1 = dpf::eval_point<F>(dpf::cmp, le1, q);
        EXPECT_EQ(dpf::reconstruct(a0, a1), x < 10 ? beta : F{}) << x;
        EXPECT_EQ(dpf::reconstruct(b0, b1), x <= 10 ? beta : F{}) << x;
    }
}

template <typename Node, typename F>
void expect_scale(Node node, F k)
{
    const auto scaled = dpf::multiply_leaf(node, k);
    unsigned char src[sizeof(Node)];
    unsigned char got[sizeof(Node)];
    std::memcpy(src, &node, sizeof(src));
    std::memcpy(got, &scaled, sizeof(got));
    if constexpr (F::bits >= 8)
    {
        constexpr std::size_t lanes = sizeof(Node) / sizeof(typename F::integral_type);
        for (std::size_t i = 0; i < lanes; ++i)
        {
            typename F::integral_type lane{};
            std::memcpy(&lane, src + i * sizeof(lane), sizeof(lane));
            typename F::integral_type out{};
            std::memcpy(&out, got + i * sizeof(out), sizeof(out));
            EXPECT_EQ(F{out}, F{lane} * k);
        }
    }
    else
    {
        constexpr unsigned w = F::bits;
        constexpr unsigned mask = (1u << w) - 1u;
        constexpr unsigned per = 8u / w;
        for (std::size_t i = 0; i < sizeof(Node); ++i)
        {
            unsigned expect = 0;
            for (unsigned lane = 0; lane < per; ++lane)
            {
                const auto a = F{(src[i] >> (lane * w)) & mask};
                expect |= static_cast<unsigned>((a * k).raw()) << (lane * w);
            }
            EXPECT_EQ(got[i], static_cast<unsigned char>(expect));
        }
    }
}

template <typename F>
void expect_algebra()
{
    const auto mod = modulus_of<F>();
    const F one{1};
    EXPECT_EQ(one + one, F{});
    EXPECT_EQ(-one, one);
    EXPECT_EQ(F{-7}, F{7});
    EXPECT_EQ(one * one, one);
    EXPECT_EQ(F{} * random_elem<F>(), F{});

    const unsigned lim = F::bits <= 8 ? (1u << F::bits) : 0u;
    if (lim != 0)
    {
        for (unsigned a = 1; a < lim; ++a)
        {
            const auto inv = poly_inv(a, mod);
            ASSERT_NE(inv, 0u) << a;
            EXPECT_EQ(F{static_cast<typename F::integral_type>(a)}
                * mask_elem<F>(inv), one) << a;
        }
        for (unsigned a = 0; a < lim; ++a)
        {
            for (unsigned b = 0; b < lim; ++b)
            {
                const F prod = F{static_cast<typename F::integral_type>(a)}
                    * F{static_cast<typename F::integral_type>(b)};
                if (a != 0 && b != 0)
                    EXPECT_NE(prod, F{}) << a << " " << b;
            }
        }
    }
    else
    {
        for (int n = 0; n < 48; ++n)
        {
            const F a = random_elem<F>();
            if (a == F{})
                continue;
            const auto inv = poly_inv(a.raw(), mod);
            ASSERT_NE(inv, 0u);
            EXPECT_EQ(a * mask_elem<F>(inv), one);
        }
    }

    for (int n = 0; n < 32; ++n)
    {
        const F a = random_elem<F>();
        const F b = random_elem<F>();
        const F c = random_elem<F>();
        EXPECT_EQ((a + b) * c, a * c + b * c);
        EXPECT_EQ((a * b) * c, a * (b * c));
        EXPECT_EQ(a + a, F{});
    }

    unsigned char poison[sizeof(F)];
    std::memset(poison, 0xff, sizeof(poison));
    F poisoned{};
    std::memcpy(&poisoned, poison, sizeof(poisoned));
    EXPECT_EQ(poisoned.raw(), mask_elem<F>(~u128{0}).raw());
    EXPECT_EQ(poisoned + F{}, mask_elem<F>(~u128{0}));

    unsigned char lo[16]{};
    unsigned char hi[16]{};
    lo[0] = 0x15;
    hi[0] = 0x15;
    if (sizeof(typename F::integral_type) < 16)
        hi[sizeof(typename F::integral_type)] = 0x5a;
    EXPECT_EQ(F::from_seed(lo, sizeof(lo)), F::from_seed(hi, sizeof(hi)));
}

template <typename F>
void expect_leaf_ops()
{
    alignas(32) unsigned char bytes[32];
    for (int i = 0; i < 32; ++i)
        bytes[i] = static_cast<unsigned char>(0xA5 ^ i);
    simde__m128i n128;
    simde__m256i n256;
    std::memcpy(&n128, bytes, sizeof(n128));
    std::memcpy(&n256, bytes, sizeof(n256));
    const F k = F::bits == 1 ? F{1} : F{2};
    expect_scale(n128, k);
    expect_scale(n128, F{});
    expect_scale(n128, F{1});
    expect_scale(n256, k);

    const auto sum = dpf::add_leaf<F>(n128, n128);
    unsigned char z[sizeof(n128)];
    std::memcpy(z, &sum, sizeof(z));
    for (unsigned char b : z)
        EXPECT_EQ(b, 0);
    unsigned char d[sizeof(n128)];
    const auto sub = dpf::subtract_leaf<F>(n128, n128);
    std::memcpy(d, &sub, sizeof(d));
    for (unsigned char b : d)
        EXPECT_EQ(b, 0);
}

template <typename F>
void expect_shares_and_prefix(F beta)
{
    const auto add = dpf::make_additive_shares(beta);
    EXPECT_EQ(dpf::reconstruct(add.first, add.second), beta);
    EXPECT_EQ(dpf::reconstruct(add.second, add.first), beta);
    const auto sub = dpf::make_subtractive_shares(beta);
    EXPECT_EQ(dpf::reconstruct(sub.first, sub.second), beta);
    EXPECT_EQ(dpf::reconstruct(sub.second, sub.first), beta);

    auto drawn = dpf::additively_share(beta);
    EXPECT_EQ(dpf::reconstruct(drawn.first, drawn.second), beta);
}

template <typename F>
void expect_verifiable(F beta)
{
    using In = std::uint8_t;
    const In alpha = 0x11;
    auto [k0, k1] = dpf::make_dpf(alpha, beta, dpf::verifiable{});
    for (int x = 0; x < 256; ++x)
    {
        const In q = static_cast<In>(x);
        dpf::proof_token a{};
        dpf::proof_token b{};
        const auto y0 = *dpf::eval_point(k0, q, dpf::prove(a));
        const auto y1 = *dpf::eval_point(k1, q, dpf::prove(b));
        EXPECT_TRUE(dpf::verify(a, b)) << +q;
        EXPECT_EQ(dpf::reconstruct(y0, y1), q == alpha ? beta : F{});
    }
    auto & leaves = const_cast<std::decay_t<decltype(k0.leaves())> &>(k0.leaves());
    auto * raw = reinterpret_cast<unsigned char *>(&std::get<0>(leaves).get());
    raw[0] = static_cast<unsigned char>(raw[0] ^ 0x1u);
    dpf::proof_token tampered{};
    dpf::proof_token honest{};
    (void)*dpf::eval_point(k0, alpha, dpf::prove(tampered));
    (void)*dpf::eval_point(k1, alpha, dpf::prove(honest));
    EXPECT_FALSE(dpf::verify(tampered, honest));
}

template <typename F>
void run_all(F beta)
{
    expect_algebra<F>();
    expect_full_domain(F{});
    expect_full_domain(F{1});
    expect_full_domain(beta);
    expect_interval_and_sequence(beta);
    expect_comparison(beta);
    expect_leaf_ops<F>();
    expect_shares_and_prefix(beta);
}

}  // namespace

TEST(Gf2Adversarial, AlgebraDomainAndLeaves)
{
    run_all(dpf::gf2{1});
    run_all(dpf::gf22{3});
    run_all(dpf::gf24{0xa});
    run_all(dpf::gf28{0x1b});
    run_all(dpf::gf216{0x2d});
    run_all(dpf::gf232{0x90200001u});
    run_all(dpf::gf264{0x11});
}

TEST(Gf2Adversarial, ProofsRejectAFlippedLeaf)
{
    expect_verifiable(dpf::gf2{1});
    expect_verifiable(dpf::gf24{0xf});
    expect_verifiable(dpf::gf28{0xff});
    expect_verifiable(dpf::gf264{~std::uint64_t{0}});
}

TEST(Gf2Adversarial, PrefixLongerThanTheLane)
{
    using In = std::uint8_t;
    const In alpha = 0x2a;
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::at<8>(dpf::gf28{0x1b}), dpf::gf28{0x5a});
    const auto p0 = *dpf::eval_point<0>(k0, alpha);
    const auto p1 = *dpf::eval_point<0>(k1, alpha);
    const auto l0 = *dpf::eval_point<1>(k0, alpha);
    const auto l1 = *dpf::eval_point<1>(k1, alpha);
    EXPECT_EQ(dpf::reconstruct(p0, p1), dpf::gf28{0x1b});
    EXPECT_EQ(dpf::reconstruct(l0, l1), dpf::gf28{0x5a});
    const auto off0 = *dpf::eval_point<1>(k0, In{0});
    const auto off1 = *dpf::eval_point<1>(k1, In{0});
    EXPECT_EQ(dpf::reconstruct(off0, off1), dpf::gf28{});
}

TEST(Gf2Adversarial, MultiOutputDoesNotBleed)
{
    using In = std::uint8_t;
    const In alpha = 0x2a;
    const dpf::gf28 lane{0x1b};
    const std::uint8_t wide = 9;
    auto [k0, k1] = dpf::make_dpf(alpha, lane, wide);
    for (int x = 0; x < 64; ++x)
    {
        const In q = static_cast<In>(x);
        const auto a0 = *dpf::eval_point<0>(k0, q);
        const auto a1 = *dpf::eval_point<0>(k1, q);
        const auto b0 = *dpf::eval_point<1>(k0, q);
        const auto b1 = *dpf::eval_point<1>(k1, q);
        EXPECT_EQ(dpf::reconstruct(a0, a1), q == alpha ? lane : dpf::gf28{});
        EXPECT_EQ(dpf::reconstruct(b0, b1), q == alpha ? wide : std::uint8_t{0});
    }
}

TEST(Gf2Adversarial, NakedLeafMasksHighBits)
{
    const auto leaf = dpf::make_naked_leaf<simde__m128i>(std::uint8_t{3}, dpf::gf24{0xF5});
    EXPECT_EQ((dpf::extract_leaf<simde__m128i, dpf::gf24>(leaf, std::uint8_t{3})),
        dpf::gf24{0x5});
    EXPECT_EQ((dpf::extract_leaf<simde__m128i, dpf::gf24>(leaf, std::uint8_t{2})),
        dpf::gf24{});
    EXPECT_EQ((dpf::extract_leaf<simde__m128i, dpf::gf24>(leaf, std::uint8_t{4})),
        dpf::gf24{});
    EXPECT_EQ((dpf::extract_leaf<simde__m128i, dpf::gf24>(leaf, std::uint8_t{31})),
        dpf::gf24{});
}

TEST(Gf2Adversarial, UnassignedWildcardThrows)
{
    auto [w0, w1] = dpf::make_dpf(std::uint8_t{3}, dpf::wildcard_value<dpf::gf28>{});
    EXPECT_THROW((void)*dpf::eval_point(w0, std::uint8_t{3}), std::runtime_error);
    EXPECT_THROW((void)*dpf::eval_point(w1, std::uint8_t{3}), std::runtime_error);
    auto [p0, p1] = dpf::make_dpf(std::uint8_t{1}, dpf::wildcard_value<dpf::gf2>{});
    EXPECT_THROW((void)*dpf::eval_point(p0, std::uint8_t{1}), std::runtime_error);
    (void)p1;
}

TEST(Gf2Adversarial, FromSeedDropsBytesPastTheWord)
{
    unsigned char narrow[16]{};
    unsigned char wide[16];
    std::memset(wide, 0xa5, sizeof(wide));
    narrow[0] = 0x15;
    wide[0] = 0x15;
    EXPECT_EQ(dpf::gf24::from_seed(narrow, 16), dpf::gf24{0x5});
    EXPECT_EQ(dpf::gf24::from_seed(wide, 16), dpf::gf24{0x5});
    narrow[0] = 0xab;
    wide[0] = 0xab;
    EXPECT_EQ(dpf::gf28::from_seed(narrow, 1), dpf::gf28::from_seed(wide, 16));
}