#include <gtest/gtest.h>

#include <cstdint>
#include <tuple>
#include <utility>
#include <vector>

#include "dpf.hpp"

namespace
{

template <typename K0, typename K1, typename InputT>
dpf::fp61 open_sdpf(const K0 & k0, const K1 & k1, InputT x,
    const dpf::sfss_ct & ct)
{
    return dpf::reconstruct(
        dpf::subtractive_share<dpf::fp61, 0>::from_raw(
            dpf::eval_sdpf(k0, x, ct)),
        dpf::subtractive_share<dpf::fp61, 1>::from_raw(
            dpf::eval_sdpf(k1, x, ct)));
}

}  // namespace

TEST(Sfss, PointBetaOneOnAndOff)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    auto [k0, k1, ke, st] = dpf::make_sdpf(alpha);
    EXPECT_EQ(ke.beta.raw(), 1u);

    const dpf::fp61 m{42};
    const auto ct = dpf::sfss_enc(st, ke, m);
    EXPECT_EQ(ct.j, 1u);
    EXPECT_EQ(st.ctr, 2u);

    EXPECT_EQ(open_sdpf(k0, k1, alpha, ct), m);
    EXPECT_EQ(open_sdpf(k0, k1, Input{0}, ct), dpf::fp61{0});
    EXPECT_EQ(open_sdpf(k0, k1, static_cast<Input>(alpha ^ 1), ct), dpf::fp61{0});
}

TEST(Sfss, WeightedBeta)
{
    using Input = std::uint16_t;
    const Input alpha = 1000;
    const dpf::fp61 beta{7};
    auto [k0, k1, ke, st] = dpf::make_sdpf(alpha, beta);
    EXPECT_EQ(ke.beta, beta);

    const dpf::fp61 m{11};
    const auto ct = dpf::sfss_enc(st, ke, m);
    EXPECT_EQ(open_sdpf(k0, k1, alpha, ct), beta * m);
    EXPECT_EQ(open_sdpf(k0, k1, Input{0}, ct), dpf::fp61{0});
}

TEST(Sfss, ManyStreamMessages)
{
    using Input = std::uint8_t;
    const Input alpha = 7;
    auto [k0, k1, ke, st] = dpf::make_sdpf(alpha, dpf::fp61{3});

    for (std::uint64_t i = 0; i < 20; ++i)
    {
        const dpf::fp61 m{i + 1};
        const auto ct = dpf::sfss_enc(st, ke, m);
        EXPECT_EQ(ct.j, i + 1);
        EXPECT_EQ(open_sdpf(k0, k1, alpha, ct), ke.beta * m) << i;
        EXPECT_EQ(open_sdpf(k0, k1, Input{8}, ct), dpf::fp61{0}) << i;
    }
}

TEST(Sfss, WindowTelescopes)
{
    using Input = std::uint8_t;
    const Input alpha = 3;
    auto [k0, k1, ke, st] = dpf::make_sdpf(alpha, dpf::fp61{5});

    const std::uint64_t j_lo = st.ctr;
    dpf::fp61 sum_m{};
    dpf::fp61 c_agg{};
    constexpr std::size_t window = 8;
    for (std::size_t i = 0; i < window; ++i)
    {
        const dpf::fp61 m{static_cast<std::uint64_t>(i + 2)};
        sum_m = sum_m + m;
        const auto ct = dpf::sfss_enc_window(st, ke, m);
        c_agg = c_agg + ct.c;
    }
    const std::uint64_t j_hi = j_lo + window - 1;

    const auto s0 = dpf::eval_sdpf_window(k0, alpha, c_agg, j_lo, j_hi);
    const auto s1 = dpf::eval_sdpf_window(k1, alpha, c_agg, j_lo, j_hi);
    EXPECT_EQ(dpf::reconstruct(
                  dpf::subtractive_share<dpf::fp61, 0>::from_raw(s0),
                  dpf::subtractive_share<dpf::fp61, 1>::from_raw(s1)),
        ke.beta * sum_m);

    const auto z0 = dpf::eval_sdpf_window(k0, Input{4}, c_agg, j_lo, j_hi);
    const auto z1 = dpf::eval_sdpf_window(k1, Input{4}, c_agg, j_lo, j_hi);
    EXPECT_EQ(dpf::reconstruct(
                  dpf::subtractive_share<dpf::fp61, 0>::from_raw(z0),
                  dpf::subtractive_share<dpf::fp61, 1>::from_raw(z1)),
        dpf::fp61{0});
}

TEST(Sfss, RejectZeroBeta)
{
    EXPECT_THROW((void)dpf::make_sdpf(std::uint8_t{1}, dpf::fp61{0}),
        std::invalid_argument);
}