#include <gtest/gtest.h>

#include "dpf.hpp"

#include <cstdint>
#include <cstring>
#include <vector>

namespace
{

using ht_prg = dpf::prg::aes128_ccr;
using leaf_prg = dpf::prg::aes128;

simde__m128i g_roots[8];
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 < 8; ++i)
        g_roots[i] = simde_mm_set_epi64x(0x1111 * (i + 1), 0xA5A50000u + i * 17);
}

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(HalfTree, TraitsSelectedByCcrPrg)
{
    static_assert(dpf::tree_traits<ht_prg>::is_half_tree);
    static_assert(!dpf::tree_traits<leaf_prg>::is_half_tree);
    static_assert(dpf::tree_traits<ht_prg>::last_level_differs);
    static_assert(!dpf::tree_traits<ht_prg>::stores_mid_advice);
}

TEST(HalfTree, DealerPointAndFullDomain)
{
    using in_t = std::uint8_t;
    using out_t = std::uint32_t;
    const in_t alpha = 0x2a;
    const out_t beta = 0x01020304;

    auto [k0, k1] = dpf::make_dpf<ht_prg, leaf_prg>(alpha, beta);
    using key_t = std::decay_t<decltype(k0)>;
    static_assert(key_t::tree::is_half_tree);

    // Mid-level advice unused; last level may pack BGI-style advice.
    for (std::size_t i = 0; i + 1 < key_t::depth; ++i)
        EXPECT_EQ(k0.correction_advice(i), 0) << "mid advice " << i;

    for (int i = 0; i < 256; ++i)
    {
        const in_t q = static_cast<in_t>(i);
        const out_t got = recon(ev(k0, q), ev(k1, q));
        EXPECT_EQ(got, q == alpha ? beta : out_t{}) << i;
    }

    auto [buf0, it0] = dpf::eval_full(k0);
    auto [buf1, it1] = dpf::eval_full(k1);
    (void)it0;
    (void)it1;
    for (int i = 0; i < 256; ++i)
    {
        const out_t got = recon(buf0[i], buf1[i]);
        EXPECT_EQ(got, static_cast<in_t>(i) == alpha ? beta : out_t{}) << i;
    }
}

TEST(HalfTree, OutputParityVsBgi)
{
    using in_t = std::uint8_t;
    using out_t = std::uint16_t;
    const in_t alpha = 0x7e;
    const out_t beta = 0xabcd;

    auto [h0, h1] = dpf::make_dpf<ht_prg, leaf_prg>(alpha, beta);
    auto [b0, b1] = dpf::make_dpf<leaf_prg, leaf_prg>(alpha, beta);

    // Keys differ (different tree), but reconstructed outputs match.
    EXPECT_NE(std::memcmp(&h0.root(), &b0.root(), sizeof(simde__m128i)), 0);

    for (int i = 0; i < 256; ++i)
    {
        const in_t q = static_cast<in_t>(i);
        const out_t ht = recon(ev(h0, q), ev(h1, q));
        const out_t bgi = recon(ev(b0, q), ev(b1, q));
        EXPECT_EQ(ht, bgi) << i;
        EXPECT_EQ(ht, q == alpha ? beta : out_t{}) << i;
    }
}

TEST(HalfTree, DoernerShelatXorMatchesDealer)
{
    using in_t = std::uint8_t;
    using out_t = std::uint16_t;
    const in_t alpha = 0x11;
    const in_t x0 = 0x55;
    const in_t x1 = static_cast<in_t>(alpha ^ x0);
    const out_t y = 0x42;

    reset_roots();
    auto dealer = dpf::make_dpf<ht_prg, leaf_prg>(alpha,
        dpf::root_sampler_t<ht_prg>{take_root}, y);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> ds_rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat<ht_prg, leaf_prg>(x0, x1, ds_rng, y);

    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;
    }

    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 ? y : out_t{})
            << i;
    }
}

TEST(HalfTree, DoernerShelatArithMatchesDealer)
{
    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 y = 9;

    reset_roots();
    auto dealer = dpf::make_dpf<ht_prg, leaf_prg>(alpha,
        dpf::root_sampler_t<ht_prg>{take_root}, y);
    reset_roots();
    HEDLEY_PRAGMA(GCC diagnostic push)
    HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
    dpf::ds_randomness<simde__m128i (*)(), Pad> ds_rng{take_root, Pad{}};
    HEDLEY_PRAGMA(GCC diagnostic pop)
    auto ds = dpf::make_dpf_doerner_shelat<ht_prg, leaf_prg>(dpf::arith_input,
        x0, x1, ds_rng, y);

    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 ? y : 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(HalfTree, GenevalPointMatchesDealer)
{
    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 y = 0x77;

    reset_roots();
    auto keys = dpf::make_dpf<ht_prg, leaf_prg>(alpha,
        dpf::root_sampler_t<ht_prg>{take_root}, y);
    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<ht_prg, leaf_prg>(x0, x1, alpha, g_rng, y);

    using key_t = std::decay_t<decltype(keys.first)>;
    EXPECT_TRUE(g.leaf_live);
    EXPECT_EQ(g.live_levels, key_t::depth);
    for (std::size_t i = 0; i < g.live_levels; ++i)
    {
        EXPECT_EQ(std::memcmp(&g.correction_words[i],
            &keys.first.correction_word(i), sizeof(simde__m128i)), 0)
            << "cw " << i;
        EXPECT_EQ(g.correction_advice[i], keys.first.correction_advice(i))
            << "advice " << i;
    }
    ASSERT_EQ(g.party0.size(), 1u);
    EXPECT_EQ(g.party0[0], ev(keys.first, alpha));
    EXPECT_EQ(g.party1[0], ev(keys.second, alpha));
    EXPECT_EQ(recon(g.party0[0], g.party1[0]), y);
}

TEST(HalfTree, GenevalArithPointMatchesDealer)
{
    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 y = 3;

    reset_roots();
    auto keys = dpf::make_dpf<ht_prg, leaf_prg>(alpha,
        dpf::root_sampler_t<ht_prg>{take_root}, y);
    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<ht_prg, leaf_prg>(dpf::arith_input, x0, x1,
        alpha, g_rng, y);

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

TEST(HalfTree, IncrementalCmpSmoke)
{
    using in_t = std::uint8_t;
    using ht_prg = dpf::prg::aes128_ccr;
    using leaf_prg = dpf::prg::aes128;
    const in_t alpha = 0x40;
    const uint64_t if_true = 4;
    const uint64_t if_false = 1;

    auto ht = dpf::make_dpf<ht_prg, leaf_prg>(alpha,
        dpf::lt(if_true, if_false));
    auto bgi = dpf::make_dpf<leaf_prg, leaf_prg>(alpha,
        dpf::lt(if_true, if_false));

    using key_t = std::decay_t<decltype(ht.first)>;
    static_assert(key_t::tree::is_half_tree);
    for (std::size_t i = 0; i + 1 < key_t::depth; ++i)
        EXPECT_EQ(ht.first.correction_advice(i), 0) << "mid advice " << i;

    const uint64_t mask = ht.first.cmp().mask;
    auto recon_cmp = [&](const auto & k0, const auto & k1, in_t q) {
        return dpf::reconstruct(dpf::eval_point(dpf::cmp, k0, q),
            dpf::eval_point(dpf::cmp, k1, q))
            & mask;
    };

    for (int i = 0; i < 256; ++i)
    {
        const in_t q = static_cast<in_t>(i);
        const auto ht_got = recon_cmp(ht.first, ht.second, q);
        const auto bgi_got = recon_cmp(bgi.first, bgi.second, q);
        const uint64_t expect = (q < alpha) ? if_true : if_false;
        EXPECT_EQ(ht_got, expect) << i;
        EXPECT_EQ(ht_got, bgi_got) << "parity " << i;
    }
}

TEST(HalfTree, IncrementalPlacementSmoke)
{
    using in_t = std::uint8_t;
    const in_t alpha = 0x40;
    auto keys = dpf::make_dpf<ht_prg, leaf_prg>(alpha, dpf::at<8>(uint8_t{7}));

    using key_t = std::decay_t<decltype(keys.first)>;
    static_assert(key_t::tree::is_half_tree);
    for (std::size_t i = 0; i + 1 < key_t::depth; ++i)
        EXPECT_EQ(keys.first.correction_advice(i), 0) << "mid advice " << i;

    for (int i = 0; i < 256; ++i)
    {
        const in_t q = static_cast<in_t>(i);
        const auto got = recon(*dpf::eval_point(keys.first, q),
            *dpf::eval_point(keys.second, q));
        EXPECT_EQ(got, q == alpha ? uint8_t{7} : uint8_t{0}) << i;
    }
}