/// @file opt_in_malicious_test.cpp
/// @brief Opt-in verifiable / extractable / MAC checks stay off by default.
#include <gtest/gtest.h>
#include <tuple>

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

#include "dpf.hpp"
#include "dpf/beaver.hpp"
#include "grotto/carry.hpp"

namespace
{

using Input = std::uint8_t;
using u64 = std::uint64_t;

struct Counter
{
    int draws = 0;
    u64 operator()()
    {
        ++draws;
        return 0x9e3779b97f4a7c15ull * static_cast<u64>(draws);
    }
};

}  // namespace

TEST(OptInProfile, SemiHonestKeyHasNoProofOrSketch)
{
    auto [k0, k1] = dpf::make_dpf_profile<dpf::semi_honest>(Input{0x2a}, u64{7});
    EXPECT_FALSE(decltype(k0)::is_verifiable);
    EXPECT_FALSE(decltype(k0)::is_extractable);
    EXPECT_FALSE(decltype(k1)::is_verifiable);
    EXPECT_FALSE(decltype(k1)::is_extractable);

    auto plain = dpf::make_dpf(Input{0x2a}, u64{7});
    EXPECT_FALSE(decltype(plain.first)::is_verifiable);
    EXPECT_FALSE(decltype(plain.first)::is_extractable);
}

TEST(OptInProfile, CheckedKeyOptsIn)
{
    auto [k0, k1] = dpf::make_dpf_profile<dpf::checked>(Input{0x2a}, dpf::fp61{7});
    EXPECT_TRUE(decltype(k0)::is_verifiable);
    EXPECT_TRUE(decltype(k0)::is_extractable);
}

TEST(OptInProfile, FlagsAreIndependentOfTheOutputMacBit)
{
    using ver_only = dpf::auth_profile<true, false, false>;
    using ext_only = dpf::auth_profile<false, true, false>;
    using mac_only = dpf::auth_profile<false, false, true>;
    static_assert(dpf::checked::output_mac);
    static_assert(!dpf::semi_honest::output_mac);
    static_assert(std::tuple_size_v<decltype(dpf::key_tags_tuple<dpf::checked>())> == 2);
    static_assert(std::tuple_size_v<decltype(dpf::key_tags_tuple<mac_only>())> == 0);
    static_assert(std::tuple_size_v<decltype(dpf::key_tags<ver_only>())> == 1);

    auto ver = dpf::make_dpf_profile<ver_only>(Input{1}, u64{4});
    EXPECT_TRUE(decltype(ver.first)::is_verifiable);
    EXPECT_FALSE(decltype(ver.first)::is_extractable);
    auto ext = dpf::make_dpf_profile<ext_only>(Input{1}, dpf::fp61{4});
    EXPECT_FALSE(decltype(ext.first)::is_verifiable);
    EXPECT_TRUE(decltype(ext.first)::is_extractable);
    auto mac = dpf::make_dpf_profile<mac_only>(Input{1}, u64{4});
    EXPECT_FALSE(decltype(mac.first)::is_verifiable);
    EXPECT_FALSE(decltype(mac.first)::is_extractable);
}

TEST(OptInProfile, CheckedUint8DomainProofSketchAndSeedFlip)
{
    const Input alpha = 0x2a;
    const dpf::fp61 beta{9};
    auto [k0, k1] = dpf::make_dpf_profile<dpf::checked>(alpha, beta);
    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));
        s0[static_cast<std::size_t>(x)] = y0.raw();
        s1[static_cast<std::size_t>(x)] = y1.raw();
        EXPECT_TRUE(dpf::verify(a, b)) << x;
        EXPECT_EQ(dpf::reconstruct(y0, y1), x == alpha ? beta : dpf::fp61{0});
        r[static_cast<std::size_t>(x)] = dpf::fp61{static_cast<std::uint64_t>(3 + x)};
    }
    EXPECT_TRUE(dpf::sketch_verify(dpf::sketch_fold(s0, r), dpf::sketch_fold(s1, r)));
    auto forged = s0;
    forged[0] = forged[0] + beta;
    EXPECT_FALSE(dpf::sketch_verify(dpf::sketch_fold(forged, r), dpf::sketch_fold(s1, r)));

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

TEST(OptInMac, DefaultSessionHasNoTagVectors)
{
    dpf::beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * y);
    Counter rng;
    s.sample(rng);
    auto tape = s.export_party(0);
    EXPECT_FALSE(tape.has_mac);
    EXPECT_TRUE(tape.lambda_tag.empty());
    EXPECT_TRUE(tape.dot_cross_tag.empty());
    (void)z;
}

TEST(OptInMac, HonestHornerOpeningsVerify)
{
    auto key = dpf::sample_mac_key<u64>();
    dpf::beavers::session<u64> s;
    s.set_mac_key(key);
    auto x = s.input();
    auto z = s.horner(x, {u64{1}, u64{2}, u64{3}});
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z), 1u + 2u * 5u + 3u * 25u);
    EXPECT_TRUE(s.verify_delta(x));
    EXPECT_TRUE(s.verify_all());
}

TEST(OptInMac, FlippedTagRejects)
{
    auto key = dpf::sample_mac_key<u64>();
    dpf::beavers::session<u64> s;
    s.set_mac_key(key);
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * y);
    Counter rng;
    s.sample(rng);
    s.bind(x, u64{3}, rng);
    s.bind(y, u64{5}, rng);
    s.evaluate();
    EXPECT_EQ(s.open(z), 15u);
    auto a = s.delta_auth(x);
    dpf::beavers::auth_opening<u64> o0{a.value.p0, a.tag.p0};
    dpf::beavers::auth_opening<u64> o1{a.value.p1, a.tag.p1};
    EXPECT_TRUE(dpf::beavers::verify_auth_opening(o0, o1, key));
    o0.tag ^= 1ull;
    EXPECT_FALSE(dpf::beavers::verify_auth_opening(o0, o1, key));
}

TEST(OptInExtractable, NoteSketchNoOpOnPlainKey)
{
    using key_t = decltype(dpf::make_dpf(Input{1}, u64{1}).first);
    dpf::sketch_share sk{};
    const std::array<u64, 1> ys{9};
    const std::array<dpf::fp61, 1> rs{dpf::fp61{2}};
    dpf::note_sketch<key_t>(sk, ys, rs);
    EXPECT_EQ(sk.z1, dpf::fp61{0});
    EXPECT_EQ(sk.z2, dpf::fp61{0});
    EXPECT_EQ(sk.z3, dpf::fp61{0});
}

TEST(OptInExtractable, HonestSketchAcceptsTamperRejects)
{
    const Input alpha = 0x11;
    const dpf::fp61 beta{42};
    auto [k0, k1] = dpf::make_dpf(alpha, beta, dpf::extractable{});
    EXPECT_TRUE(decltype(k0)::is_extractable);
    EXPECT_FALSE(decltype(k0)::is_verifiable);

    std::array<Input, 4> pts{0x10, 0x11, 0x12, 0x13};
    std::array<dpf::fp61, 4> r{
        dpf::fp61{3}, dpf::fp61{5}, dpf::fp61{7}, dpf::fp61{11}};
    std::array<dpf::fp61, 4> s0{}, s1{};
    for (std::size_t i = 0; i < pts.size(); ++i)
    {
        s0[i] = (*dpf::eval_point(k0, pts[i])).raw();
        s1[i] = (*dpf::eval_point(k1, pts[i])).raw();
    }
    dpf::sketch_share sk0{}, sk1{};
    dpf::note_sketch<decltype(k0)>(sk0, s0, r);
    dpf::note_sketch<decltype(k1)>(sk1, s1, r);
    EXPECT_TRUE(dpf::sketch_verify(sk0, sk1));
    s0[0] = s0[0] + beta;
    dpf::note_sketch<decltype(k0)>(sk0, s0, r);
    EXPECT_FALSE(dpf::sketch_verify(sk0, sk1));
}

TEST(OptInMemoProve, HonestIdempotentAndTamper)
{
    const Input alpha = 0x2a;
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::block_width<4>(dpf::lt(u64{1})), dpf::verifiable{});
    using KT0 = decltype(k0);
    using KT1 = decltype(k1);
    const Input from = 0x20;
    const Input to = 0x2f;
    const auto nbits = static_cast<std::size_t>(k0.cmp().nbits);
    using integral = typename KT0::integral_type;
    const auto a = static_cast<integral>(from);
    const auto b = static_cast<integral>(to);
    const auto excl = dpf::detail::incr::cmp_exclusive_end(b);
    const auto count = dpf::detail::incr::cmp_inclusive_count(a, b);
    constexpr std::size_t stop =
        KT0::cmp_depth == 0 ? KT0::depth : KT0::cmp_depth;
    const std::size_t levels = KT0::cmp_h;

    dpf::detail::incr::cmp_full_interval_memo<KT0, stop> memo0{count};
    dpf::detail::incr::cmp_full_interval_memo<KT1, stop> memo1{count};
    dpf::detail::incr::eval_cmp_interval_impl_interior(k0, a, excl, nbits, memo0,
        levels, nullptr);
    dpf::detail::incr::eval_cmp_interval_impl_interior(k1, a, excl, nbits, memo1,
        levels, nullptr);

    const integral lane = static_cast<integral>(alpha);
    dpf::proof_token pi0{}, pi1{};
    dpf::detail::vdpf::init_proof(pi0, k0);
    dpf::detail::vdpf::init_proof(pi1, k1);
    dpf::basic_path_memoizer<KT0> path0{};
    dpf::basic_path_memoizer<KT1> path1{};
    dpf::detail::blocked::eval_share_memo(k0, lane, a, excl, memo0, &pi0,
        &path0);
    dpf::detail::blocked::eval_share_memo(k1, lane, a, excl, memo1, &pi1,
        &path1);
    EXPECT_TRUE(dpf::verify(pi0, pi1));

    const auto snap0 = pi0;
    dpf::detail::blocked::eval_share_memo(k0, lane, a, excl, memo0, &pi0,
        &path0);
    EXPECT_TRUE(dpf::detail::vdpf::proof_equal(pi0, snap0));

    // CW tamper on party 0.
    auto k0_bad = k0;
    for (auto & cs : const_cast<typename KT0::correction_seeds_array &>(
             k0_bad.correction_seeds()))
    {
        cs[0] = simde_mm_xor_si128(cs[0], simde_mm_set1_epi8(0x5a));
    }
    dpf::detail::incr::cmp_full_interval_memo<KT0, stop> memo_bad{count};
    dpf::detail::incr::eval_cmp_interval_impl_interior(k0_bad, a, excl, nbits,
        memo_bad, levels, nullptr);
    dpf::proof_token q0{}, q1{};
    dpf::detail::vdpf::init_proof(q0, k0_bad);
    dpf::detail::vdpf::init_proof(q1, k1);
    dpf::basic_path_memoizer<KT0> pb0{};
    dpf::basic_path_memoizer<KT1> pb1{};
    dpf::detail::blocked::eval_share_memo(k0_bad, lane, a, excl, memo_bad,
        &q0, &pb0);
    dpf::detail::blocked::eval_share_memo(k1, lane, a, excl, memo1, &q1,
        &pb1);
    EXPECT_FALSE(dpf::verify(q0, q1));
}

TEST(OptInMemoProve, BfsIntervalProveUnchanged)
{
    const Input alpha = 0x2a;
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::block_width<4>(dpf::lt(u64{1})), dpf::verifiable{});
    const Input from = 0x00;
    const Input to = 0x3f;
    dpf::proof_token a{}, b{};
    dpf::prove_cmp_interval(k0, from, to, dpf::prove(a));
    dpf::prove_cmp_interval(k1, from, to, dpf::prove(b));
    EXPECT_TRUE(dpf::verify(a, b));

    // Re-proving the same interval yields the same BFS transcript.
    dpf::proof_token a2{}, b2{};
    dpf::prove_cmp_interval(k0, from, to, dpf::prove(a2));
    dpf::prove_cmp_interval(k1, from, to, dpf::prove(b2));
    EXPECT_TRUE(dpf::detail::vdpf::proof_equal(a, a2));
    EXPECT_TRUE(dpf::detail::vdpf::proof_equal(b, b2));

    // Interval eval with prove still verifies (memo walk does not fold again).
    auto buf0 = dpf::eval_interval(dpf::cmp, k0, from, to, dpf::prove(a2));
    auto buf1 = dpf::eval_interval(dpf::cmp, k1, from, to, dpf::prove(b2));
    EXPECT_TRUE(dpf::verify(a2, b2));
    EXPECT_EQ(buf0.size(), buf1.size());
}

TEST(OptInCarry, VerifiableAndMacRoundTrip)
{
    grotto::carry_auth auth{};
    auth.verifiable = true;
    auth.output_mac = true;
    auto keys = grotto::make_carry_in_keys(8, 3, auth);
    ASSERT_TRUE(keys.has_mac);
    ASSERT_TRUE(keys.low_lt_v.has_value());
    const std::uint64_t opened = 0x3cu;
    dpf::proof_token t0[4]{}, t1[4]{};
    const auto n0 = grotto::prove_carry_keys(keys, 0, opened, t0, 4);
    const auto n1 = grotto::prove_carry_keys(keys, 1, opened, t1, 4);
    EXPECT_EQ(n0, n1);
    EXPECT_TRUE(dpf::verify_batch(
        std::vector<dpf::proof_token>(t0, t0 + n0),
        std::vector<dpf::proof_token>(t1, t1 + n1)));
    const auto y0 = grotto::eval_carry_in(keys, 0, opened);
    const auto y1 = grotto::eval_carry_in(keys, 1, opened);
    auto [m0, m1] = grotto::mac_carry_result(keys, y0.value, y1.value);
    EXPECT_TRUE(dpf::mac_verify(m0, m1, keys.mac, t0[0], t1[0]));
}