#include <gtest/gtest.h>

#include <array>
#include <cstdint>
#include <cstring>
#include <exception>
#include <filesystem>
#include <stdexcept>
#include <string>
#include <thread>
#include <unistd.h>
#include <utility>
#include <vector>

#include "dpf.hpp"
#include "dpf/iknp.hpp"
#include "dpf/prg_aes_ccr.hpp"
#include "flow_util.hpp"
#include "iknp_deal.hpp"

namespace
{

using dpf::net::role;
using dpf::net::trio;

template <typename Fn0, typename Fn1>
void run_pair(Fn0 && fn0, Fn1 && fn1)
{
    const auto dir = std::filesystem::temp_directory_path()
        / ("libdpf_iknp_" + std::to_string(::getpid()));
    std::filesystem::create_directories(dir);
    std::exception_ptr ep0;
    std::exception_ptr ep1;
    std::thread t0([&] {
        try
        {
            auto net = trio::connect_pair(role::p0, dir.string());
            fn0(net);
        }
        catch (...)
        {
            ep0 = std::current_exception();
        }
    });
    std::thread t1([&] {
        try
        {
            auto net = trio::connect_pair(role::p1, dir.string());
            fn1(net);
        }
        catch (...)
        {
            ep1 = std::current_exception();
        }
    });
    t0.join();
    t1.join();
    std::filesystem::remove_all(dir);
    if (ep0)
        std::rethrow_exception(ep0);
    if (ep1)
        std::rethrow_exception(ep1);
}

simde__m128i xor_block(simde__m128i a, simde__m128i b)
{
    return simde_mm_xor_si128(a, b);
}

TEST(iknp, pads_match_dealer_relations)
{
    constexpr std::size_t nblock = 4;
    constexpr std::size_t nbit = 5;
    constexpr std::size_t nb2a = 3;
    constexpr std::size_t ncw = 2;
    dpf::iknp::material m0, m1;
    run_pair(
        [&](trio & net) {
            m0 = dpf::iknp::sample(net.to(role::p1), 0, nblock, nbit, nb2a, ncw);
        },
        [&](trio & net) {
            m1 = dpf::iknp::sample(net.to(role::p0), 1, nblock, nbit, nb2a, ncw);
        });
    ASSERT_EQ(m0.blocks.size(), nblock);
    ASSERT_EQ(m1.blocks.size(), nblock);
    for (std::size_t i = 0; i < nblock; ++i)
    {
        const auto a = static_cast<std::uint8_t>(m0.blocks[i].a ^ m1.blocks[i].a);
        const auto b = xor_block(m0.blocks[i].b, m1.blocks[i].b);
        const auto c = xor_block(m0.blocks[i].c, m1.blocks[i].c);
        const auto expect = (a & 1u) ? b : simde_mm_setzero_si128();
        EXPECT_EQ(std::memcmp(&c, &expect, sizeof(c)), 0);
    }
    for (std::size_t i = 0; i < nbit; ++i)
    {
        const auto a = static_cast<std::uint8_t>(m0.bits[i].a ^ m1.bits[i].a);
        const auto b = static_cast<std::uint8_t>(m0.bits[i].b ^ m1.bits[i].b);
        const auto c = static_cast<std::uint8_t>(m0.bits[i].c ^ m1.bits[i].c);
        EXPECT_EQ(c, static_cast<std::uint8_t>((a & b) & 1u));
    }
    for (std::size_t i = 0; i < nb2a; ++i)
    {
        const auto bit = static_cast<std::uint64_t>(m0.b2a[i].r ^ m1.b2a[i].r);
        EXPECT_EQ(m0.b2a[i].add + m1.b2a[i].add, bit);
    }
    // gamma0 ⊕ gamma1 = (bit1·rand0) ⊕ (bit0·rand1) (XOR shares; neither
    // party learns the peer pad bit).
    for (std::size_t i = 0; i < ncw; ++i)
    {
        const auto prod = xor_block(
            (m1.cws[i].bit & 1u) ? m0.cws[i].rand : simde_mm_setzero_si128(),
            (m0.cws[i].bit & 1u) ? m1.cws[i].rand : simde_mm_setzero_si128());
        const auto got = xor_block(m0.cws[i].gamma, m1.cws[i].gamma);
        EXPECT_EQ(std::memcmp(&got, &prod, sizeof(prod)), 0);
    }
}

template <typename Point>
void expect_point(Point alpha, Point x0, Point x1, std::uint64_t beta, bool additive)
{
    using namespace dpf::party;
    run_pair(
        [&](trio & net) {
            dist_with_point_key_iknp<dpf::prg::aes128, dpf::prg::aes128, true>(
                net, role::p0, x0, x1, beta,
                [&](const auto & key) {
                    const auto y = util::share_bits(*dpf::eval_point(key, alpha));
                    const auto open = util::open_subtractive(net, role::p0, y);
                    EXPECT_EQ(open, beta);
                    const Point other = static_cast<Point>(
                        static_cast<unsigned>(alpha) ^ 1u);
                    const auto z = util::share_bits(*dpf::eval_point(key, other));
                    const auto oz = util::open_subtractive(net, role::p0, z);
                    EXPECT_EQ(oz, 0u);
                },
                [](const auto &) {}, false, additive);
            EXPECT_TRUE(net.dealer_inbox_done());
        },
        [&](trio & net) {
            dist_with_point_key_iknp<dpf::prg::aes128, dpf::prg::aes128, true>(
                net, role::p1, x0, x1, beta,
                [](const auto &) {},
                [&](const auto & key) {
                    const auto y = util::share_bits(*dpf::eval_point(key, alpha));
                    (void)util::open_subtractive(net, role::p1, y);
                    const Point other = static_cast<Point>(
                        static_cast<unsigned>(alpha) ^ 1u);
                    const auto z = util::share_bits(*dpf::eval_point(key, other));
                    (void)util::open_subtractive(net, role::p1, z);
                }, false, additive);
            EXPECT_TRUE(net.dealer_inbox_done());
        });
}

TEST(iknp, point_reveal)
{
    const std::uint8_t alpha = 0x2a;
    expect_point<std::uint8_t>(alpha, 0x10, static_cast<std::uint8_t>(alpha ^ 0x10),
        7, false);
}

TEST(iknp, point_additive)
{
    const std::uint8_t alpha = 0x2a;
    const std::uint8_t x0 = 0x10;
    const std::uint8_t x1 = static_cast<std::uint8_t>(alpha - x0);
    expect_point<std::uint8_t>(alpha, x0, x1, 11, true);
}

TEST(iknp, half_tree_reveal)
{
    using namespace dpf::party;
    using Ht = dpf::prg::aes128_ccr;
    const std::uint8_t alpha = 7;
    const std::uint8_t x0 = 0x25;
    const std::uint8_t x1 = static_cast<std::uint8_t>(alpha ^ x0);
    const std::uint64_t beta = 9;
    run_pair(
        [&](trio & net) {
            dist_with_point_key_iknp<Ht, Ht, true>(net, role::p0, x0, x1, beta,
                [&](const auto & key) {
                    const auto open = util::open_subtractive(net, role::p0,
                        util::share_bits(*dpf::eval_point(key, alpha)));
                    EXPECT_EQ(open, beta);
                },
                [](const auto &) {});
            EXPECT_TRUE(net.dealer_inbox_done());
        },
        [&](trio & net) {
            dist_with_point_key_iknp<Ht, Ht, true>(net, role::p1, x0, x1, beta,
                [](const auto &) {},
                [&](const auto & key) {
                    (void)util::open_subtractive(net, role::p1,
                        util::share_bits(*dpf::eval_point(key, alpha)));
                });
            EXPECT_TRUE(net.dealer_inbox_done());
        });
}

TEST(iknp, comparison_reveal_and_oblivious)
{
    using namespace dpf::party;
    const std::uint8_t alpha = 0x40;
    const std::uint8_t x0 = 0x14;
    const std::uint8_t x1 = static_cast<std::uint8_t>(alpha ^ x0);
    const std::uint64_t beta = 9;
    auto check = [&](bool reveal) {
        run_pair(
            [&](trio & net) {
                auto on = [&](const auto & key) {
                    const auto mask = key.cmp().mask;
                    const auto below = util::open_additive(net, role::p0,
                        util::share_bits(dpf::eval_point(dpf::cmp, key,
                            static_cast<std::uint8_t>(alpha - 1)))) & mask;
                    const auto at = util::open_additive(net, role::p0,
                        util::share_bits(dpf::eval_point(dpf::cmp, key, alpha)))
                        & mask;
                    EXPECT_EQ(below, beta);
                    EXPECT_EQ(at, 0u);
                };
                if (reveal)
                    dist_with_cmp_key_iknp<dpf::prg::aes128, dpf::prg::aes128, true>(
                        net, role::p0, x0, x1, dpf::lt(beta), on, [](const auto &) {});
                else
                    dist_with_cmp_key_iknp(net, role::p0, x0, x1, dpf::lt(beta),
                        on, [](const auto &) {});
                EXPECT_TRUE(net.dealer_inbox_done());
            },
            [&](trio & net) {
                auto on = [&](const auto & key) {
                    const auto mask = key.cmp().mask;
                    (void)(util::open_additive(net, role::p1,
                        util::share_bits(dpf::eval_point(dpf::cmp, key,
                            static_cast<std::uint8_t>(alpha - 1)))) & mask);
                    (void)(util::open_additive(net, role::p1,
                        util::share_bits(dpf::eval_point(dpf::cmp, key, alpha)))
                        & mask);
                };
                if (reveal)
                    dist_with_cmp_key_iknp<dpf::prg::aes128, dpf::prg::aes128, true>(
                        net, role::p1, x0, x1, dpf::lt(beta), [](const auto &) {}, on);
                else
                    dist_with_cmp_key_iknp(net, role::p1, x0, x1, dpf::lt(beta),
                        [](const auto &) {}, on);
                EXPECT_TRUE(net.dealer_inbox_done());
            });
    };
    check(true);
    check(false);
}

TEST(iknp, interval_shared_and_reveal)
{
    using namespace dpf::party;
    const std::uint8_t r = 40, p = 7, q = 90;
    const std::uint8_t r0 = 0x13;
    const std::uint8_t r1 = static_cast<std::uint8_t>(r ^ r0);
    const std::uint32_t if_true = 11, if_false = 2;
    auto check = [&](bool reveal) {
        run_pair(
            [&](trio & net) {
                auto on = [&](const auto & key) {
                    const auto gmask = key.group_mask;
                    const auto nmask = key.input_mask;
                    for (unsigned x : {0u, 7u, 40u, 90u, 200u})
                    {
                        const auto mine = util::share_bits(
                            dpf::eval_point(dpf::ic, key, static_cast<std::uint8_t>(x)));
                        const auto peer = net.exchange_with(role::p1, mine);
                        const auto w = (x - static_cast<unsigned>(r)) & nmask;
                        const bool inside = w >= p && w <= q;
                        const auto want = (inside ? if_true : if_false) & gmask;
                        EXPECT_EQ((mine + peer) & gmask, want);
                    }
                };
                if (reveal)
                    dist_with_ic_key_iknp<dpf::prg::aes128, dpf::prg::aes128, true>(
                        net, role::p0, r0, r1, dpf::ic(p, q, if_true, if_false),
                        on, [](const auto &) {});
                else
                    dist_with_ic_key_iknp(net, role::p0, r0, r1,
                        dpf::ic(p, q, if_true, if_false), on, [](const auto &) {});
                EXPECT_TRUE(net.dealer_inbox_done());
            },
            [&](trio & net) {
                auto on = [&](const auto & key) {
                    for (unsigned x : {0u, 7u, 40u, 90u, 200u})
                    {
                        const auto mine = util::share_bits(
                            dpf::eval_point(dpf::ic, key, static_cast<std::uint8_t>(x)));
                        (void)net.exchange_with(role::p0, mine);
                    }
                };
                if (reveal)
                    dist_with_ic_key_iknp<dpf::prg::aes128, dpf::prg::aes128, true>(
                        net, role::p1, r0, r1, dpf::ic(p, q, if_true, if_false),
                        [](const auto &) {}, on);
                else
                    dist_with_ic_key_iknp(net, role::p1, r0, r1,
                        dpf::ic(p, q, if_true, if_false), [](const auto &) {}, on);
                EXPECT_TRUE(net.dealer_inbox_done());
            });
    };
    check(false);
    check(true);
}

template <typename Key, typename Share>
void assign_wildcard(Key & key, Share my_share, trio & net, role self)
{
    const role peer = self == role::p0 ? role::p1 : role::p0;
    auto & wrap = std::get<0>(key.leaf_nodes);
    if (wrap.is_ready())
        wrap.begin_update();
    auto blinded = wrap.compute_and_get_blinded_output_share(my_share);
    auto peer_blinded = net.exchange_with(peer, blinded);
    auto leaf = wrap.compute_and_get_leaf_share(peer_blinded);
    auto peer_leaf = net.exchange_with(peer, leaf);
    wrap.reconstruct_correction_word(peer_leaf);
}

TEST(iknp, wildcard_assign)
{
    using namespace dpf::party::util;
    const std::uint8_t alpha = 0xAA;
    const std::uint8_t x0 = 0x31;
    const std::uint8_t x1 = static_cast<std::uint8_t>(alpha ^ x0);
    const std::uint32_t y = 0xAAAAAAAAu;
    const std::uint32_t y0 = 0x12345678u;
    const std::uint32_t y1 = y - y0;
    using out_t = dpf::wildcard_value<std::uint32_t>;
    run_pair(
        [&](trio & net) {
            dpf::party::dist_with_point_key_iknp(net, role::p0, x0, x1, out_t{},
                [&](auto key) {
                    assign_wildcard(key, y0, net, role::p0);
                    const auto on = open_subtractive(net, role::p0,
                        share_bits(*dpf::eval_point(key, alpha)));
                    const auto off = open_subtractive(net, role::p0,
                        share_bits(*dpf::eval_point(key,
                            static_cast<std::uint8_t>(alpha ^ 1u))));
                    EXPECT_EQ(on, y);
                    EXPECT_EQ(off, 0u);
                },
                [](auto) {});
            EXPECT_TRUE(net.dealer_inbox_done());
        },
        [&](trio & net) {
            dpf::party::dist_with_point_key_iknp(net, role::p1, x0, x1, out_t{},
                [](auto) {},
                [&](auto key) {
                    assign_wildcard(key, y1, net, role::p1);
                    (void)open_subtractive(net, role::p1,
                        share_bits(*dpf::eval_point(key, alpha)));
                    (void)open_subtractive(net, role::p1,
                        share_bits(*dpf::eval_point(key,
                            static_cast<std::uint8_t>(alpha ^ 1u))));
                });
            EXPECT_TRUE(net.dealer_inbox_done());
        });
}

TEST(iknp, extractable_sketch)
{
    const std::uint8_t alpha = 0x2a;
    const std::uint8_t x0 = 0x10;
    const std::uint8_t x1 = static_cast<std::uint8_t>(alpha ^ x0);
    const dpf::fp61 beta{7};
    run_pair(
        [&](trio & net) {
            dpf::party::dist_with_extractable_point_key_iknp(
                net, role::p0, x0, x1, beta,
                [&](const auto & key) {
                    const std::array<dpf::fp61, 1> rs{dpf::fp61{3}};
                    dpf::sketch_share local{};
                    auto sk = dpf::sketch(local, rs);
                    (void)*dpf::eval_point(key, alpha, sk);
                    auto theirs = net.exchange_with(role::p1, local,
                        dpf::net::msg::sketch_share);
                    EXPECT_TRUE(dpf::sketch_verify(local, theirs));
                },
                [](const auto &) {});
            EXPECT_TRUE(net.dealer_inbox_done());
        },
        [&](trio & net) {
            dpf::party::dist_with_extractable_point_key_iknp(
                net, role::p1, x0, x1, beta,
                [](const auto &) {},
                [&](const auto & key) {
                    const std::array<dpf::fp61, 1> rs{dpf::fp61{3}};
                    dpf::sketch_share local{};
                    auto sk = dpf::sketch(local, rs);
                    (void)*dpf::eval_point(key, alpha, sk);
                    auto theirs = net.exchange_with(role::p0, local,
                        dpf::net::msg::sketch_share);
                    EXPECT_TRUE(dpf::sketch_verify(theirs, local));
                });
            EXPECT_TRUE(net.dealer_inbox_done());
        });
}

TEST(iknp, oblivious_point_hash)
{
    using namespace dpf::party;
    // Default (non-reveal) point keygen hashes the shared prefix, so this
    // exercises the per-level AND tape as well as the correction words.
    const std::uint8_t alpha = 0x2a;
    const std::uint8_t x0 = 0x10;
    const std::uint8_t x1 = static_cast<std::uint8_t>(alpha ^ x0);
    const std::uint64_t beta = 4;
    run_pair(
        [&](trio & net) {
            dist_with_point_key_iknp(net, role::p0, x0, x1, beta,
                [&](const auto & key) {
                    const auto on = util::open_subtractive(net, role::p0,
                        util::share_bits(*dpf::eval_point(key, alpha)));
                    const auto off = util::open_subtractive(net, role::p0,
                        util::share_bits(*dpf::eval_point(key,
                            static_cast<std::uint8_t>(alpha ^ 1u))));
                    EXPECT_EQ(on, beta);
                    EXPECT_EQ(off, 0u);
                },
                [](const auto &) {});
            EXPECT_TRUE(net.dealer_inbox_done());
        },
        [&](trio & net) {
            dist_with_point_key_iknp(net, role::p1, x0, x1, beta,
                [](const auto &) {},
                [&](const auto & key) {
                    (void)util::open_subtractive(net, role::p1,
                        util::share_bits(*dpf::eval_point(key, alpha)));
                    (void)util::open_subtractive(net, role::p1,
                        util::share_bits(*dpf::eval_point(key,
                            static_cast<std::uint8_t>(alpha ^ 1u))));
                });
            EXPECT_TRUE(net.dealer_inbox_done());
        });
}

TEST(iknp, transfer_labels_follows_choice_and_reuses_base)
{
    constexpr std::size_t n = 7;
    std::vector<simde__m128i> m0(n), m1(n);
    std::vector<std::uint8_t> choices(n);
    for (std::size_t i = 0; i < n; ++i)
    {
        m0[i] = dpf::uniform_sample<simde__m128i>();
        m1[i] = dpf::uniform_sample<simde__m128i>();
        choices[i] = static_cast<std::uint8_t>(i & 1u);
    }
    std::vector<simde__m128i> again0(2), again1(2);
    again0[0] = dpf::uniform_sample<simde__m128i>();
    again0[1] = dpf::uniform_sample<simde__m128i>();
    again1[0] = dpf::uniform_sample<simde__m128i>();
    again1[1] = dpf::uniform_sample<simde__m128i>();
    const std::uint8_t again_choice[2] = {1, 0};

    std::vector<simde__m128i> got, got2;
    run_pair(
        [&](trio & net) {
            dpf::iknp::detail::role_state st;
            std::vector<simde__m128i> unused;
            dpf::iknp::transfer_labels(net.to(role::p1), 0, true, st, m0, m1,
                {}, unused);
            EXPECT_TRUE(unused.empty());
            std::vector<simde__m128i> empty_out;
            dpf::iknp::transfer_labels(net.to(role::p1), 0, true, st, {}, {},
                {}, empty_out);
            dpf::iknp::transfer_labels(net.to(role::p1), 0, true, st, again0,
                again1, {}, unused);
        },
        [&](trio & net) {
            dpf::iknp::detail::role_state st;
            dpf::iknp::transfer_labels(net.to(role::p0), 1, false, st, {}, {},
                choices, got);
            std::vector<std::uint8_t> none;
            std::vector<simde__m128i> empty_out;
            dpf::iknp::transfer_labels(net.to(role::p0), 1, false, st, {}, {},
                none, empty_out);
            EXPECT_TRUE(empty_out.empty());
            dpf::iknp::transfer_labels(net.to(role::p0), 1, false, st, {}, {},
                std::vector<std::uint8_t>(again_choice, again_choice + 2), got2);
        });

    ASSERT_EQ(got.size(), n);
    for (std::size_t i = 0; i < n; ++i)
    {
        const auto & expect = choices[i] ? m1[i] : m0[i];
        EXPECT_EQ(std::memcmp(&got[i], &expect, sizeof(expect)), 0) << i;
    }
    ASSERT_EQ(got2.size(), 2u);
    EXPECT_EQ(std::memcmp(&got2[0], &again1[0], sizeof(again1[0])), 0);
    EXPECT_EQ(std::memcmp(&got2[1], &again0[1], sizeof(again0[1])), 0);
}

}  // namespace