#include <gtest/gtest.h>

#include "dpf.hpp"

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstring>
#include <set>
#include <stdexcept>
#include <vector>

namespace
{

using block = dpf::prg::aes128::block_type;

block seed_block(std::uint64_t lo, std::uint64_t hi)
{
    alignas(16) std::uint64_t lanes[2] = {lo, hi};
    return simde_mm_load_si128(reinterpret_cast<const simde__m128i *>(lanes));
}

bool same_root(const block & a, const block & b)
{
    alignas(16) unsigned char aa[16], bb[16];
    simde_mm_store_si128(reinterpret_cast<simde__m128i *>(aa), a);
    simde_mm_store_si128(reinterpret_cast<simde__m128i *>(bb), b);
    return std::memcmp(aa, bb, 16) == 0;
}

template <typename Scheme>
void expect_point_programming(const typename Scheme::public_params & pp,
    const typename Scheme::commitment & com, const typename Scheme::state & st,
    typename Scheme::input_type xi)
{
    ASSERT_TRUE(Scheme::check_well_formed(pp, com, st));
    const auto baseline = Scheme::eval(Scheme::open(st, 0, 0, xi));
    std::set<std::uint64_t> seen;
    for (std::uint64_t tau = 0; tau <= Scheme::value_mask(); ++tau)
    {
        const auto op = Scheme::open(st, 0, tau, xi);
        const auto x = Scheme::eval(op);
        const auto rotated = Scheme::eval_rotated(op);

        EXPECT_TRUE(Scheme::accept(pp, com, op, x, xi));
        EXPECT_EQ(op.delta, Scheme::sub(xi, st.i));
        EXPECT_EQ(x[Scheme::index_of(st.i)], tau);
        EXPECT_EQ(rotated[Scheme::index_of(xi)], tau);
        EXPECT_EQ(Scheme::column_sum(x), Scheme::column_sum(rotated));
        for (std::size_t y = 0; y < Scheme::domain_size; ++y)
        {
            if (y == Scheme::index_of(st.i))
                continue;
            EXPECT_EQ(x[y], baseline[y]);
        }

        std::uint64_t pack = 0;
        const auto bits = Scheme::disclosure(op);
        for (std::size_t j = 0; j < Scheme::width; ++j)
            if (bits[j])
                pack |= std::uint64_t{1} << j;
        EXPECT_TRUE(seen.insert(pack).second);

        for (std::size_t j = 0; j < Scheme::width; ++j)
        {
            const bool side0 = same_root(op.keys[j]->root(), st.keys[j][0]->root());
            const bool side1 = same_root(op.keys[j]->root(), st.keys[j][1]->root());
            EXPECT_NE(side0, side1);
            EXPECT_TRUE(same_root(op.coins[j], side1 ? st.coins[j][1] : st.coins[j][0]));
        }
    }
    EXPECT_EQ(seen.size(), Scheme::value_mask() + 1);

    const auto point_open = Scheme::open(st, 0, 1, xi);
    const auto point_vector = Scheme::eval(point_open);
    const auto sum = Scheme::column_sum(point_vector);
    const auto sum_open = Scheme::open(st, 1, sum, xi);
    EXPECT_EQ(point_vector, Scheme::eval(sum_open));
    EXPECT_TRUE(Scheme::accept(pp, com, sum_open, point_vector, xi));
}

template <typename Scheme>
void expect_sum_programming(const typename Scheme::public_params & pp,
    const typename Scheme::commitment & com, const typename Scheme::state & st,
    typename Scheme::input_type xi)
{
    const auto baseline = Scheme::eval(Scheme::open(st, 0, 0, xi));
    const auto hidden = Scheme::index_of(st.i);
    std::uint64_t previous = 0;
    bool have_previous = false;
    for (std::uint64_t tau = 0; tau <= Scheme::value_mask(); ++tau)
    {
        const auto op = Scheme::open(st, 1, tau, xi);
        const auto x = Scheme::eval(op);
        EXPECT_TRUE(Scheme::accept(pp, com, op, x, xi));
        EXPECT_EQ(Scheme::column_sum(x), tau);
        EXPECT_EQ(Scheme::column_sum(Scheme::eval_rotated(op)), tau);
        for (std::size_t y = 0; y < Scheme::domain_size; ++y)
        {
            if (y == hidden)
                continue;
            EXPECT_EQ(x[y], baseline[y]);
        }
        if (have_previous)
        {
            EXPECT_EQ((x[hidden] - previous) & Scheme::value_mask(),
                (tau - (tau - 1)) & Scheme::value_mask());
        }
        previous = x[hidden];
        have_previous = true;
    }
}

}  // namespace

TEST(Ppvc, ProgramsEveryCoordinate)
{
    using narrow = dpf::ppvc<std::uint8_t, 1, 8>;
    using wide = dpf::ppvc<std::uint8_t, 4, 8>;
    const auto pp1 = narrow::setup_from_seed(seed_block(7, 9));
    const auto made1 = narrow::commit_from_seed(pp1, seed_block(11, 13));
    expect_point_programming<narrow>(pp1, made1.first, made1.second, narrow::point(40));

    const auto pp4 = wide::setup_from_seed(seed_block(7, 9));
    const auto made4 = wide::commit_from_seed(pp4, seed_block(11, 13));
    expect_point_programming<wide>(pp4, made4.first, made4.second, wide::point(40));
}

TEST(Ppvc, ProgramsEverySum)
{
    using narrow = dpf::ppvc<std::uint8_t, 1, 8>;
    using wide = dpf::ppvc<std::uint8_t, 4, 8>;
    const auto pp1 = narrow::setup_from_seed(seed_block(7, 9));
    const auto made1 = narrow::commit_from_seed(pp1, seed_block(11, 13));
    expect_sum_programming<narrow>(pp1, made1.first, made1.second, narrow::point(40));

    const auto pp4 = wide::setup_from_seed(seed_block(7, 9));
    const auto made4 = wide::commit_from_seed(pp4, seed_block(11, 13));
    expect_sum_programming<wide>(pp4, made4.first, made4.second, wide::point(40));
}

TEST(Ppvc, FreshCommitOpens)
{
    using scheme = dpf::ppvc<std::uint8_t, 8, 8>;
    const auto pp = scheme::setup();
    const auto [com, st] = scheme::commit(pp);
    const auto xi = scheme::point(200);
    const auto op = scheme::open(st, 0, 0x5a, xi);
    const auto x = scheme::eval(op);
    EXPECT_TRUE(scheme::accept(pp, com, op, x, xi));
    EXPECT_EQ(x[scheme::index_of(st.i)], 0x5au);
    EXPECT_EQ(scheme::eval_rotated(op)[200], 0x5au);
}

TEST(Ppvc, PrescribedIndexAndWrappingShift)
{
    using scheme = dpf::ppvc<std::uint8_t, 4, 8>;
    const auto pp = scheme::setup_from_seed(seed_block(8, 8));
    for (std::size_t y : {std::size_t{0}, std::size_t{1}, std::size_t{255}})
        EXPECT_EQ(scheme::index_of(scheme::point(y)), y);

    const auto hidden = scheme::point(255);
    const auto [com, st] = scheme::commit_at(pp, hidden);
    EXPECT_EQ(st.i, hidden);
    ASSERT_TRUE(scheme::check_well_formed(pp, com, st));

    const auto wrapped = scheme::open(st, 0, 7, scheme::point(0));
    EXPECT_EQ(scheme::index_of(wrapped.delta), 1u);
    EXPECT_EQ(scheme::eval(wrapped)[255], 7u);
    EXPECT_EQ(scheme::eval_rotated(wrapped)[0], 7u);
    EXPECT_TRUE(scheme::accept(pp, com, wrapped, scheme::eval(wrapped), scheme::point(0)));

    const auto unshifted = scheme::open(st, 0, 7, hidden);
    EXPECT_EQ(scheme::index_of(unshifted.delta), 0u);
    EXPECT_EQ(scheme::eval_rotated(unshifted)[255], 7u);
}

TEST(Ppvc, SeedExpansionIsDeterministic)
{
    using scheme = dpf::ppvc<std::uint8_t, 4, 8>;
    const auto seed = seed_block(3, 4);
    const auto pp = scheme::setup_from_seed(seed);
    EXPECT_EQ(pp.columns, scheme::setup_from_seed(seed).columns);
    EXPECT_NE(pp.columns, scheme::setup_from_seed(seed_block(3, 5)).columns);

    const auto replica = seed_block(5, 6);
    const auto first = scheme::commit_from_seed(pp, replica);
    const auto second = scheme::commit_from_seed(pp, replica);
    EXPECT_EQ(first.first, second.first);
    EXPECT_EQ(first.second.i, second.second.i);
    EXPECT_TRUE(same_root(first.second.keys[0][0]->root(), second.second.keys[0][0]->root()));
    EXPECT_EQ(0, std::memcmp(first.second.keys[0][0]->correction_words().data(),
                      second.second.keys[0][0]->correction_words().data(),
                      sizeof(first.second.keys[0][0]->correction_words())));
    EXPECT_TRUE(scheme::audit(pp, first.first, replica));
    EXPECT_FALSE(scheme::audit(pp, first.first, seed_block(5, 7)));

    const auto fixed = scheme::point(3);
    const auto at = scheme::commit_at_from_seed(pp, replica, fixed);
    const auto at_again = scheme::commit_at_from_seed(pp, replica, fixed);
    EXPECT_EQ(at.first, at_again.first);
    EXPECT_EQ(at.second.i, fixed);
    EXPECT_TRUE(scheme::check_well_formed(pp, at.first, at.second));
}

TEST(Ppvc, RejectsATamperedOpening)
{
    using scheme = dpf::ppvc<std::uint8_t, 4, 8>;
    const auto pp = scheme::setup_from_seed(seed_block(3, 4));
    auto [com, st] = scheme::commit_from_seed(pp, seed_block(5, 6));
    const auto xi = scheme::point(9);
    auto op = scheme::open(st, 0, 3, xi);
    ASSERT_TRUE(scheme::verify(pp, com, op));

    alignas(16) std::uint8_t raw[16];
    simde_mm_store_si128(reinterpret_cast<simde__m128i *>(raw), op.coins[0]);
    raw[0] = static_cast<std::uint8_t>(raw[0] ^ 0x1u);
    op.coins[0] = simde_mm_load_si128(reinterpret_cast<const simde__m128i *>(raw));
    EXPECT_FALSE(scheme::verify(pp, com, op));

    op = scheme::open(st, 0, 3, xi);
    const unsigned side = scheme::disclosure(op)[0] ? 1u : 0u;
    op.coins[0] = st.coins[0][side ^ 1u];
    EXPECT_FALSE(scheme::verify(pp, com, op));

    op = scheme::open(st, 0, 3, xi);
    com.slots[0][side].bytes[0] =
        static_cast<std::uint8_t>(com.slots[0][side].bytes[0] ^ 0x1u);
    EXPECT_FALSE(scheme::verify(pp, com, op));

    op.keys[0].reset();
    EXPECT_FALSE(scheme::verify(pp, com, op));
    EXPECT_THROW(scheme::eval(op), std::invalid_argument);
    EXPECT_THROW(scheme::disclosure(op), std::invalid_argument);

    op = scheme::open(st, 0, 3, xi);
    const auto x = scheme::eval(op);
    EXPECT_FALSE(scheme::check_statement(op, std::vector<scheme::value_type>(3), xi));
    auto wrong = x;
    std::fill(wrong.begin(), wrong.end(), scheme::value_type{0});
    EXPECT_FALSE(scheme::check_statement(op, wrong, xi));
    op.mu = 3;
    EXPECT_FALSE(scheme::check_statement(op, x, xi));
}

TEST(Ppvc, RejectsAMalformedReplica)
{
    using scheme = dpf::ppvc<std::uint8_t, 4, 8>;
    const auto pp = scheme::setup_from_seed(seed_block(21, 22));
    auto [com, st] = scheme::commit_from_seed(pp, seed_block(23, 24));
    ASSERT_TRUE(scheme::check_well_formed(pp, com, st));
    EXPECT_EQ(0, std::memcmp(st.keys[0][0]->correction_words().data(),
                      st.keys[0][1]->correction_words().data(),
                      sizeof(st.keys[0][0]->correction_words())));
    EXPECT_EQ(st.keys[0][0]->correction_advice(), st.keys[0][1]->correction_advice());

    const auto moved = scheme::point(scheme::index_of(st.i) + 1);
    auto fresh = dpf::make_dpf(moved, dpf::bit::one);
    st.keys[0][0] = fresh.first.key();
    st.keys[0][1] = fresh.second.key();
    st.coins[0][0] = dpf::uniform_sample<block>();
    st.coins[0][1] = dpf::uniform_sample<block>();
    com.slots[0][0] = scheme::commit_root(pp, st.keys[0][0]->root(), st.coins[0][0]);
    com.slots[0][1] = scheme::commit_root(pp, st.keys[0][1]->root(), st.coins[0][1]);
    EXPECT_FALSE(scheme::check_well_formed(pp, com, st));

    auto [com2, st2] = scheme::commit_from_seed(pp, seed_block(23, 24));
    auto replacement = dpf::make_dpf(st2.i, dpf::bit::one);
    st2.keys[1][0] = replacement.first.key();
    st2.coins[1][0] = dpf::uniform_sample<block>();
    com2.slots[1][0] = scheme::commit_root(pp, st2.keys[1][0]->root(), st2.coins[1][0]);
    EXPECT_FALSE(scheme::check_well_formed(pp, com2, st2));

    EXPECT_FALSE(scheme::check_well_formed(pp, scheme::commitment{}, scheme::state{}));
}

TEST(Ppvc, RejectsAnOutOfRangeProgram)
{
    using scheme = dpf::ppvc<std::uint8_t, 4, 8>;
    const auto pp = scheme::setup_from_seed(seed_block(1, 1));
    const auto [com, st] = scheme::commit_from_seed(pp, seed_block(2, 2));
    (void)com;
    EXPECT_THROW(scheme::open(st, 2, 0, scheme::point(0)), std::invalid_argument);
    EXPECT_THROW(scheme::open(st, 0, 16, scheme::point(0)), std::invalid_argument);
    EXPECT_THROW(scheme::open(st, 1, scheme::value_mask() + 1, scheme::point(0)),
        std::invalid_argument);
}

TEST(Ppvc, DeeperDomain)
{
    using scheme = dpf::ppvc<std::uint16_t, 2, 8>;
    EXPECT_GT(scheme::domain_bits, 8u);
    const auto pp = scheme::setup_from_seed(seed_block(40, 41));
    const auto hidden = scheme::point(1000);
    const auto [com, st] = scheme::commit_at_from_seed(pp, seed_block(42, 43), hidden);
    ASSERT_TRUE(scheme::check_well_formed(pp, com, st));
    const auto xi = scheme::point(7);
    const auto op = scheme::open(st, 0, 3, xi);
    const auto x = scheme::eval(op);
    EXPECT_TRUE(scheme::accept(pp, com, op, x, xi));
    EXPECT_EQ(scheme::eval_rotated(op)[7], 3u);

    const auto summed = scheme::open(st, 1, 1, xi);
    EXPECT_EQ(scheme::column_sum(scheme::eval(summed)), 1u);
    EXPECT_TRUE(scheme::audit(pp, scheme::commit_from_seed(pp, seed_block(42, 43)).first,
        seed_block(42, 43)));
}

TEST(Ppvc, FullWidthValue)
{
    using scheme = dpf::ppvc<std::uint8_t, 64, 8>;
    EXPECT_EQ(scheme::value_mask(), ~std::uint64_t{0});
    const auto pp = scheme::setup_from_seed(seed_block(50, 51));
    const auto [com, st] = scheme::commit_from_seed(pp, seed_block(52, 53));
    const auto xi = scheme::point(12);
    for (std::uint64_t tau : {std::uint64_t{0}, std::uint64_t{1},
             std::uint64_t{1} << 63, ~std::uint64_t{0}})
    {
        const auto point = scheme::open(st, 0, tau, xi);
        const auto x = scheme::eval(point);
        EXPECT_TRUE(scheme::accept(pp, com, point, x, xi));
        EXPECT_EQ(x[scheme::index_of(st.i)], tau);

        const auto summed = scheme::open(st, 1, tau, xi);
        EXPECT_EQ(scheme::column_sum(scheme::eval(summed)), tau);
    }
}

TEST(Ppvc, DefaultSigmaUsesWholeBlocks)
{
    using scheme = dpf::ppvc<std::uint8_t, 2>;
    EXPECT_EQ(scheme::sigma, 128u);
    EXPECT_EQ(scheme::m_bits % 128, 0u);
    EXPECT_EQ(scheme::commitment_bits, 2u * scheme::width * scheme::m_bits);
    const auto pp = scheme::setup_from_seed(seed_block(60, 61));
    const auto [com, st] = scheme::commit_from_seed(pp, seed_block(62, 63));
    const auto xi = scheme::point(4);
    const auto op = scheme::open(st, 0, 2, xi);
    EXPECT_TRUE(scheme::accept(pp, com, op, scheme::eval(op), xi));
    EXPECT_TRUE(scheme::audit(pp, com, seed_block(62, 63)));
}

TEST(KPpvc, OneCopyMatchesTheSingleScheme)
{
    using scheme = dpf::k_ppvc<1, std::uint8_t, 4, 8>;
    const auto pp = scheme::one::setup_from_seed(seed_block(70, 71));
    const auto [com, st] = scheme::commit_from_seed(pp, seed_block(72, 73));
    ASSERT_TRUE(scheme::check_well_formed(pp, com, st));
    const std::array<std::uint8_t, 1> xi{scheme::one::point(15)};
    const std::array<scheme::value_type, 1> tau{6};
    const auto op = scheme::open(st, 0, tau, xi);
    EXPECT_TRUE(scheme::verify(pp, com, op));
    EXPECT_EQ(scheme::combine_rotated(op), scheme::one::eval_rotated(op.copies[0]));
}

TEST(KPpvc, ReprogramsOneCoordinateOfTheSum)
{
    using scheme = dpf::k_ppvc<2, std::uint8_t, 3, 8>;
    const auto pp = scheme::one::setup_from_seed(seed_block(30, 31));
    const auto seed = seed_block(32, 33);
    const auto [com, st] = scheme::commit_from_seed(pp, seed);
    const auto again = scheme::commit_from_seed(pp, seed);
    EXPECT_EQ(com, again.first);
    EXPECT_EQ(st.copies[0].i, again.second.copies[0].i);
    EXPECT_EQ(st.copies[1].i, again.second.copies[1].i);
    EXPECT_NE(st.copies[0].i, st.copies[1].i);
    ASSERT_TRUE(scheme::check_well_formed(pp, com, st));
    ASSERT_TRUE(scheme::audit(pp, com, seed));
    EXPECT_FALSE(scheme::audit(pp, com, seed_block(32, 34)));

    const std::array<std::uint8_t, 2> xi{scheme::one::point(4), scheme::one::point(90)};
    const std::array<scheme::value_type, 2> tau{1, 5};
    const auto op = scheme::open(st, 0, tau, xi);
    ASSERT_TRUE(scheme::verify(pp, com, op));
    const auto sum = scheme::combine_rotated(op);

    std::array<scheme::value_type, 2> tau2{3, 5};
    const auto sum2 = scheme::combine_rotated(scheme::open(st, 0, tau2, xi));
    for (std::size_t y = 0; y < scheme::one::domain_size; ++y)
    {
        if (y == scheme::one::index_of(xi[0]))
            EXPECT_NE(sum2[y], sum[y]);
        else
            EXPECT_EQ(sum2[y], sum[y]);
    }
    EXPECT_EQ((sum2[scheme::one::index_of(xi[0])] - sum[scheme::one::index_of(xi[0])])
            & scheme::one::value_mask(),
        (tau2[0] - tau[0]) & scheme::one::value_mask());

    std::array<scheme::value_type, 2> tau3{1, 2};
    const auto sum3 = scheme::combine_rotated(scheme::open(st, 0, tau3, xi));
    for (std::size_t y = 0; y < scheme::one::domain_size; ++y)
    {
        if (y == scheme::one::index_of(xi[1]))
            EXPECT_NE(sum3[y], sum[y]);
        else
            EXPECT_EQ(sum3[y], sum[y]);
    }
}

TEST(KPpvc, SumModeAndFreshCommit)
{
    using scheme = dpf::k_ppvc<2, std::uint8_t, 3, 8>;
    const auto pp = scheme::one::setup();
    const auto [com, st] = scheme::commit(pp);
    EXPECT_NE(st.copies[0].i, st.copies[1].i);
    ASSERT_TRUE(scheme::check_well_formed(pp, com, st));

    const std::array<std::uint8_t, 2> xi{scheme::one::point(1), scheme::one::point(2)};
    const std::array<scheme::value_type, 2> tau{2, 4};
    const auto op = scheme::open(st, 1, tau, xi);
    EXPECT_TRUE(scheme::verify(pp, com, op));
    for (std::size_t r = 0; r < 2; ++r)
        EXPECT_EQ(scheme::one::column_sum(scheme::one::eval(op.copies[r])), tau[r]);

    auto bad = op;
    bad.copies[1].coins[0] = seed_block(1, 2);
    EXPECT_FALSE(scheme::verify(pp, com, bad));
    EXPECT_THROW(scheme::open(st, 3, tau, xi), std::invalid_argument);
}

TEST(KPpvc, RejectsASharedHiddenIndex)
{
    using scheme = dpf::k_ppvc<2, std::uint8_t, 2, 8>;
    const auto pp = scheme::one::setup_from_seed(seed_block(80, 81));
    const auto index = scheme::one::point(7);
    auto [c0, s0] = scheme::one::commit_at(pp, index);
    auto [c1, s1] = scheme::one::commit_at(pp, index);
    ASSERT_TRUE(scheme::one::check_well_formed(pp, c0, s0));
    ASSERT_TRUE(scheme::one::check_well_formed(pp, c1, s1));

    scheme::commitment com;
    com.copies[0] = std::move(c0);
    com.copies[1] = std::move(c1);
    scheme::state st;
    st.copies[0] = std::move(s0);
    st.copies[1] = std::move(s1);
    EXPECT_EQ(st.copies[0].i, st.copies[1].i);
    EXPECT_FALSE(scheme::check_well_formed(pp, com, st));
}