#include <gtest/gtest.h>

#include "dpf.hpp"

#include <cstdint>
#include <cstring>

namespace
{

simde__m128i g_roots[16];
int g_ri = 0;
simde__m128i take_root() { return g_roots[g_ri++]; }

struct Pad
{
    uint64_t n = 1;
    simde__m128i block()
    {
        auto v = simde_mm_set_epi64x(static_cast<long long>(n),
            static_cast<long long>(n * 9 + 3));
        n += 2;
        return v;
    }
    uint8_t bit() { return static_cast<uint8_t>(n++ & 1u); }
};

void reset_roots()
{
    g_ri = 0;
    for (int i = 0; i < 16; ++i)
        g_roots[i] = simde_mm_set_epi64x(0x2222 * (i + 1), 0xBEEF0000u + i * 13);
}

template <typename T>
T bare(const T & v)
{
    return v;
}

template <typename T, std::size_t Party, dpf::sharing Scheme>
T bare(const dpf::secret_share<T, Party, Scheme> & s)
{
    return s.raw();
}

template <typename A, typename B>
auto recon(const A & a, const B & b)
{
    using T = decltype(bare(a));
    return static_cast<T>(bare(a) - bare(b));
}

template <typename Key, typename In>
auto ev(const Key & key, In x)
{
    return bare(*dpf::eval_point(key, x));
}

}  // namespace

TEST(ArithPayload, DsXorIndexMatchesDealer)
{
    using in_t = std::uint8_t;
    using out_t = std::uint32_t;
    const in_t alpha = 0x2a;
    const in_t x0 = 0x55;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);
    const out_t beta = 0x01020304;
    const out_t y0 = 0x00010002;
    const out_t y1 = static_cast<out_t>(beta - y0);

    reset_roots();
    auto dealer = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        beta);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat(dpf::arith_output, x0, x1, rng, y0, y1);

    using key_t = std::decay_t<decltype(dealer.first)>;
    EXPECT_EQ(std::memcmp(&dealer.first.root(), &ds.first.root(),
        sizeof(simde__m128i)), 0);
    EXPECT_EQ(std::memcmp(&dealer.second.root(), &ds.second.root(),
        sizeof(simde__m128i)), 0);
    for (std::size_t i = 0; i < key_t::depth; ++i)
    {
        EXPECT_EQ(std::memcmp(&dealer.first.correction_word(i),
            &ds.first.correction_word(i), sizeof(simde__m128i)), 0)
            << "cw " << i;
        EXPECT_EQ(dealer.first.correction_advice(i),
            ds.first.correction_advice(i))
            << "advice " << i;
    }
    EXPECT_EQ(std::memcmp(&dealer.first.leaf(), &ds.first.leaf(),
        sizeof(dealer.first.leaf())), 0);

    for (int i = 0; i < 256; ++i)
    {
        const in_t q = static_cast<in_t>(i);
        EXPECT_EQ(recon(ev(ds.first, q), ev(ds.second, q)),
            q == alpha ? beta : out_t{})
            << i;
        EXPECT_EQ(ev(ds.first, q), ev(dealer.first, q)) << "p0 " << i;
        EXPECT_EQ(ev(ds.second, q), ev(dealer.second, q)) << "p1 " << i;
    }
}

TEST(ArithPayload, DsArithIndexMatchesDealer)
{
    using in_t = std::uint8_t;
    using out_t = std::uint16_t;
    const in_t alpha = 0xc0;
    const in_t x0 = 0x40;
    const in_t x1 = static_cast<in_t>(alpha - x0);
    const out_t beta = 9;
    const out_t y0 = 3;
    const out_t y1 = 6;

    reset_roots();
    auto dealer = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        beta);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat(dpf::arith_input, dpf::arith_output,
        x0, x1, rng, y0, y1);

    for (int i = 0; i < 256; ++i)
    {
        const in_t q = static_cast<in_t>(i);
        EXPECT_EQ(recon(ev(ds.first, q), ev(ds.second, q)),
            q == alpha ? beta : out_t{})
            << i;
        EXPECT_EQ(ev(ds.first, q), ev(dealer.first, q)) << "p0 " << i;
        EXPECT_EQ(ev(ds.second, q), ev(dealer.second, q)) << "p1 " << i;
    }
}

TEST(ArithPayload, GenevalXorSharedBeta)
{
    using in_t = std::uint8_t;
    using out_t = std::uint16_t;
    const in_t alpha = 0x33;
    const in_t x0 = 0x0f;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);
    const out_t beta = 0x77;
    const out_t y0 = 0x10;
    const out_t y1 = static_cast<out_t>(beta - y0);

    reset_roots();
    auto keys = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        beta);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> g_rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto g = dpf::geneval_point(dpf::arith_output, x0, x1, alpha, g_rng, y0, y1);

    EXPECT_TRUE(g.leaf_live);
    ASSERT_EQ(g.party0.size(), 1u);
    EXPECT_EQ(recon(g.party0[0], g.party1[0]), beta);
    EXPECT_EQ(g.party0[0], ev(keys.first, alpha));
    EXPECT_EQ(g.party1[0], ev(keys.second, alpha));
}

TEST(ArithPayload, GenevalArithSharedBeta)
{
    using in_t = std::uint8_t;
    using out_t = std::uint16_t;
    const in_t alpha = 0x90;
    const in_t x0 = 0x20;
    const in_t x1 = static_cast<in_t>(alpha - x0);
    const out_t beta = 3;
    const out_t y0 = 1;
    const out_t y1 = 2;

    reset_roots();
    auto keys = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        beta);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> g_rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto g = dpf::geneval_point(dpf::arith_input, dpf::arith_output, x0, x1,
        alpha, g_rng, y0, y1);

    EXPECT_TRUE(g.leaf_live);
    ASSERT_EQ(g.party0.size(), 1u);
    EXPECT_EQ(recon(g.party0[0], g.party1[0]), beta);
    EXPECT_EQ(g.party0[0], ev(keys.first, alpha));
    EXPECT_EQ(g.party1[0], ev(keys.second, alpha));
}

TEST(ArithPayload, XorWrapperSharesMatchDealer)
{
    using in_t = std::uint8_t;
    using out_t = dpf::xor_wrapper<std::uint32_t>;
    const in_t alpha = 0x11;
    const in_t x0 = 0x55;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);
    const out_t beta{0x0a0b0c0du};
    const out_t y0{0x01020304u};
    const out_t y1 = beta ^ y0;

    reset_roots();
    auto dealer = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        beta);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat(dpf::arith_output, x0, x1, rng, y0, y1);

    for (int i = 0; i < 256; ++i)
    {
        const in_t q = static_cast<in_t>(i);
        const auto got = dpf::reconstruct(*dpf::eval_point(ds.first, q),
            *dpf::eval_point(ds.second, q));
        const auto expect = dpf::reconstruct(*dpf::eval_point(dealer.first, q),
            *dpf::eval_point(dealer.second, q));
        EXPECT_EQ(got, expect) << i;
        EXPECT_EQ(got, q == alpha ? beta : out_t{}) << i;
    }

    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> g_rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto g = dpf::geneval_point(dpf::arith_output, x0, x1, alpha, g_rng, y0, y1);
    EXPECT_TRUE(g.leaf_live);
    ASSERT_EQ(g.party0.size(), 1u);
    EXPECT_EQ(g.party0[0] ^ g.party1[0], beta);
}

TEST(ArithPayload, MultiBlockUint256MatchesDealer)
{
    using in_t = std::uint8_t;
    using out_t = uint256_t;
    const in_t alpha = 0x2a;
    const in_t x0 = 0x0f;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);
    const out_t beta = (out_t{1} << 200) + out_t{0xabcdefu};
    const out_t y0 = (out_t{1} << 180) + out_t{0x1111u};
    const out_t y1 = beta - y0;

    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    EXPECT_GT((dpf::block_length_of_leaf_v<out_t, simde__m128i>), 1u);
    HEDLEY_PRAGMA(GCC diagnostic pop)

    reset_roots();
    auto dealer = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        beta);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat(dpf::arith_output, x0, x1, rng, y0, y1);

    EXPECT_EQ(std::memcmp(&dealer.first.leaf(), &ds.first.leaf(),
        sizeof(dealer.first.leaf())), 0);

    for (int i = 0; i < 256; i += 17)
    {
        const in_t q = static_cast<in_t>(i);
        EXPECT_EQ(recon(ev(ds.first, q), ev(ds.second, q)),
            q == alpha ? beta : out_t{})
            << i;
        EXPECT_EQ(ev(ds.first, q), ev(dealer.first, q)) << "p0 " << i;
        EXPECT_EQ(ev(ds.second, q), ev(dealer.second, q)) << "p1 " << i;
    }
}

TEST(ArithPayload, IncrementalMultiSlotArithBeta)
{
    using in_t = std::uint8_t;
    const in_t alpha = 0x44;
    const in_t x0 = 0x12;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);
    const std::uint16_t b0 = 0x1111;
    const std::uint16_t b1 = 0x2222;
    const std::uint16_t y00 = 0x0100;
    const std::uint16_t y01 = static_cast<std::uint16_t>(b0 - y00);
    const std::uint16_t y10 = 0x0003;
    const std::uint16_t y11 = static_cast<std::uint16_t>(b1 - y10);

    reset_roots();
    auto dealer = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<8>(b0, b1));
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat(dpf::arith_output, x0, x1, rng,
        dpf::at<8>(dpf::arith_beta<std::uint16_t>{y00, y01},
            dpf::arith_beta<std::uint16_t>{y10, y11}));

    auto r0 = [&](auto & k0, auto & k1, in_t q) {
        return recon(*dpf::eval_point(dpf::out<0, 8>, k0, q),
            *dpf::eval_point(dpf::out<0, 8>, k1, q));
    };
    auto r1 = [&](auto & k0, auto & k1, in_t q) {
        return recon(*dpf::eval_point(dpf::out<1, 8>, k0, q),
            *dpf::eval_point(dpf::out<1, 8>, k1, q));
    };

    for (int i = 0; i < 256; i += 13)
    {
        const in_t q = static_cast<in_t>(i);
        EXPECT_EQ(r0(ds.first, ds.second, q), r0(dealer.first, dealer.second, q))
            << "s0 " << i;
        EXPECT_EQ(r1(ds.first, ds.second, q), r1(dealer.first, dealer.second, q))
            << "s1 " << i;
        EXPECT_EQ(r0(ds.first, ds.second, q),
            q == alpha ? b0 : std::uint16_t{0})
            << i;
        EXPECT_EQ(r1(ds.first, ds.second, q),
            q == alpha ? b1 : std::uint16_t{0})
            << i;
    }
}

TEST(ArithPayload, MixedArithBetaAndWildcard)
{
    using in_t = std::uint8_t;
    const in_t alpha = 0x70;
    const in_t x0 = 0x01;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);
    const std::uint8_t beta = 9;
    const std::uint8_t y0 = 2;
    const std::uint8_t y1 = 7;

    reset_roots();
    auto dealer = dpf::make_dpf(alpha, dpf::root_sampler_t<dpf::prg::aes128>{take_root},
        dpf::at<8>(beta, dpf::wildcard_value<std::uint8_t>{}));
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat(dpf::arith_output, x0, x1, rng,
        dpf::at<8>(dpf::arith_beta<std::uint8_t>{y0, y1},
            dpf::wildcard_value<std::uint8_t>{}));

    using key_t = std::decay_t<decltype(ds.first)>;
    static_assert(dpf::is_wildcard_v<typename key_t::template output_type_t<1>>);

    for (int i = 0; i < 256; i += 19)
    {
        const in_t q = static_cast<in_t>(i);
        const auto got = recon(*dpf::eval_point(dpf::out<0, 8>, ds.first, q),
            *dpf::eval_point(dpf::out<0, 8>, ds.second, q));
        const auto expect =
            recon(*dpf::eval_point(dpf::out<0, 8>, dealer.first, q),
                *dpf::eval_point(dpf::out<0, 8>, dealer.second, q));
        EXPECT_EQ(got, expect) << i;
        EXPECT_EQ(got, q == alpha ? beta : std::uint8_t{0}) << i;
    }
}