#include <gtest/gtest.h>

#include <cstdint>
#include <vector>

#include "dpf.hpp"

namespace
{

using input_t = std::uint8_t;               // 256-point domain
constexpr std::size_t kDomain = 256;

}  // namespace

// eval_full_add_into(buf, key, fold): the fold sees every written share once,
// in the same pass that adds it into the buffer (Express / Prio audit hook).
TEST(CallerFold, FullAddIntoFoldSeesEveryShare)
{
    const input_t alpha = 42;
    const std::uint64_t beta = 0xdeadbeefull;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);

    std::vector<std::uint64_t> buf0(kDomain, 0), buf1(kDomain, 0);
    std::uint64_t fold0 = 0, fold1 = 0;
    std::size_t calls0 = 0, calls1 = 0;

    dpf::eval_full_add_into(buf0, k0,
        [&](std::size_t, std::uint64_t g) { fold0 += g; ++calls0; });
    dpf::eval_full_add_into(buf1, k1,
        [&](std::size_t, std::uint64_t g) { fold1 += g; ++calls1; });

    EXPECT_EQ(calls0, kDomain);
    EXPECT_EQ(calls1, kDomain);

    // The buffer opens to a point function at alpha.
    for (std::size_t i = 0; i < kDomain; ++i)
    {
        const std::uint64_t got = buf0[i] - buf1[i];
        EXPECT_EQ(got, i == alpha ? beta : 0ull) << "i=" << i;
    }
    // The fold accumulated exactly the same shares: sum reconstructs to beta.
    EXPECT_EQ(static_cast<std::uint64_t>(fold0 - fold1), beta);
}

// eval_full_fold allocates its own buffer and folds each share.
TEST(CallerFold, FullFoldReconstructsAudit)
{
    const input_t alpha = 200;
    const std::uint64_t beta = 7;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);

    // Weighted fold: sum_i (i+1) * share_i. Reconstructs to (alpha+1)*beta.
    std::uint64_t w0 = 0, w1 = 0;
    (void)dpf::eval_full_fold(k0,
        [&](std::size_t i, std::uint64_t g) { w0 += (i + 1) * g; });
    (void)dpf::eval_full_fold(k1,
        [&](std::size_t i, std::uint64_t g) { w1 += (i + 1) * g; });

    EXPECT_EQ(static_cast<std::uint64_t>(w0 - w1),
        static_cast<std::uint64_t>((alpha + 1) * beta));
}

// eval_point_fold calls the fold exactly once with the written share.
TEST(CallerFold, PointFoldCalledOnce)
{
    const input_t alpha = 5;
    const std::uint64_t beta = 99;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);

    std::uint64_t seen0 = 0, seen1 = 0;
    int calls0 = 0, calls1 = 0;
    auto o0 = dpf::eval_point_fold(k0, alpha,
        [&](std::size_t, std::uint64_t g) { seen0 = g; ++calls0; });
    auto o1 = dpf::eval_point_fold(k1, alpha,
        [&](std::size_t, std::uint64_t g) { seen1 = g; ++calls1; });

    EXPECT_EQ(calls0, 1);
    EXPECT_EQ(calls1, 1);
    // What the fold saw equals what eval_point returned.
    EXPECT_EQ(seen0, static_cast<std::uint64_t>((*o0).raw()));
    EXPECT_EQ(seen1, static_cast<std::uint64_t>((*o1).raw()));
    EXPECT_EQ(static_cast<std::uint64_t>(seen0 - seen1), beta);
}

// The fold is generic over the leaf group: a blob<N> row folds by XOR.
TEST(CallerFold, FoldOverBlobLeaf)
{
    using blob_t = dpf::blob<24>;
    const input_t alpha = 17;
    blob_t beta{};
    for (std::size_t i = 0; i < blob_t::size; ++i)
        beta.bytes[i] = static_cast<unsigned char>(0x30 + i);

    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    std::vector<blob_t> buf0(kDomain), buf1(kDomain);

    blob_t x0{}, x1{};
    std::size_t calls = 0;
    dpf::eval_full_add_into(buf0, k0,
        [&](std::size_t, const blob_t & g) { x0 = x0 ^ g; ++calls; });
    dpf::eval_full_add_into(buf1, k1,
        [&](std::size_t, const blob_t & g) { x1 = x1 ^ g; });

    EXPECT_EQ(calls, kDomain);
    // XOR of all shares reconstructs to beta (only alpha is nonzero).
    blob_t recon{};
    for (std::size_t i = 0; i < blob_t::size; ++i)
        recon.bytes[i] = static_cast<unsigned char>(x0.bytes[i] ^ x1.bytes[i]);
    EXPECT_EQ(recon, beta);
}

// eval_sequence_xor: keyword PIR without materializing the bit vector.
TEST(CallerFold, SequenceXorMatchesRecordAtAlpha)
{
    const input_t alpha = 123;
    // Point function with a bit payload (1 at alpha).
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::bit::one);

    std::vector<std::uint64_t> records(kDomain);
    for (std::size_t i = 0; i < kDomain; ++i)
        records[i] = 0x1000ull * (i + 1) + 7;

    const std::uint64_t acc0 = dpf::eval_sequence_xor(k0, records);
    const std::uint64_t acc1 = dpf::eval_sequence_xor(k1, records);

    EXPECT_EQ(acc0 ^ acc1, records[alpha]);
}