#include <gtest/gtest.h>

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

#include "dpf.hpp"
#include "dpf/yao.hpp"
#include "dpf/yao_share.hpp"

namespace
{

template <typename Ring>
Ring xor_bits(const std::vector<std::uint8_t> & a, const std::vector<std::uint8_t> & b)
{
    Ring v{};
    for (std::size_t i = 0; i < a.size() && i < 8u * sizeof(Ring); ++i)
        v = static_cast<Ring>(v | (static_cast<Ring>(a[i] ^ b[i]) << i));
    return v;
}

template <typename Ring>
Ring low_mask(unsigned width)
{
    if (width == 0 || width >= 8u * sizeof(Ring))
        return static_cast<Ring>(~Ring{0});
    return static_cast<Ring>((Ring{1} << width) - Ring{1});
}

template <typename Ring>
void expect_additive(Ring secret, unsigned width)
{
    const unsigned bits = width == 0 ? static_cast<unsigned>(8u * sizeof(Ring)) : width;
    const Ring r = dpf::uniform_sample<Ring>();
    const auto a0 = dpf::additive_share<Ring, 0>::from_raw(r);
    const auto a1 = dpf::additive_share<Ring, 1>::from_raw(static_cast<Ring>(secret - r));
    const auto [y0, y1] = dpf::yao::a2y(a0, a1, width);
    ASSERT_EQ(y0.size(), bits);
    ASSERT_EQ(y1.size(), bits);
    EXPECT_EQ(xor_bits<Ring>(y0, y1), static_cast<Ring>(secret & low_mask<Ring>(bits)));
    const auto [b0, b1] = dpf::yao::y2a<Ring>(y0, y1, bits);
    EXPECT_EQ(dpf::reconstruct(b0, b1), static_cast<Ring>(secret & low_mask<Ring>(bits)));
}

template <typename Ring>
void expect_subtractive(Ring s0, Ring s1, unsigned width)
{
    const unsigned bits = width == 0 ? static_cast<unsigned>(8u * sizeof(Ring)) : width;
    const Ring opened = static_cast<Ring>(s0 - s1);
    const auto b0 = dpf::subtractive_share<Ring, 0>::from_raw(s0);
    const auto b1 = dpf::subtractive_share<Ring, 1>::from_raw(s1);
    const auto [y0, y1] = dpf::yao::b2y(b0, b1, width);
    EXPECT_EQ(xor_bits<Ring>(y0, y1), static_cast<Ring>(opened & low_mask<Ring>(bits)));
    const auto [z0, z1] = dpf::yao::y2b<Ring>(y0, y1, bits);
    EXPECT_EQ(dpf::reconstruct(z0, z1), static_cast<Ring>(opened & low_mask<Ring>(bits)));

    const auto f0 = b0.as_fss();
    const auto f1 = b1.as_fss();
    const auto [g0, g1] = dpf::yao::fss2y(f0, f1, width);
    EXPECT_EQ(xor_bits<Ring>(g0, g1), static_cast<Ring>(opened & low_mask<Ring>(bits)));
    const auto [h0, h1] = dpf::yao::y2fss<Ring>(g0, g1, bits);
    EXPECT_EQ(dpf::reconstruct(h0, h1), static_cast<Ring>(opened & low_mask<Ring>(bits)));
}

}  // namespace

TEST(YaoShare, WidthSweep)
{
    const unsigned widths64[] = {1, 2, 7, 8, 9, 16, 31, 32, 63, 64, 0};
    for (unsigned w : widths64)
        expect_additive<std::uint64_t>(0x1234abcd5ull, w);
    const unsigned widths32[] = {1, 8, 31, 32, 0};
    for (unsigned w : widths32)
        expect_additive<std::uint32_t>(0x89abcdefu, w);
    expect_additive<std::uint16_t>(0xBEEFu, 0);
    expect_additive<std::uint16_t>(0xBEEFu, 9);
    expect_additive<std::uint8_t>(0xA5u, 0);
    expect_additive<std::uint8_t>(0xA5u, 3);
}

TEST(YaoShare, CornerSecrets)
{
    expect_additive<std::uint64_t>(0, 64);
    expect_additive<std::uint64_t>(1, 64);
    expect_additive<std::uint64_t>(~std::uint64_t{0}, 64);
    expect_additive<std::uint64_t>(std::uint64_t{1} << 63, 64);
    expect_additive<std::uint64_t>(0x105u, 8);
    expect_subtractive<std::uint64_t>(0, 0, 64);
    expect_subtractive<std::uint64_t>(1, 3, 64);
    expect_subtractive<std::uint64_t>(3, 1, 16);
    expect_subtractive<std::uint32_t>(0, 1, 32);
    expect_subtractive<std::uint8_t>(0x10, 0x10, 8);
}

TEST(YaoShare, RandomRoundTrips)
{
    for (int i = 0; i < 24; ++i)
    {
        const auto secret = dpf::uniform_sample<std::uint64_t>();
        expect_additive<std::uint64_t>(secret, 64);
        const auto s0 = dpf::uniform_sample<std::uint32_t>();
        const auto s1 = dpf::uniform_sample<std::uint32_t>();
        expect_subtractive<std::uint32_t>(s0, s1, 32);
    }
}

TEST(YaoShare, BitFlipChangesThatPlace)
{
    const auto a0 = dpf::additive_share<std::uint64_t, 0>::from_raw(0x111u);
    const auto a1 = dpf::additive_share<std::uint64_t, 1>::from_raw(0x222u);
    auto [y0, y1] = dpf::yao::a2y(a0, a1, 12);
    const auto opened = xor_bits<std::uint64_t>(y0, y1);
    for (unsigned i = 0; i < 12; ++i)
    {
        auto flipped = y0;
        flipped[i] = static_cast<std::uint8_t>(flipped[i] ^ 1u);
        EXPECT_EQ(xor_bits<std::uint64_t>(flipped, y1), opened ^ (std::uint64_t{1} << i));
        const auto [b0, b1] = dpf::yao::y2a<std::uint64_t>(flipped, y1, 12);
        EXPECT_EQ(dpf::reconstruct(b0, b1), opened ^ (std::uint64_t{1} << i));
    }
}

TEST(YaoShare, DefaultWidthUsesTheVector)
{
    const auto a0 = dpf::additive_share<std::uint16_t, 0>::from_raw(0x00FFu);
    const auto a1 = dpf::additive_share<std::uint16_t, 1>::from_raw(0);
    const auto [y0, y1] = dpf::yao::a2y(a0, a1, 8);
    const auto [b0, b1] = dpf::yao::y2a<std::uint16_t>(y0, y1);
    EXPECT_EQ(dpf::reconstruct(b0, b1), 0x00FFu);
}

TEST(YaoShare, RejectsBadInputs)
{
    const auto a0 = dpf::additive_share<std::uint64_t, 0>::from_raw(1);
    const auto a1 = dpf::additive_share<std::uint64_t, 1>::from_raw(0);
    EXPECT_THROW(dpf::yao::a2y(a0, a1, 65), std::invalid_argument);
    EXPECT_THROW(dpf::yao::a2y(
        dpf::additive_share<std::uint8_t, 0>::from_raw(1),
        dpf::additive_share<std::uint8_t, 1>::from_raw(0), 9),
        std::invalid_argument);

    std::vector<std::uint8_t> bits(8, 0);
    std::vector<std::uint8_t> shortv(3, 0);
    EXPECT_THROW(dpf::yao::y2a<std::uint64_t>(shortv, bits, 8), std::invalid_argument);
    bits[4] = 2;
    EXPECT_THROW(dpf::yao::y2a<std::uint64_t>(bits, std::vector<std::uint8_t>(8, 0), 8),
        std::invalid_argument);

    const auto rss = dpf::make_replicated_shares<std::uint32_t>(1, 2, 3);
    auto r0 = std::get<0>(rss);
    r0.next = 99;
    EXPECT_THROW(dpf::yao::rss2y(r0, std::get<1>(rss)), std::invalid_argument);
}

TEST(YaoShare, ReplicatedManyAndFresh)
{
    for (int i = 0; i < 12; ++i)
    {
        const auto x0 = dpf::uniform_sample<std::uint32_t>();
        const auto x1 = dpf::uniform_sample<std::uint32_t>();
        const auto x2 = dpf::uniform_sample<std::uint32_t>();
        const auto secret = static_cast<std::uint32_t>(x0 + x1 + x2);
        const auto shares = dpf::make_replicated_shares(x0, x1, x2);
        const auto [y0, y1] = dpf::yao::rss2y(std::get<0>(shares), std::get<1>(shares));
        EXPECT_EQ(xor_bits<std::uint32_t>(y0, y1), secret);
        const auto back = dpf::yao::y2rss<std::uint32_t>(y0, y1);
        EXPECT_EQ(dpf::reconstruct(std::get<0>(back), std::get<1>(back)), secret);
        EXPECT_EQ(dpf::reconstruct(std::get<1>(back), std::get<2>(back)), secret);
        EXPECT_EQ(dpf::reconstruct(std::get<0>(back), std::get<2>(back)), secret);
        const auto again = dpf::yao::y2rss<std::uint32_t>(y0, y1);
        EXPECT_NE(std::get<0>(back).own, std::get<0>(again).own);
    }
}

TEST(YaoShare, ConcreteAdditiveSum)
{
    const auto a0 = dpf::additive_share<std::uint16_t, 0>::from_raw(0x1234);
    const auto a1 = dpf::additive_share<std::uint16_t, 1>::from_raw(0x0100);
    const auto [y0, y1] = dpf::yao::a2y(a0, a1, 16);
    EXPECT_EQ(xor_bits<std::uint16_t>(y0, y1), static_cast<std::uint16_t>(0x1334));
}

TEST(YaoShare, PointLeafOnAndOff)
{
    for (int n = 0; n < 16; ++n)
    {
        const auto alpha = dpf::uniform_sample<std::uint8_t>();
        const auto beta = dpf::uniform_sample<std::uint32_t>();
        auto [k0, k1] = dpf::make_dpf(alpha, beta);
        auto s0 = *dpf::eval_point(k0, alpha);
        auto s1 = *dpf::eval_point(k1, alpha);
        const auto [y0, y1] = dpf::yao::b2y(s0, s1, 32);
        EXPECT_EQ(xor_bits<std::uint32_t>(y0, y1), beta);
        const auto [z0, z1] = dpf::yao::y2b<std::uint32_t>(y0, y1, 32);
        EXPECT_EQ(dpf::reconstruct(z0, z1), beta);

        const auto other = static_cast<std::uint8_t>(alpha + 1u);
        auto t0 = *dpf::eval_point(k0, other);
        auto t1 = *dpf::eval_point(k1, other);
        const auto [u0, u1] = dpf::yao::b2y(t0, t1, 32);
        EXPECT_EQ(xor_bits<std::uint32_t>(u0, u1), 0u);
    }
}

TEST(YaoShare, PointLeafPlusPublicAndSecondLeaf)
{
    const std::uint8_t alpha = 9;
    const std::uint32_t beta = 40;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    auto s0 = *dpf::eval_point(k0, alpha);
    auto s1 = *dpf::eval_point(k1, alpha);
    s0 += std::uint32_t{5};
    const auto [y0, y1] = dpf::yao::b2y(s0, s1, 32);
    EXPECT_EQ(xor_bits<std::uint32_t>(y0, y1), 45u);

    auto [j0, j1] = dpf::make_dpf(alpha, std::uint32_t{7});
    auto u0 = *dpf::eval_point(k0, alpha) + *dpf::eval_point(j0, alpha);
    auto u1 = *dpf::eval_point(k1, alpha) + *dpf::eval_point(j1, alpha);
    const auto [v0, v1] = dpf::yao::b2y(u0, u1, 32);
    EXPECT_EQ(xor_bits<std::uint32_t>(v0, v1), 47u);
}

TEST(YaoShare, ComparisonLeaf)
{
    const std::uint32_t alpha = 100u;
    const std::uint64_t yt = 5u;
    const std::uint64_t yf = 9u;
    auto [k0, k1] = dpf::make_dpf(alpha, dpf::lt(yt, yf));
    const std::uint64_t mask = k0.cmp().mask;
    auto c0 = dpf::eval_point(dpf::cmp, k0, 50u);
    auto c1 = dpf::eval_point(dpf::cmp, k1, 50u);
    const auto [y0, y1] = dpf::yao::a2y(c0, c1, 64);
    const auto bits = xor_bits<std::uint64_t>(y0, y1);
    EXPECT_EQ(bits, dpf::reconstruct(c0, c1));
    EXPECT_EQ(bits & mask, yt);

    auto d0 = dpf::eval_point(dpf::cmp, k0, alpha);
    auto d1 = dpf::eval_point(dpf::cmp, k1, alpha);
    const auto [z0, z1] = dpf::yao::a2y(d0, d1, 64);
    EXPECT_EQ(xor_bits<std::uint64_t>(z0, z1) & mask, yf);
}

TEST(YaoShare, LeafBitsThroughANetlist)
{
    const std::uint8_t alpha = 42;
    const std::uint32_t beta = 0x6bu;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    auto s0 = *dpf::eval_point(k0, alpha);
    auto s1 = *dpf::eval_point(k1, alpha);
    const auto [p0, p1] = dpf::yao::b2y(s0, s1, 8);

    dpf::yao::netlist n;
    dpf::yao::bit in[8];
    for (int i = 0; i < 8; ++i)
        in[i] = n.shared_in();
    n.out(n.and_(in[0], in[1]));
    for (int i = 0; i < 8; ++i)
        n.out(in[i]);

    auto [o0, o1] = dpf::yao::eval_pair(n, p0.data(), p1.data());
    ASSERT_EQ(o0.size(), 9u);
    EXPECT_EQ(o0[0] ^ o1[0], static_cast<std::uint8_t>((beta & 1u) & ((beta >> 1) & 1u)));
    for (int i = 0; i < 8; ++i)
        EXPECT_EQ(o0[static_cast<std::size_t>(i + 1)] ^ o1[static_cast<std::size_t>(i + 1)],
            static_cast<std::uint8_t>((beta >> i) & 1u));

    std::vector<std::uint8_t> and0{o0[0]};
    std::vector<std::uint8_t> and1{o1[0]};
    const auto [z0, z1] = dpf::yao::y2b<std::uint32_t>(and0, and1, 1);
    EXPECT_EQ(dpf::reconstruct(z0, z1), 1u);
}

TEST(YaoShare, RandomNetlistMatchesPlain)
{
    std::uint32_t rng = 0xC0FFEEu;
    auto nxt = [&] {
        rng = rng * 1664525u + 1013904223u;
        return rng;
    };
    for (int trial = 0; trial < 20; ++trial)
    {
        dpf::yao::netlist n;
        dpf::yao::bit w[24];
        std::uint8_t sem[4];
        std::uint8_t p0[4];
        std::uint8_t p1[4];
        int nwire = 4;
        for (int i = 0; i < 4; ++i)
        {
            w[i] = n.shared_in();
            sem[i] = static_cast<std::uint8_t>(nxt() & 1u);
            p0[i] = static_cast<std::uint8_t>(nxt() & 1u);
            p1[i] = static_cast<std::uint8_t>(sem[i] ^ p0[i]);
        }
        for (int g = 0; g < 32; ++g)
        {
            const int ia = static_cast<int>(nxt() % static_cast<std::uint32_t>(nwire));
            const int ib = static_cast<int>(nxt() % static_cast<std::uint32_t>(nwire));
            dpf::yao::bit d;
            switch (nxt() % 3u)
            {
            case 0: d = n.xor_(w[ia], w[ib]); break;
            case 1: d = n.and_(w[ia], w[ib]); break;
            default: d = n.not_(w[ia]); break;
            }
            if (nwire < 24)
                w[nwire++] = d;
        }
        for (int i = 0; i < 4; ++i)
            n.out(w[nwire - 1 - i]);
        const auto plain = dpf::yao::eval_plain(n, sem);
        EXPECT_EQ(dpf::yao::eval_local(n, sem), plain) << trial;
        auto [a, b] = dpf::yao::eval_pair(n, p0, p1);
        for (std::size_t i = 0; i < plain.size(); ++i)
            EXPECT_EQ(static_cast<std::uint8_t>(a[i] ^ b[i]), plain[i]) << trial;
    }
}

TEST(YaoShare, MixedPrivateInputs)
{
    dpf::yao::netlist n;
    const auto a = n.priv_in(0);
    const auto b = n.priv_in(1);
    const auto c = n.shared_in();
    n.out(n.xor_(n.and_(a, b), c));
    n.out(n.xor_public(c, 0));
    n.out(n.not_(n.not_(a)));
    const std::uint8_t p0[3] = {1, 0, 1};
    const std::uint8_t p1[3] = {9, 1, 0};
    auto [o0, o1] = dpf::yao::eval_pair(n, p0, p1);
    EXPECT_EQ(static_cast<std::uint8_t>(o0[0] ^ o1[0]), 0u);
    EXPECT_EQ(static_cast<std::uint8_t>(o0[1] ^ o1[1]), 1u);
    EXPECT_EQ(static_cast<std::uint8_t>(o0[2] ^ o1[2]), 1u);
}

TEST(YaoShare, NetlistRejects)
{
    dpf::yao::netlist n;
    EXPECT_THROW(n.priv_in(2), std::invalid_argument);
    EXPECT_THROW(n.xor_(dpf::yao::bit{3}, dpf::yao::bit{0}), std::invalid_argument);
    const auto a = n.shared_in();
    n.out(a);
    EXPECT_THROW(n.shared_in(), std::logic_error);
    std::uint8_t bad = 2;
    EXPECT_THROW(dpf::yao::eval_local(n, &bad), std::invalid_argument);
}