#include <gtest/gtest.h>

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

#include "dpf.hpp"

using Interior = dpf::prg::aes128;
using Exterior = dpf::prg::aes128;

TEST(Verifiable, HonestPointAccepts)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const std::uint64_t beta = 7;
    auto [k0, k1] = dpf::make_dpf<Interior, Exterior>(alpha, beta, dpf::verifiable{});
    EXPECT_TRUE(decltype(k0)::is_verifiable);
    EXPECT_FALSE(decltype(k0)::is_multilevel);

    dpf::proof_token pi0{}, pi1{};
    const auto y0 = *dpf::eval_point(k0, alpha, dpf::prove(pi0));
    const auto y1 = *dpf::eval_point(k1, alpha, dpf::prove(pi1));
    EXPECT_EQ(dpf::reconstruct(y0, y1), beta);
    EXPECT_TRUE(dpf::verify(pi0, pi1));

    dpf::proof_token q0{}, q1{};
    const Input other = static_cast<Input>(alpha ^ 1);
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(k0, other, dpf::prove(q0)),
                  *dpf::eval_point(k1, other, dpf::prove(q1))),
        0);
    EXPECT_TRUE(dpf::verify(q0, q1));
}

TEST(Verifiable, TamperedCwRejects)
{
    using Input = std::uint8_t;
    auto [k0, k1] = dpf::make_dpf<Interior, Exterior>(Input{3},
        std::uint64_t{1}, dpf::verifiable{});

    // Flip one bit of a public correction word on party 0's view of the
    // shared CW array by rebuilding an otherwise-identical key is hard;
    // instead flip cs after the fact via const_cast of the seed storage.
    auto & cs = const_cast<dpf::cs_block &>(k0.correction_seeds()[0]);
    cs[0] = simde_mm_xor_si128(cs[0], simde_mm_set1_epi8(1));

    dpf::proof_token pi0{}, pi1{};
    (void)*dpf::eval_point(k0, Input{3}, dpf::prove(pi0));
    (void)*dpf::eval_point(k1, Input{3}, dpf::prove(pi1));
    EXPECT_FALSE(dpf::verify(pi0, pi1));
}

TEST(Verifiable, BatchVerify)
{
    using Input = std::uint8_t;
    auto [k0, k1] = dpf::make_dpf<Interior, Exterior>(Input{1},
        std::uint64_t{9}, dpf::verifiable{});

    std::vector<dpf::proof_token> left, right;
    for (Input x = 0; x < 8; ++x)
    {
        dpf::proof_token a{}, b{};
        (void)*dpf::eval_point(k0, x, dpf::prove(a));
        (void)*dpf::eval_point(k1, x, dpf::prove(b));
        left.push_back(a);
        right.push_back(b);
    }
    EXPECT_TRUE(dpf::verify_batch(left, right));
    left[2][0] = simde_mm_xor_si128(left[2][0], simde_mm_set1_epi8(0xff));
    EXPECT_FALSE(dpf::verify_batch(left, right));
    right.pop_back();
    EXPECT_FALSE(dpf::verify_batch(left, right));
}

TEST(Verifiable, HalfTreeXorPayload)
{
    using Input = std::uint16_t;
    using Ht = dpf::prg::aes128_ccr;
    const Input alpha = 0x0101;
    auto [k0, k1] = dpf::make_dpf<Ht, Ht>(alpha,
        dpf::xor_wrapper<std::uint64_t>{0xdeadbeefull}, dpf::verifiable{});
    EXPECT_TRUE(decltype(k0)::tree::is_half_tree);

    dpf::proof_token pi0{}, pi1{};
    const auto y0 = *dpf::eval_point(k0, alpha, dpf::prove(pi0));
    const auto y1 = *dpf::eval_point(k1, alpha, dpf::prove(pi1));
    EXPECT_EQ(dpf::reconstruct(y0, y1), dpf::xor_wrapper<std::uint64_t>{0xdeadbeefull});
    EXPECT_TRUE(dpf::verify(pi0, pi1));
}

TEST(Verifiable, SamePublicPart)
{
    using Input = std::uint8_t;
    auto [k0, k1] = dpf::make_dpf<Interior, Exterior>(Input{5},
        std::uint64_t{2}, dpf::verifiable{});
    EXPECT_TRUE(dpf::same_public_part(k0, k1));

    auto & cw = const_cast<typename std::decay_t<decltype(k0)>::interior_node &>(
        k0.correction_words()[0]);
    cw = simde_mm_xor_si128(cw, simde_mm_set1_epi8(1));
    EXPECT_FALSE(dpf::same_public_part(k0, k1));
}

TEST(Verifiable, DefaultKeyUnchangedLayout)
{
    using Input = std::uint8_t;
    auto [a0, a1] = dpf::make_dpf<Interior, Exterior>(Input{1}, std::uint64_t{3});
    auto [b0, b1] = dpf::make_dpf<Interior, Exterior>(Input{1}, std::uint64_t{3},
        dpf::verifiable{});
    EXPECT_FALSE(decltype(a0)::is_verifiable);
    EXPECT_TRUE(decltype(b0)::is_verifiable);
    EXPECT_EQ(sizeof(a0.correction_words()), sizeof(b0.correction_words()));
    EXPECT_EQ(std::tuple_size_v<typename decltype(a0)::correction_seeds_array>, 0u);
    EXPECT_GT(std::tuple_size_v<typename decltype(b0)::correction_seeds_array>, 0u);
}

TEST(Extractable, Fp61ReconstructAndSketch)
{
    using Input = std::uint8_t;
    const Input alpha = 0x11;
    const dpf::fp61 beta{42};
    auto [k0, k1] = dpf::make_dpf<Interior, Exterior>(alpha, beta,
        dpf::extractable{}, dpf::verifiable{});
    EXPECT_TRUE(decltype(k0)::is_extractable);
    EXPECT_TRUE(decltype(k0)::is_verifiable);
    EXPECT_TRUE(dpf::same_public_part(k0, k1));

    const auto y0 = *dpf::eval_point(k0, alpha);
    const auto y1 = *dpf::eval_point(k1, alpha);
    EXPECT_EQ(dpf::reconstruct(y0, y1), beta);

    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();
    }
    auto sk0 = dpf::sketch_fold(s0, r);
    auto sk1 = dpf::sketch_fold(s1, r);
    EXPECT_TRUE(dpf::sketch_verify(sk0, sk1));

    // Two hot points: forge by XORing a second beta into another share.
    s0[0] = s0[0] + beta;
    sk0 = dpf::sketch_fold(s0, r);
    sk1 = dpf::sketch_fold(s1, r);
    EXPECT_FALSE(dpf::sketch_verify(sk0, sk1));
}

TEST(Extractable, IncrementalPrefix)
{
    using Input = std::uint16_t;
    const Input alpha = 0x00ab;
    auto [k0, k1] = dpf::make_dpf<Interior, Exterior>(alpha,
        dpf::at<8>(dpf::fp61{1}), dpf::extractable{});
    EXPECT_TRUE(decltype(k0)::is_extractable);
    EXPECT_TRUE(decltype(k0)::is_multilevel);

    const auto p0 = *dpf::eval_point(dpf::out<0>, k0, alpha);
    const auto p1 = *dpf::eval_point(dpf::out<0>, k1, alpha);
    EXPECT_EQ(dpf::reconstruct(p0, p1), dpf::fp61{1});
}

TEST(Verifiable, IntervalProve)
{
    using Input = std::uint8_t;
    auto [k0, k1] = dpf::make_dpf<Interior, Exterior>(Input{0x20},
        std::uint64_t{1}, dpf::verifiable{});
    dpf::proof_token a{}, b{};
    dpf::prove_interval(k0, Input{0x1c}, Input{0x24}, dpf::prove(a));
    dpf::prove_interval(k1, Input{0x1c}, Input{0x24}, dpf::prove(b));
    EXPECT_TRUE(dpf::verify(a, b));
}

TEST(Verifiable, DoernerShelatProve)
{
    using Input = std::uint8_t;
    const Input alpha = 0x44;
    const std::uint64_t beta = 5;
    Input x0 = 0x12;
    Input x1 = static_cast<Input>(alpha ^ x0);
    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;
        }
        std::uint8_t bit() { return static_cast<std::uint8_t>(n++ & 1u); }
    };
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> rng{
        dpf::uniform_sample<simde__m128i>, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto [s0, s1] = dpf::make_dpf_doerner_shelat<Interior, Exterior>(
        x0, x1, rng, beta, dpf::verifiable{});

    EXPECT_TRUE(decltype(s0)::is_verifiable);
    dpf::proof_token a{}, b{};
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(s0, alpha, dpf::prove(a)),
                  *dpf::eval_point(s1, alpha, dpf::prove(b))),
        beta);
    EXPECT_TRUE(dpf::verify(a, b));
}