#include <gtest/gtest.h>

#include <cstdint>
#include <cstring>
#include <type_traits>
#include <utility>

#include "dpf.hpp"

namespace
{

TEST(SecretShare, LayoutMatchesValueType)
{
    using T = std::uint64_t;
    EXPECT_EQ(sizeof(dpf::additive_share<T, 0>), sizeof(T));
    EXPECT_EQ(sizeof(dpf::subtractive_share<T, 1>), sizeof(T));
    EXPECT_TRUE((std::is_trivially_copyable_v<dpf::additive_share<T, 0>>));
    EXPECT_TRUE((std::is_standard_layout_v<dpf::subtractive_share<T, 1>>));
}

TEST(SecretShare, PlaintextSplitOpens)
{
    const std::uint32_t secret = 0xdeadbeefu;
    auto [a0, a1] = dpf::make_additive_shares(secret);
    EXPECT_EQ(dpf::reconstruct(a0, a1), secret);
    EXPECT_EQ(a1.raw(), 0u);

    auto [s0, s1] = dpf::make_subtractive_shares(secret);
    EXPECT_EQ(dpf::reconstruct(s0, s1), secret);
    EXPECT_EQ(s1.raw(), 0u);
}

TEST(SecretShare, SameSchemeLinearCombo)
{
    auto [a0, a1] = dpf::make_additive_shares(std::uint32_t{10});
    auto [b0, b1] = dpf::make_additive_shares(std::uint32_t{3});
    auto c0 = a0 * 2 + b0;
    auto c1 = a1 * 2 + b1;
    EXPECT_EQ(dpf::reconstruct(c0, c1), 23u);

    auto [s0, s1] = dpf::make_subtractive_shares(std::uint32_t{10});
    auto [t0, t1] = dpf::make_subtractive_shares(std::uint32_t{3});
    auto u0 = s0 * 2 - t0;
    auto u1 = s1 * 2 - t1;
    EXPECT_EQ(dpf::reconstruct(u0, u1), 17u);
}

TEST(SecretShare, CrossSchemeSigns)
{
    auto [a0, a1] = dpf::make_additive_shares(std::int32_t{20});
    auto [s0, s1] = dpf::make_subtractive_shares(std::int32_t{7});

    // party 0: raw add; party 1: flip the differing-scheme operand
    auto r0 = a0 + s0;
    auto r1 = a1 + s1;
    EXPECT_EQ(dpf::reconstruct(r0, r1), 27);

    auto q0 = s0 + a0;
    auto q1 = s1 + a1;
    EXPECT_EQ(dpf::reconstruct(q0, q1), 27);
}

TEST(SecretShare, PlaintextAbsorbOnParty0)
{
    auto [s0, s1] = dpf::make_subtractive_shares(std::uint32_t{5});
    s0 += std::uint32_t{10};
    s1 += std::uint32_t{10};  // no-op for party 1
    EXPECT_EQ(dpf::reconstruct(s0, s1), 15u);
}

TEST(SecretShare, AsAdditivePreservesSecret)
{
    auto [s0, s1] = dpf::make_subtractive_shares(std::int32_t{42});
    auto a0 = s0.as_additive();
    auto a1 = s1.as_additive();
    EXPECT_EQ(dpf::reconstruct(a0, a1), 42);
}

TEST(SecretShare, DpfLeafRoundTrip)
{
    const std::uint8_t alpha = 0x2a;
    const std::uint32_t beta = 0x01020304;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    EXPECT_TRUE(dpf::is_party_key_v<decltype(k0)>);
    EXPECT_EQ(decltype(k0)::party, 0u);
    EXPECT_EQ(decltype(k1)::party, 1u);

    auto y0 = *dpf::eval_point(k0, alpha);
    auto y1 = *dpf::eval_point(k1, alpha);
    EXPECT_TRUE((std::is_same_v<decltype(y0), dpf::subtractive_share<std::uint32_t, 0>>));
    EXPECT_EQ(dpf::reconstruct(y0, y1), beta);

    auto z0 = *dpf::eval_point(k0, static_cast<std::uint8_t>(alpha ^ 1));
    auto z1 = *dpf::eval_point(k1, static_cast<std::uint8_t>(alpha ^ 1));
    EXPECT_EQ(dpf::reconstruct(z0, z1), 0u);
}

TEST(SecretShare, DpfXorLeafRoundTrip)
{
    using X = dpf::xor_wrapper<std::uint32_t>;
    const std::uint8_t alpha = 7;
    const X beta{0xA5A5A5A5u};
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    auto y0 = *dpf::eval_point(k0, alpha);
    auto y1 = *dpf::eval_point(k1, alpha);
    EXPECT_EQ(dpf::reconstruct(y0, y1), beta);
}

TEST(SecretShare, CmpAdditiveRoundTrip)
{
    const std::uint32_t alpha = 100u;
    const std::uint64_t yt = 5u, yf = 9u;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::lt(yt, yf));
    const std::uint64_t mask = k0.cmp().mask;
    auto below = dpf::reconstruct(
        dpf::eval_point(dpf::cmp, k0, 50u),
        dpf::eval_point(dpf::cmp, k1, 50u)) & mask;
    auto at = dpf::reconstruct(
        dpf::eval_point(dpf::cmp, k0, alpha),
        dpf::eval_point(dpf::cmp, k1, alpha)) & mask;
    EXPECT_EQ(below, yt);
    EXPECT_EQ(at, yf);
}

TEST(SecretShare, PrgExpandSubtractive)
{
    using prg = dpf::prg::aes128;
    const auto seed0 = dpf::uniform_sample<prg::block_type>();
    const auto seed1 = dpf::uniform_sample<prg::block_type>();
    auto s0 = prg::expand<std::uint64_t, 0>(seed0, 0);
    auto s1 = prg::expand<std::uint64_t, 1>(seed1, 0);
    // Same pos, different seeds: reconstruct is raw0 - raw1 (bit pattern).
    EXPECT_EQ(dpf::reconstruct(s0, s1),
        static_cast<std::uint64_t>(s0.raw() - s1.raw()));

    // Same seed → same bits → reconstruct 0
    auto t0 = prg::expand<std::uint32_t, 0>(seed0, 3);
    auto t1 = prg::expand<std::uint32_t, 1>(seed0, 3);
    EXPECT_EQ(dpf::reconstruct(t0, t1), 0u);
}

}  // namespace