#include <gtest/gtest.h>

#include "dpf.hpp"

#include <array>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <type_traits>
#include <utility>

namespace
{

template <typename F>
F pow_elem(F base, unsigned long long exp)
{
    F r{1};
    while (exp != 0)
    {
        if ((exp & 1ull) != 0)
            r = r * base;
        exp >>= 1;
        if (exp != 0)
            base = base * base;
    }
    return r;
}

template <typename F>
void expect_field_laws()
{
    constexpr unsigned bits = F::bits;
    EXPECT_EQ(F{0} + F{1}, F{1});
    EXPECT_EQ(F{1} + F{1}, F{0});
    EXPECT_EQ(-F{1}, F{1});
    EXPECT_EQ(F{-1}, F{1});
    EXPECT_EQ(F{1} * F{1}, F{1});
    EXPECT_EQ(F{0} * F{5}, F{0});
    EXPECT_TRUE(std::is_trivially_copyable_v<F>);
    EXPECT_TRUE(std::is_standard_layout_v<F>);
    EXPECT_TRUE(dpf::utils::has_characteristic_two_v<F>);
    EXPECT_EQ(dpf::utils::bitlength_of_v<F>, bits);
    EXPECT_EQ(dpf::utils::make_default_v<F>, F{1});

    if constexpr (bits > 1)
    {
        const F x{2};
        F xpow{1};
        for (unsigned i = 0; i < bits; ++i)
            xpow = xpow * x;
        if constexpr (bits == 2 || bits == 4)
            EXPECT_EQ(xpow, F{0x3});
        else if constexpr (bits == 8)
            EXPECT_EQ(xpow, F{0x1b});
        else if constexpr (bits == 16)
            EXPECT_EQ(xpow, F{0x2d});
        else if constexpr (bits == 32)
            EXPECT_EQ(xpow.raw(), 0x90200001u);
        else
        {
            const auto tail = static_cast<typename F::integral_type>(
                (typename F::integral_type{1} << 63)
                | (typename F::integral_type{1} << 62)
                | (typename F::integral_type{1} << 53)
                | typename F::integral_type{1});
            EXPECT_EQ(xpow.raw(), tail);
        }
    }

    const F a{0x13};
    const F b{0x2a};
    const F c{0x7};
    EXPECT_EQ((a + b) * c, a * c + b * c);
    EXPECT_EQ((a * b) * c, a * (b * c));
    if constexpr (bits <= 32)
    {
        if (a != F{0})
            EXPECT_EQ(a * pow_elem(a, (1ull << bits) - 2ull), F{1});
    }

    if constexpr (bits <= 8)
    {
        const unsigned lim = 1u << bits;
        for (unsigned i = 1; i < lim; ++i)
            EXPECT_EQ(F{i} * pow_elem(F{i}, (1ull << bits) - 2ull), F{1}) << i;
    }
}

template <typename Key0, typename Key1, typename In>
auto open_at(const Key0 & k0, const Key1 & k1, In x)
{
    const auto y0 = *dpf::eval_point(k0, x);
    const auto y1 = *dpf::eval_point(k1, x);
    return dpf::reconstruct(y0, y1);
}

template <typename Out, typename In>
void expect_point_payload(In alpha, Out beta)
{
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    for (int x = 0; x < 32; ++x)
    {
        const In q = static_cast<In>(x);
        const Out got = open_at(k0, k1, q);
        EXPECT_EQ(got, q == alpha ? beta : Out{}) << static_cast<unsigned>(x);
    }
    EXPECT_EQ(open_at(k0, k1, alpha), beta);
}

template <typename Out, typename Spec>
void expect_cmp(std::uint8_t alpha, Out beta, Spec spec, bool leq)
{
    auto [k0, k1] = dpf::make_dpf(alpha, spec);
    for (int x = 0; x < 24; ++x)
    {
        const auto q = static_cast<std::uint8_t>(x);
        const auto y0 = dpf::eval_point<Out>(dpf::cmp, k0, q);
        const auto y1 = dpf::eval_point<Out>(dpf::cmp, k1, q);
        const bool hot = leq ? x <= static_cast<int>(alpha)
                             : x < static_cast<int>(alpha);
        EXPECT_EQ(dpf::reconstruct(y0, y1), hot ? beta : Out{}) << x;
    }
}

}  // namespace

TEST(Gf2, FieldLaws)
{
    expect_field_laws<dpf::gf2>();
    expect_field_laws<dpf::gf22>();
    expect_field_laws<dpf::gf24>();
    expect_field_laws<dpf::gf28>();
    expect_field_laws<dpf::gf216>();
    expect_field_laws<dpf::gf232>();
    expect_field_laws<dpf::gf264>();
}

TEST(Gf2, SubByteAdditionIsXor)
{
    EXPECT_EQ(dpf::gf22{3} + dpf::gf22{1}, dpf::gf22{2});
    EXPECT_EQ(dpf::gf24{0xf} + dpf::gf24{1}, dpf::gf24{0xe});
    EXPECT_EQ(dpf::gf24{2} * dpf::gf24{2}, dpf::gf24{4});
    EXPECT_EQ(dpf::gf22{2} * dpf::gf22{2}, dpf::gf22{3});
}

TEST(Gf2, LeafScaleIsPerLane)
{
    alignas(16) unsigned char bytes[16];
    for (int i = 0; i < 16; ++i)
        bytes[i] = 0xe4;
    simde__m128i node;
    std::memcpy(&node, bytes, sizeof(node));
    const auto scaled = dpf::multiply_leaf(node, dpf::gf24{2});
    unsigned char out[16];
    std::memcpy(out, &scaled, sizeof(out));
    const auto lo = static_cast<unsigned>((dpf::gf24{0x4} * dpf::gf24{2}).raw());
    const auto hi = static_cast<unsigned>((dpf::gf24{0xe} * dpf::gf24{2}).raw());
    const auto expect = static_cast<unsigned char>(lo | (hi << 4));
    for (unsigned char b : out)
        EXPECT_EQ(b, expect);

    const auto summed = dpf::add_leaf<dpf::gf24>(node, node);
    unsigned char z[16];
    std::memcpy(z, &summed, sizeof(z));
    for (unsigned char b : z)
        EXPECT_EQ(b, 0);
}

TEST(Gf2, ShufbScalarVectorMatchesFieldMul)
{
    alignas(32) unsigned char bytes[33];
    for (int i = 0; i < 33; ++i)
        bytes[i] = static_cast<unsigned char>(i * 17 + 3);
    const unsigned scalars[] = {0u, 1u, 2u, 0x1bu, 0x80u, 0xffu, 0x2du, 0x8000u, 0xffffu};

    for (unsigned s : scalars)
    {
        if (s > 0xffu)
            continue;
        simde__m128i n128;
        simde__m256i n256;
        std::memcpy(&n128, bytes, 16);
        std::memcpy(&n256, bytes, 32);
        const auto a128 = dpf::multiply_leaf(n128, dpf::gf28{s});
        const auto a256 = dpf::multiply_leaf(n256, dpf::gf28{s});
        unsigned char o128[16];
        unsigned char o256[32];
        std::memcpy(o128, &a128, 16);
        std::memcpy(o256, &a256, 32);
        for (int i = 0; i < 16; ++i)
            EXPECT_EQ(o128[i], (dpf::gf28{bytes[i]} * dpf::gf28{s}).raw()) << s << " " << i;
        for (int i = 0; i < 32; ++i)
            EXPECT_EQ(o256[i], (dpf::gf28{bytes[i]} * dpf::gf28{s}).raw()) << s << " " << i;

        unsigned char tail[33];
        std::memcpy(tail, bytes, 33);
        dpf::gf2_detail::scale_gf28(tail, bytes, 33, static_cast<std::uint8_t>(s));
        for (int i = 0; i < 33; ++i)
            EXPECT_EQ(tail[i], (dpf::gf28{bytes[i]} * dpf::gf28{s}).raw()) << "tail " << i;
    }

    alignas(32) unsigned char lanes[32];
    for (int i = 0; i < 16; ++i)
    {
        const auto lane = static_cast<std::uint16_t>(i * 19 + 1);
        lanes[2 * i] = static_cast<unsigned char>(lane);
        lanes[2 * i + 1] = static_cast<unsigned char>(lane >> 8);
    }
    for (unsigned s : scalars)
    {
        simde__m256i node;
        std::memcpy(&node, lanes, 32);
        const auto scaled = dpf::multiply_leaf(node, dpf::gf216{s});
        unsigned char got[32];
        std::memcpy(got, &scaled, 32);
        for (int i = 0; i < 16; ++i)
        {
            const auto lane = static_cast<std::uint16_t>(lanes[2 * i] | (lanes[2 * i + 1] << 8));
            const auto prod = (dpf::gf216{lane} * dpf::gf216{s}).raw();
            EXPECT_EQ(got[2 * i], static_cast<unsigned char>(prod)) << s << " " << i;
            EXPECT_EQ(got[2 * i + 1], static_cast<unsigned char>(prod >> 8)) << s << " " << i;
        }
    }
}

TEST(Gf2, PointPayload)
{
    expect_point_payload<dpf::gf2>(std::uint8_t{0x2a}, dpf::gf2{1});
    expect_point_payload<dpf::gf22>(std::uint8_t{0x11}, dpf::gf22{3});
    expect_point_payload<dpf::gf24>(std::uint8_t{0x05}, dpf::gf24{0xa});
    expect_point_payload<dpf::gf28>(std::uint8_t{0x2a}, dpf::gf28{0x1b});
    expect_point_payload<dpf::gf216>(std::uint8_t{0x07}, dpf::gf216{0x2d});
    expect_point_payload<dpf::gf232>(std::uint8_t{0x13}, dpf::gf232{0x90200001u});
    expect_point_payload<dpf::gf264>(std::uint8_t{0x04}, dpf::gf264{0x11});
}

TEST(Gf2, ComparisonPayload)
{
    expect_cmp(std::uint8_t{10}, dpf::gf24{0x7}, dpf::lt(dpf::gf24{0x7}), false);
    expect_cmp(std::uint8_t{10}, dpf::gf28{0x1b}, dpf::leq(dpf::gf28{0x1b}), true);
    expect_cmp(std::uint8_t{4}, dpf::gf264{0x53}, dpf::lt(dpf::gf264{0x53}), false);
}

template <typename F>
void expect_inv_exhaustive()
{
    EXPECT_THROW(dpf::detail::shamir_field<F>::inv(F{0}), std::invalid_argument);
    const unsigned long long n = 1ull << F::bits;
    for (unsigned long long i = 1; i < n; ++i)
    {
        const F a{static_cast<typename F::integral_type>(i)};
        const F b = dpf::detail::shamir_field<F>::inv(a);
        EXPECT_EQ(a * b, F{1}) << i;
    }
}

template <typename F>
void expect_shamir23(F secret, F slope)
{
    const auto shares = dpf::shamir::deal<F, 2, 3>(secret, std::array<F, 1>{{slope}});
    EXPECT_EQ((dpf::shamir::reconstruct(std::get<0>(shares), std::get<1>(shares))), secret);
    EXPECT_EQ((dpf::shamir::reconstruct(std::get<1>(shares), std::get<2>(shares))), secret);
    EXPECT_EQ((dpf::shamir::reconstruct(std::get<2>(shares), std::get<0>(shares))), secret);
    EXPECT_EQ((dpf::shamir::reconstruct(
        std::get<0>(shares), std::get<1>(shares), std::get<2>(shares))), secret);
}

TEST(Gf2, InverseAndShamir)
{
    expect_inv_exhaustive<dpf::gf2>();
    expect_inv_exhaustive<dpf::gf22>();
    expect_inv_exhaustive<dpf::gf24>();
    expect_inv_exhaustive<dpf::gf28>();

    const dpf::gf216 wide[] = {dpf::gf216{1}, dpf::gf216{2}, dpf::gf216{0x2d},
        dpf::gf216{0xffff}};
    for (auto a : wide)
        EXPECT_EQ(a * dpf::detail::shamir_field<dpf::gf216>::inv(a), dpf::gf216{1});
    const dpf::gf232 mid[] = {dpf::gf232{1}, dpf::gf232{2}, dpf::gf232{0x190200001u}};
    for (auto a : mid)
        EXPECT_EQ(a * dpf::detail::shamir_field<dpf::gf232>::inv(a), dpf::gf232{1});
    const dpf::gf264 big[] = {dpf::gf264{1}, dpf::gf264{2}, dpf::gf264{~std::uint64_t{0}}};
    for (auto a : big)
        EXPECT_EQ(a * dpf::detail::shamir_field<dpf::gf264>::inv(a), dpf::gf264{1});

    expect_shamir23(dpf::gf22{1}, dpf::gf22{2});
    expect_shamir23(dpf::gf28{0x1b}, dpf::gf28{0x5a});
    expect_shamir23(dpf::gf264{0x11}, dpf::gf264{0x53});

    const auto shares = dpf::shamir::deal<dpf::gf28, 3, 5>(
        dpf::gf28{0x1b}, std::array<dpf::gf28, 2>{{dpf::gf28{2}, dpf::gf28{9}}});
    EXPECT_EQ((dpf::shamir::reconstruct(
        std::get<0>(shares), std::get<2>(shares), std::get<4>(shares))),
        dpf::gf28{0x1b});

    // Point 2 is 0 in GF(2), so that party would hold the secret.
    const auto leaked = dpf::shamir::deal<dpf::gf2, 2, 3>(
        dpf::gf2{1}, std::array<dpf::gf2, 1>{{dpf::gf2{1}}});
    EXPECT_EQ(std::get<1>(leaked).raw(), dpf::gf2{1}.raw());
    EXPECT_THROW((dpf::shamir::reconstruct(std::get<0>(leaked), std::get<1>(leaked))),
        std::invalid_argument);

    const auto one = dpf::shamir::deal<dpf::gf2, 1, 1>(dpf::gf2{1}, std::array<dpf::gf2, 0>{});
    EXPECT_EQ(dpf::shamir::reconstruct(std::get<0>(one)), dpf::gf2{1});
}