/// @file vdpf_adversarial_test.cpp
/// @brief Brute-force and corruption checks for verifiable evaluation.
/// @details Covers bug classes seen while bringing the proofs up: subtractive
///          reconstruction is not commutative, a single untouched control bit
///          can hide a seed flip, and a proof convention must still reject a
///          flipped token. Small domains are checked at every point.
#include <gtest/gtest.h>

#include <array>
#include <cstdint>
#include <vector>

#include "dpf.hpp"

namespace
{

using Input = std::uint8_t;

struct Pad
{
    std::uint64_t n = 1;
    simde__m128i block()
    {
        auto v = simde_mm_set_epi64x(static_cast<long long>(n),
            static_cast<long long>(n * 9 + 3));
        n += 2;
        return v;
    }
    void fill(void * p, std::size_t nbytes)
    {
        auto * b = static_cast<unsigned char *>(p);
        for (std::size_t i = 0; i < nbytes; ++i)
            b[i] = static_cast<unsigned char>(n + i * 17);
        n += nbytes;
    }
    std::uint8_t bit() { return static_cast<std::uint8_t>(n++ & 1u); }
};

dpf::ds_randomness<simde__m128i (*)(), Pad> tape()
{
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    return {dpf::uniform_sample<simde__m128i>, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
}

template <typename Y0, typename Y1, typename Want>
void expect_ordered_reconstruct(const Y0 & y0, const Y1 & y1, const Want & want)
{
    EXPECT_EQ(dpf::reconstruct(y0, y1), want);
    // The typed overload accepts either party order. A raw subtraction does not.
    EXPECT_EQ(dpf::reconstruct(y1, y0), want);
    const auto swapped = static_cast<Want>(y1.raw() - y0.raw());
    if (y0.raw() != y1.raw())
        EXPECT_NE(swapped, want);
}

}  // namespace

TEST(VdpfAdversarial, PointFullDomainValuesAndProofs)
{
    const Input alpha = 0;
    const std::uint64_t beta = 9;
    auto [k0, k1] = dpf::make_dpf(alpha, beta, dpf::verifiable{});
    for (int x = 0; x < 256; ++x)
    {
        const Input q = static_cast<Input>(x);
        dpf::proof_token a{}, b{};
        const auto y0 = *dpf::eval_point(k0, q, dpf::prove(a));
        const auto y1 = *dpf::eval_point(k1, q, dpf::prove(b));
        const std::uint64_t want = q == alpha ? beta : 0;
        expect_ordered_reconstruct(y0, y1, want);
        EXPECT_TRUE(dpf::verify(a, b)) << int(q);
    }
}

TEST(VdpfAdversarial, HalfTreeFullDomainValuesAndProofs)
{
    using Ht = dpf::prg::aes128_ccr;
    const Input alpha = 255;
    const std::uint64_t beta = 0x1001;
    auto [k0, k1] = dpf::make_dpf<Ht, Ht>(alpha, beta, dpf::verifiable{});
    EXPECT_TRUE(decltype(k0)::tree::is_half_tree);
    for (int x = 0; x < 256; ++x)
    {
        const Input q = static_cast<Input>(x);
        dpf::proof_token a{}, b{};
        const auto y0 = *dpf::eval_point(k0, q, dpf::prove(a));
        const auto y1 = *dpf::eval_point(k1, q, dpf::prove(b));
        expect_ordered_reconstruct(y0, y1, q == alpha ? beta : 0ull);
        EXPECT_TRUE(dpf::verify(a, b)) << int(q);
    }
}

TEST(VdpfAdversarial, ComparisonAndBlockedFullDomain)
{
    const Input alpha = 0x40;
    auto native = dpf::make_dpf(alpha, dpf::lt(std::uint64_t{1}), dpf::verifiable{});
    auto blocked = dpf::make_dpf(alpha,
        dpf::block_width<4>(dpf::lt(std::uint64_t{1})), dpf::verifiable{});
    for (int x = 0; x < 256; ++x)
    {
        const Input q = static_cast<Input>(x);
        dpf::proof_token n0{}, n1{}, b0{}, b1{};
        const auto ny0 = dpf::eval_point(dpf::cmp, native.first, q, dpf::prove(n0));
        const auto ny1 = dpf::eval_point(dpf::cmp, native.second, q, dpf::prove(n1));
        const auto by0 = dpf::eval_point(dpf::cmp, blocked.first, q, dpf::prove(b0));
        const auto by1 = dpf::eval_point(dpf::cmp, blocked.second, q, dpf::prove(b1));
        const std::uint64_t want = q < alpha ? 1u : 0u;
        EXPECT_EQ(dpf::reconstruct(ny0, ny1) & native.first.cmp().mask, want) << int(q);
        EXPECT_EQ(dpf::reconstruct(by0, by1) & blocked.first.cmp().mask, want) << int(q);
        EXPECT_TRUE(dpf::verify(n0, n1)) << int(q);
        EXPECT_TRUE(dpf::verify(b0, b1)) << int(q);
    }
}

TEST(VdpfAdversarial, ExtractableFp61FullDomainSketch)
{
    const Input alpha = 0x7f;
    const dpf::fp61 beta{42};
    auto [k0, k1] = dpf::make_dpf(alpha, beta, dpf::extractable{}, dpf::verifiable{});
    std::array<dpf::fp61, 256> s0{}, s1{}, r{};
    for (int x = 0; x < 256; ++x)
    {
        const Input q = static_cast<Input>(x);
        dpf::proof_token a{}, b{};
        const auto y0 = *dpf::eval_point(k0, q, dpf::prove(a));
        const auto y1 = *dpf::eval_point(k1, q, dpf::prove(b));
        EXPECT_EQ(dpf::reconstruct(y0, y1), q == alpha ? beta : dpf::fp61{0}) << int(q);
        EXPECT_TRUE(dpf::verify(a, b)) << int(q);
        s0[static_cast<std::size_t>(x)] = y0.raw();
        s1[static_cast<std::size_t>(x)] = y1.raw();
        r[static_cast<std::size_t>(x)] = dpf::fp61{static_cast<std::uint64_t>(3 * x + 1)};
    }
    const auto honest0 = s0;
    EXPECT_TRUE(dpf::sketch_verify(dpf::sketch_fold(s0, r), dpf::sketch_fold(s1, r)));
    // A second hot point is weight 2. Changing only the magnitude of the
    // single hot point stays weight 1 and must still verify.
    s0[0] = s0[0] + beta;
    EXPECT_FALSE(dpf::sketch_verify(dpf::sketch_fold(s0, r), dpf::sketch_fold(s1, r)));

    auto r_bad = r;
    r_bad[alpha] = r_bad[alpha] + dpf::fp61{1};
    EXPECT_FALSE(dpf::sketch_verify(
        dpf::sketch_fold(honest0, r_bad), dpf::sketch_fold(s1, r)));
}

TEST(VdpfAdversarial, SeedAndWordCorruptionAreVisible)
{
    const Input alpha = 0x2a;
    auto [k0, k1] = dpf::make_dpf(alpha, std::uint64_t{5}, dpf::verifiable{});
    EXPECT_TRUE(dpf::same_public_part(k0, k1));

    for (auto & cs : const_cast<typename std::decay_t<decltype(k0)>::correction_seeds_array &>(
             k0.correction_seeds()))
        cs[0] = simde_mm_xor_si128(cs[0], simde_mm_set1_epi8(1));
    EXPECT_FALSE(dpf::same_public_part(k0, k1));
    int proof_fail = 0;
    for (int x = 0; x < 256; ++x)
    {
        const Input q = static_cast<Input>(x);
        dpf::proof_token a{}, b{};
        (void)*dpf::eval_point(k0, q, dpf::prove(a));
        (void)*dpf::eval_point(k1, q, dpf::prove(b));
        if (!dpf::verify(a, b))
            ++proof_fail;
    }
    EXPECT_GT(proof_fail, 0);

    auto [h0, h1] = dpf::make_dpf(alpha, std::uint64_t{5}, dpf::verifiable{});
    // A level-0 word is invisible when that party's root control bit is 0.
    // Flip every correction word so a later on-path level is corrupted.
    auto & words = const_cast<typename std::decay_t<decltype(h0)>::correction_words_array &>(
        h0.correction_words());
    for (auto & word : words)
        word = simde_mm_xor_si128(word, simde_mm_set1_epi8(0x5a));
    EXPECT_FALSE(dpf::same_public_part(h0, h1));
    int value_fail = 0;
    for (int x = 0; x < 256; ++x)
    {
        const Input q = static_cast<Input>(x);
        const auto got = dpf::reconstruct(*dpf::eval_point(h0, q), *dpf::eval_point(h1, q));
        const std::uint64_t want = q == alpha ? 5u : 0u;
        if (got != want)
            ++value_fail;
    }
    EXPECT_GT(value_fail, 0);
}

TEST(VdpfAdversarial, EmptySequenceProofVerifies)
{
    auto [k0, k1] = dpf::make_dpf(Input{1}, std::uint64_t{1}, dpf::verifiable{});
    const std::vector<Input> none;
    dpf::proof_token a{}, b{};
    dpf::prove_sequence(k0, none.begin(), none.end(), dpf::prove(a));
    dpf::prove_sequence(k1, none.begin(), none.end(), dpf::prove(b));
    EXPECT_TRUE(dpf::verify(a, b));
}

TEST(VdpfAdversarial, GenevalFullDomainProofAndPartyOrder)
{
    const Input alpha = 0x3c;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const std::uint64_t y = 0x7e;
    auto g = dpf::geneval_full(x0, x1, tape(), y);
    ASSERT_EQ(g.party0.size(), 256u);
    EXPECT_TRUE(dpf::verify(g.proof0, g.proof1));
    for (int q = 0; q < 256; ++q)
    {
        const auto want = static_cast<Input>(q) == alpha ? y : 0ull;
        EXPECT_EQ(static_cast<std::uint64_t>(g.party0[static_cast<std::size_t>(q)]
                       - g.party1[static_cast<std::size_t>(q)]),
            want) << q;
        if (g.party0[static_cast<std::size_t>(q)] != g.party1[static_cast<std::size_t>(q)])
        {
            EXPECT_NE(static_cast<std::uint64_t>(g.party1[static_cast<std::size_t>(q)]
                           - g.party0[static_cast<std::size_t>(q)]),
                want) << q;
        }
    }
    g.proof0[0] = simde_mm_xor_si128(g.proof0[0], simde_mm_set1_epi8(1));
    EXPECT_FALSE(dpf::verify(g.proof0, g.proof1));
}

TEST(VdpfAdversarial, MacDetectsValueAndTagCorruption)
{
    const auto key = dpf::sample_mac_key<dpf::fp61>();
    auto [a, b] = dpf::mac_share_value(dpf::fp61{20}, key);
    EXPECT_TRUE(dpf::mac_verify(a, b, key));
    const auto scaled = dpf::mac_scale(a, dpf::fp61{2});
    const auto scaled_b = dpf::mac_scale(b, dpf::fp61{2});
    EXPECT_TRUE(dpf::mac_verify(scaled, scaled_b, key));
    EXPECT_EQ((scaled.value + scaled_b.value).raw(), (dpf::fp61{20} * dpf::fp61{2}).raw());

    a.value = a.value + dpf::fp61{1};
    EXPECT_FALSE(dpf::mac_verify(a, b, key));
    a.value = a.value - dpf::fp61{1};
    a.tag = a.tag + dpf::fp61{1};
    EXPECT_FALSE(dpf::mac_verify(a, b, key));

    a.tag = a.tag - dpf::fp61{1};
    std::array<dpf::mac_share<dpf::fp61>, 1> left{a};
    std::array<dpf::mac_share<dpf::fp61>, 2> right{b, b};
    std::array<dpf::fp61, 1> coeffs{dpf::fp61{1}};
    EXPECT_FALSE(dpf::mac_verify_batch(left, right, coeffs, key));
}