#include <gtest/gtest.h>

#include <cstdint>
#include <cstring>
#include <limits>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

#include "dpf.hpp"

namespace
{

// Full AES-128 block payload: Boyle packing leaves lg(outputs_per_leaf) = 0,
// so a 128-bit domain walks depth 128 (Express `domainSize` = 128).
using input_t = simde_uint128;
using output_t = simde_uint128;
using dpf_key_t = dpf::utils::dpf_type_t<dpf::prg::aes128, dpf::prg::aes128,
    input_t, output_t>;

output_t make_payload()
{
    output_t y{};
    auto * bytes = reinterpret_cast<unsigned char *>(&y);
    for (std::size_t i = 0; i < sizeof(y); ++i)
        bytes[i] = static_cast<unsigned char>(0xa0 + i);
    return y;
}

input_t make_alpha()
{
    // High bit set, plus a distinctive low pattern.
    return (input_t{1} << 127) | (input_t{0x0123456789abcdefull} << 64)
        | input_t{0xfedcba9876543210ull};
}

}  // namespace

TEST(Domain128Point, DepthIs128ForFullBlockPayload)
{
    static_assert(dpf::utils::bitlength_of_v<input_t> == 128);
    static_assert(dpf_key_t::lg_outputs_per_leaf == 0);
    static_assert(dpf_key_t::outputs_per_leaf == 1);
    static_assert(dpf_key_t::depth == 128);
    EXPECT_EQ(dpf_key_t::depth, 128u);
    EXPECT_EQ(dpf::utils::bitlength_of_v<input_t>, 128u);
}

TEST(Domain128Point, PointEvalShares)
{
    const input_t alpha = make_alpha();
    const output_t beta = make_payload();
    const output_t zero{};

    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    EXPECT_EQ(k0.correction_words().size(), 128u);
    EXPECT_EQ(k1.correction_words().size(), 128u);
    EXPECT_EQ(std::decay_t<decltype(k0)>::depth, 128u);

    auto open = [&](input_t x) {
        return dpf::reconstruct(*dpf::eval_point(k0, x),
            *dpf::eval_point(k1, x));
    };

    // Programmed point: both parties' shares open to the payload.
    {
        const auto s0 = *dpf::eval_point(k0, alpha);
        const auto s1 = *dpf::eval_point(k1, alpha);
        EXPECT_EQ(dpf::reconstruct(s0, s1), beta);
        EXPECT_EQ(open(alpha), beta);
    }

    const std::vector<input_t> off_points = {
        input_t{0},
        input_t{1},
        alpha + 1,
        alpha - 1,
        input_t{1} << 127,
        (input_t{1} << 127) | input_t{1},
        alpha ^ (input_t{1} << 64),
        alpha ^ input_t{1},
        ~input_t{0},
    };

    for (const input_t x : off_points)
    {
        if (x == alpha)
            continue;
        const auto s0 = *dpf::eval_point(k0, x);
        const auto s1 = *dpf::eval_point(k1, x);
        EXPECT_EQ(dpf::reconstruct(s0, s1), zero) << "off-point open";
    }

    // Every individual bit of alpha must be on the path.
    for (int i = 0; i < 128; ++i)
    {
        const input_t flipped = alpha ^ (input_t{1} << i);
        const auto s0 = *dpf::eval_point(k0, flipped);
        const auto s1 = *dpf::eval_point(k1, flipped);
        EXPECT_EQ(dpf::reconstruct(s0, s1), zero) << "bit " << i;
    }
}

TEST(Domain128Point, Uint128ClassInput)
{
    using in_t = uint128_t;
    using out_t = simde_uint128;
    using kt = dpf::utils::dpf_type_t<dpf::prg::aes128, dpf::prg::aes128,
        in_t, out_t>;
    static_assert(kt::depth == 128);
    static_assert(dpf::utils::bitlength_of_v<in_t> == 128);

    // uint128_t(upper, lower): bit 127 set.
    const in_t alpha{std::uint64_t{1} << 63, 0};
    out_t beta{};
    std::memset(&beta, 0x5c, sizeof(beta));

    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    EXPECT_EQ(k0.correction_words().size(), 128u);

    const auto s0 = *dpf::eval_point(k0, alpha);
    const auto s1 = *dpf::eval_point(k1, alpha);
    EXPECT_EQ(dpf::reconstruct(s0, s1), beta);

    const in_t neigh = alpha + in_t{1};
    EXPECT_EQ(dpf::reconstruct(*dpf::eval_point(k0, neigh),
        *dpf::eval_point(k1, neigh)), out_t{});
}

TEST(Domain128Point, FullDomainExpansionThrows)
{
    const input_t alpha = make_alpha();
    const output_t beta = make_payload();
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    (void)k1;

    EXPECT_THROW(
        (void)dpf::make_output_buffer_for_full(k0),
        std::length_error);

    EXPECT_THROW(
        (void)dpf::eval_full(k0),
        std::length_error);

    EXPECT_THROW(
        (void)dpf::eval_interval(k0,
            std::numeric_limits<input_t>::min(),
            std::numeric_limits<input_t>::max()),
        std::length_error);

    EXPECT_THROW(
        (void)dpf::make_basic_full_memoizer(k0),
        std::length_error);
}