#include <gtest/gtest.h>

#include <cstdint>
#include <vector>

#include "dpf.hpp"

namespace
{

std::vector<dpf::fp61> random_challenges(std::size_t n)
{
    std::vector<dpf::fp61> rs(n);
    for (auto & r : rs)
        r = dpf::uniform_sample<dpf::fp61>();
    return rs;
}

}  // namespace

TEST(PathSketch, HonestIdpfAccepts)
{
    using Input = std::uint8_t;
    const Input alpha = 0xA0;  // prefix 1010_0000 → length-3 prefix 101
    auto [k0, k1] = dpf::make_dpf(alpha,
        dpf::idpf(std::uint64_t{5}, std::uint64_t{5}, std::uint64_t{5}));

    const auto rs = random_challenges(dpf::path_sketch_challenge_count(3));
    EXPECT_TRUE((dpf::verify_idpf_path<3>(k0, k1, rs)));
}

TEST(PathSketch, HonestUnitWeightAccepts)
{
    using Input = std::uint8_t;
    auto [k0, k1] = dpf::make_dpf(Input{0b10100101},
        dpf::idpf(std::uint64_t{1}, std::uint64_t{1}, std::uint64_t{1},
            std::uint64_t{1}));
    const auto rs = random_challenges(dpf::path_sketch_challenge_count(4));
    EXPECT_TRUE((dpf::verify_idpf_path<4>(k0, k1, rs)));
}

TEST(PathSketch, LevelSketchRejectsTwoHot)
{
    // Manufacture a two-hot vector and feed sketch_fold directly.
    std::vector<dpf::fp61> y0 = {dpf::fp61{3}, dpf::fp61{0}, dpf::fp61{0},
        dpf::fp61{0}};
    std::vector<dpf::fp61> y1 = {dpf::fp61{0}, dpf::fp61{0}, dpf::fp61{2},
        dpf::fp61{0}};
    // Opened values: 3, 0, -2, 0 — two nonzeros.
    std::vector<dpf::fp61> r = {dpf::fp61{9}, dpf::fp61{4}, dpf::fp61{7},
        dpf::fp61{2}};
    const auto s0 = dpf::sketch_fold(y0, r);
    const auto s1 = dpf::sketch_fold(y1, r);
    EXPECT_FALSE(dpf::sketch_verify(s0, s1));
}

TEST(PathSketch, ParentGapRejectsBrokenRelation)
{
    using Input = std::uint8_t;
    auto [k0, k1] = dpf::make_dpf(Input{0xA0},
        dpf::idpf(std::uint64_t{5}, std::uint64_t{5}, std::uint64_t{5}));

    auto gamma = random_challenges(2);
    // Honest parent gap opens to 0.
    const auto g0 = dpf::sketch_path_parent(dpf::out<0, 1>, dpf::out<1, 2>,
        k0, gamma.data(), gamma.size());
    const auto g1 = dpf::sketch_path_parent(dpf::out<0, 1>, dpf::out<1, 2>,
        k1, gamma.data(), gamma.size());
    EXPECT_TRUE(dpf::sketch_path_verify_parent(g0, g1));

    // Bias party 0's gap — must reject.
    EXPECT_FALSE(dpf::sketch_path_verify_parent(g0 + dpf::fp61{1}, g1));
}

TEST(PathSketch, ChallengeCountFormula)
{
    // Depth 1: only level sketch at N=1 → 2 challenges.
    EXPECT_EQ(dpf::path_sketch_challenge_count(1), 2u);
    // Depth 2: level1 (2) + parent(2) + level2 (4) = 8.
    EXPECT_EQ(dpf::path_sketch_challenge_count(2), 8u);
    // Depth 3: 8 + parent@2 (4) + level3 (8) = 20.
    EXPECT_EQ(dpf::path_sketch_challenge_count(3), 20u);
}

TEST(PathSketch, MasticClientsVerify)
{
    auto [a0, a1] = dpf::make_dpf(std::uint8_t{0xA0},
        dpf::idpf(std::uint64_t{5}, std::uint64_t{5}, std::uint64_t{5}));
    auto [b0, b1] = dpf::make_dpf(std::uint8_t{0xB0},
        dpf::idpf(std::uint64_t{3}, std::uint64_t{3}, std::uint64_t{3}));
    const auto rs = random_challenges(dpf::path_sketch_challenge_count(3));
    EXPECT_TRUE((dpf::verify_idpf_path<3>(a0, a1, rs)));
    EXPECT_TRUE((dpf::verify_idpf_path<3>(b0, b1, rs)));
}