#include <gtest/gtest.h>

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

#include "dpf/compose.hpp"
#include "dpf/shuffle.hpp"

namespace
{

template <typename T>
std::vector<T> replay(const std::vector<T> & v, const dpf::rss::seed_bundle & bundle,
    std::uint64_t index, unsigned passes)
{
    const dpf::rss::seed_block seeds[3] = {bundle.k01, bundle.k12, bundle.k20};
    auto out = v;
    for (unsigned p = 0; p < passes; ++p)
    {
        out = dpf::shuffle::permute(out,
            dpf::shuffle::permutation_from_seed(seeds[p], v.size(), index));
    }
    return out;
}

template <typename T>
void deal(const std::vector<T> & clear, dpf::shuffle::shuffle_party_view<T> held[3])
{
    const std::size_t n = clear.size();
    for (unsigned p = 0; p < 3; ++p)
    {
        held[p].own.assign(n, T{});
        held[p].next.assign(n, T{});
    }
    for (std::size_t i = 0; i < n; ++i)
    {
        const T a = dpf::uniform_sample<T>();
        const T b = dpf::uniform_sample<T>();
        const T c = static_cast<T>(clear[i] - a - b);
        held[0].own[i] = a;
        held[0].next[i] = b;
        held[1].own[i] = b;
        held[1].next[i] = c;
        held[2].own[i] = c;
        held[2].next[i] = a;
    }
}

template <typename T>
std::vector<T> open_owns(const dpf::shuffle::shuffle_party_view<T> held[3])
{
    std::vector<T> out(held[0].own.size());
    for (std::size_t i = 0; i < out.size(); ++i)
        out[i] = static_cast<T>(held[0].own[i] + held[1].own[i] + held[2].own[i]);
    return out;
}

template <typename T>
void run_passes(dpf::shuffle::shuffle_party_view<T> held[3],
    const dpf::rss::seed_bundle & bundle, std::uint64_t index, unsigned passes)
{
    const unsigned order[3] = {2u, 0u, 1u};
    for (unsigned step = 0; step < passes; ++step)
    {
        const unsigned left = order[step];
        const unsigned u = dpf::shuffle::hidden_u_party(left);
        const unsigned side = dpf::shuffle::hidden_side_party(left);
        auto u_step = dpf::shuffle::shuffle_hidden_pass<T>(
            u, dpf::rss::party_seeds::from_bundle(bundle, u), held[u], index, left,
            nullptr);
        auto s_step = dpf::shuffle::shuffle_hidden_pass<T>(
            side, dpf::rss::party_seeds::from_bundle(bundle, side), held[side],
            index, left, &u_step.out.data);
        auto l_step = dpf::shuffle::shuffle_hidden_pass<T>(
            left, dpf::rss::party_seeds::from_bundle(bundle, left), held[left],
            index, left, &s_step.out.data);
        EXPECT_EQ(u_step.out.data.size(), held[u].own.size());
        EXPECT_EQ(s_step.out.to, left);
        EXPECT_EQ(held[u].own.size(), l_step.view.own.size());
        held[u] = std::move(u_step.view);
        held[side] = std::move(s_step.view);
        held[left] = std::move(l_step.view);
        ASSERT_EQ(held[0].next, held[1].own);
        ASSERT_EQ(held[1].next, held[2].own);
        ASSERT_EQ(held[2].next, held[0].own);
    }
}

}  // namespace

TEST(ShuffleHidden, ThreePassesMatchTheSeedPermutations)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    std::vector<std::uint64_t> v(8);
    std::iota(v.begin(), v.end(), 0);
    const std::uint64_t index = 4;
    auto opened = dpf::shuffle::shuffle_hidden_triple(v, bundle, index);
    auto expect = dpf::shuffle::permute(v,
        dpf::shuffle::permutation_from_seed(bundle.k01, v.size(), index));
    expect = dpf::shuffle::permute(expect,
        dpf::shuffle::permutation_from_seed(bundle.k12, v.size(), index));
    expect = dpf::shuffle::permute(expect,
        dpf::shuffle::permutation_from_seed(bundle.k20, v.size(), index));
    EXPECT_EQ(opened, expect);
    EXPECT_TRUE(dpf::shuffle::is_permutation_of(v, opened));
}

TEST(ShuffleHidden, LeftOutPartyUsesOnlyItsTwoSeeds)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    const std::size_t n = 4;
    dpf::shuffle::shuffle_party_view<std::uint64_t> in{
        std::vector<std::uint64_t>(n, 1),
        std::vector<std::uint64_t>(n, 2),
    };
    auto seeds = dpf::rss::party_seeds::from_bundle(bundle, 2);
    std::vector<std::uint64_t> inbound(n, 9);
    auto step = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
        2, seeds, in, 0, 2, &inbound);
    EXPECT_FALSE(step.out.sends);
    EXPECT_EQ(step.view.own, inbound);
    EXPECT_EQ(step.view.next.size(), n);
}

TEST(ShuffleHidden, ReplicationHoldsAfterEachPass)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    const std::size_t n = 5;
    dpf::shuffle::shuffle_party_view<std::uint64_t> held[3];
    for (unsigned p = 0; p < 3; ++p)
    {
        held[p].own.assign(n, 0);
        held[p].next.assign(n, 0);
    }
    for (std::size_t i = 0; i < n; ++i)
    {
        const auto a = dpf::uniform_sample<std::uint64_t>();
        const auto b = dpf::uniform_sample<std::uint64_t>();
        held[0].own[i] = a;
        held[0].next[i] = b;
        held[1].own[i] = b;
        held[1].next[i] = static_cast<std::uint64_t>(i - a - b);
        held[2].own[i] = held[1].next[i];
        held[2].next[i] = a;
    }
    for (unsigned left : {2u, 0u, 1u})
    {
        const unsigned u = dpf::shuffle::hidden_u_party(left);
        const unsigned side = dpf::shuffle::hidden_side_party(left);
        auto u_step = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
            u, dpf::rss::party_seeds::from_bundle(bundle, u), held[u], 1, left,
            nullptr);
        auto s_step = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
            side, dpf::rss::party_seeds::from_bundle(bundle, side), held[side],
            1, left, &u_step.out.data);
        auto l_step = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
            left, dpf::rss::party_seeds::from_bundle(bundle, left), held[left],
            1, left, &s_step.out.data);
        EXPECT_TRUE(u_step.out.sends);
        EXPECT_EQ(u_step.out.to, side);
        EXPECT_TRUE(s_step.out.sends);
        EXPECT_EQ(s_step.out.to, left);
        EXPECT_FALSE(l_step.out.sends);
        held[u] = std::move(u_step.view);
        held[side] = std::move(s_step.view);
        held[left] = std::move(l_step.view);
        EXPECT_EQ(held[0].next, held[1].own);
        EXPECT_EQ(held[1].next, held[2].own);
        EXPECT_EQ(held[2].next, held[0].own);
    }
}

TEST(ShuffleHidden, ComposerRecordsThreeSends)
{
    dpf::protocol::composer c;
    auto out = c.shuffle_hidden(8, sizeof(std::uint64_t));
    auto plan = c.schedule();
    EXPECT_EQ(plan.rounds(), 3u);
    std::vector<std::uint32_t> left_out;
    for (std::size_t w = 0; w < plan.waves(); ++w)
    {
        for (auto ex : plan.wave(w).exchanges)
        {
            EXPECT_EQ(plan.opcode_of(ex.id), dpf::protocol::opcodes::shuffle_send);
            left_out.push_back(plan.aux_of(ex.id));
            EXPECT_EQ(dpf::protocol::detail::exchange_channel(plan, ex.id),
                dpf::protocol::edge_channel::rss_next);
        }
    }
    EXPECT_EQ(left_out, (std::vector<std::uint32_t>{2u, 0u, 1u}));
    EXPECT_EQ(plan.value_bytes_of(out.id), 8u * sizeof(std::uint64_t));
    std::size_t prev = 0;
    bool seen = false;
    for (std::size_t w = 0; w < plan.waves(); ++w)
    {
        if (plan.wave(w).exchanges.empty())
            continue;
        if (seen)
            EXPECT_GT(w, prev);
        prev = w;
        seen = true;
    }
}

TEST(ShuffleHidden, EachPassMatchesOneSeedPermutation)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    std::vector<std::uint64_t> v(9);
    for (std::size_t i = 0; i < v.size(); ++i)
        v[i] = 1000u + static_cast<std::uint64_t>(i * 17u);
    dpf::shuffle::shuffle_party_view<std::uint64_t> held[3];
    deal(v, held);
    for (unsigned passes = 1; passes <= 3; ++passes)
    {
        deal(v, held);
        run_passes(held, bundle, 9, passes);
        EXPECT_EQ(open_owns(held), replay(v, bundle, 9, passes));
    }
}

TEST(ShuffleHidden, LengthsRingsAndHighBits)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    for (std::size_t n : {0u, 1u, 2u, 3u, 7u, 16u, 64u})
    {
        std::vector<std::uint64_t> v(n);
        for (std::size_t i = 0; i < n; ++i)
            v[i] = (i * 0x9E3779B97F4A7C15ull) ^ 0x8000000000000000ull;
        auto opened = dpf::shuffle::shuffle_hidden_triple(v, bundle, 3);
        EXPECT_EQ(opened, replay(v, bundle, 3, 3));
        if (n == 1)
            EXPECT_EQ(opened, v);
    }
    std::vector<std::uint8_t> bytes{0, 255, 1, 128, 7, 7, 255};
    EXPECT_EQ(dpf::shuffle::shuffle_hidden_triple(bytes, bundle, 2),
        replay(bytes, bundle, 2, 3));
    std::vector<std::uint32_t> words{0u, 0xffffffffu, 1u, 0x80000000u};
    EXPECT_EQ(dpf::shuffle::shuffle_hidden_triple(words, bundle, 5),
        replay(words, bundle, 5, 3));
    std::vector<std::uint64_t> same(12, 42);
    EXPECT_EQ(dpf::shuffle::shuffle_hidden_triple(same, bundle, 1), same);
}

TEST(ShuffleHidden, ReplayIsStableAndIndexChangesTheOrder)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    std::vector<std::uint64_t> v(16);
    std::iota(v.begin(), v.end(), 0);
    auto a = dpf::shuffle::shuffle_hidden_triple(v, bundle, 0);
    auto again = dpf::shuffle::shuffle_hidden_triple(v, bundle, 0);
    auto other = dpf::shuffle::shuffle_hidden_triple(v, bundle, 99);
    EXPECT_EQ(a, again);
    EXPECT_EQ(a, replay(v, bundle, 0, 3));
    EXPECT_NE(a, other);
}

TEST(ShuffleHidden, UPartyIgnoresInboundAndABadMessageBreaksTheOpen)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    std::vector<std::uint64_t> v(8);
    std::iota(v.begin(), v.end(), 3);
    dpf::shuffle::shuffle_party_view<std::uint64_t> held[3];
    deal(v, held);
    const unsigned left = 2;
    const unsigned u = dpf::shuffle::hidden_u_party(left);
    const unsigned side = dpf::shuffle::hidden_side_party(left);
    std::vector<std::uint64_t> junk(v.size(), 123);
    auto ignored = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
        u, dpf::rss::party_seeds::from_bundle(bundle, u), held[u], 0, left, &junk);
    auto clean = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
        u, dpf::rss::party_seeds::from_bundle(bundle, u), held[u], 0, left, nullptr);
    EXPECT_EQ(ignored.view.own, clean.view.own);
    EXPECT_EQ(ignored.view.next, clean.view.next);
    junk[0] ^= 1u;
    auto bad = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
        side, dpf::rss::party_seeds::from_bundle(bundle, side), held[side], 0,
        left, &junk);
    auto good = dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
        side, dpf::rss::party_seeds::from_bundle(bundle, side), held[side], 0,
        left, &clean.out.data);
    EXPECT_NE(bad.view.own, good.view.own);
}

TEST(ShuffleHidden, LeftOutMaskMatchesAndMissesThePermutationSeed)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    auto party2 = dpf::rss::party_seeds::from_bundle(bundle, 2);
    EXPECT_EQ(std::memcmp(&party2.with_prev, &bundle.k12, sizeof(bundle.k12)), 0);
    EXPECT_EQ(std::memcmp(&party2.with_next, &bundle.k20, sizeof(bundle.k20)), 0);
    EXPECT_NE(std::memcmp(&party2.with_prev, &bundle.k01, sizeof(bundle.k01)), 0);
    EXPECT_NE(std::memcmp(&party2.with_next, &bundle.k01, sizeof(bundle.k01)), 0);
    const auto pi_known = dpf::shuffle::permutation_from_seed(bundle.k01, 8, 0);
    const auto pi_prev = dpf::shuffle::permutation_from_seed(party2.with_prev, 8, 0);
    const auto pi_next = dpf::shuffle::permutation_from_seed(party2.with_next, 8, 0);
    EXPECT_NE(pi_known, pi_prev);
    EXPECT_NE(pi_known, pi_next);
}

TEST(ShuffleHidden, RejectsBadPartiesAndShortMessages)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    dpf::shuffle::shuffle_party_view<std::uint64_t> in{
        {1, 2},
        {3, 4},
    };
    auto seeds = dpf::rss::party_seeds::from_bundle(bundle, 0);
    EXPECT_THROW(
        dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(1, seeds, in, 0, 2, nullptr),
        std::invalid_argument);
    EXPECT_THROW(
        dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(0, seeds, in, 0, 3, nullptr),
        std::invalid_argument);
    in.next.pop_back();
    EXPECT_THROW(
        dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(0, seeds, in, 0, 2, nullptr),
        std::invalid_argument);
    in.next.push_back(4);
    auto side = dpf::rss::party_seeds::from_bundle(bundle, 1);
    EXPECT_THROW(
        dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(1, side, in, 0, 2, nullptr),
        std::invalid_argument);
    std::vector<std::uint64_t> short_msg{1};
    EXPECT_THROW(dpf::shuffle::shuffle_hidden_pass<std::uint64_t>(
                     1, side, in, 0, 2, &short_msg),
        std::invalid_argument);
    EXPECT_THROW(dpf::shuffle::permute(std::vector<std::uint64_t>{1},
                     std::vector<std::size_t>{}),
        std::invalid_argument);
}

TEST(ShuffleHidden, PublicShuffleStillFollowsK01)
{
    auto bundle = dpf::rss::sample_seed_bundle();
    std::vector<std::uint64_t> v{4, 9, 1, 7, 3};
    EXPECT_EQ(dpf::shuffle::shuffle_party_triple(v, bundle, 6),
        dpf::shuffle::shuffle_clear(v, bundle, 6));
}

TEST(ShuffleHidden, ComposerRejectsAnEmptyColumn)
{
    dpf::protocol::composer c;
    EXPECT_THROW(c.shuffle_hidden(0, 8), std::invalid_argument);
    EXPECT_THROW(c.shuffle_hidden(4, 0), std::invalid_argument);
}