#include <gtest/gtest.h>

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

#include "dpf.hpp"

namespace
{

template <typename Block>
bool eq_block(const Block & a, const Block & b)
{
    return std::memcmp(&a, &b, sizeof(Block)) == 0;
}

template <typename InputT>
void check_pprf_spread(InputT alpha)
{
    auto master = dpf::make_pprf_master<InputT>();
    auto punctured = dpf::puncture(master, alpha, /*program_alpha=*/true);

    const auto master_at_alpha = dpf::pprf_eval(master, alpha);
    EXPECT_TRUE(eq_block(master_at_alpha, *punctured.programmed));
    EXPECT_TRUE(eq_block(dpf::pprf_eval(punctured, alpha), master_at_alpha));

    std::vector<InputT> points = {
        InputT{0},
        InputT{1},
        static_cast<InputT>(alpha + 1),
        static_cast<InputT>(alpha ^ InputT{1}),
    };
    if constexpr (dpf::utils::bitlength_of_v<InputT> > 8)
    {
        points.push_back(static_cast<InputT>(
            InputT{1} << (dpf::utils::bitlength_of_v<InputT> / 2)));
        points.push_back(static_cast<InputT>(alpha ^ (InputT{1} << 7)));
    }

    for (InputT x : points)
    {
        if (x == alpha)
            continue;
        EXPECT_TRUE(eq_block(dpf::pprf_eval(master, x),
            dpf::pprf_eval(punctured, x)));
    }

    auto bare = dpf::puncture(master, alpha, /*program_alpha=*/false);
    EXPECT_FALSE(bare.programmed.has_value());
    EXPECT_THROW((void)dpf::pprf_eval(bare, alpha), std::invalid_argument);
    for (InputT x : points)
    {
        if (x == alpha)
            continue;
        EXPECT_TRUE(eq_block(dpf::pprf_eval(master, x),
            dpf::pprf_eval(bare, x)));
    }

    EXPECT_EQ(punctured.siblings.size(), dpf::utils::bitlength_of_v<InputT>);
}

bool node_has_seed(const dpf::pprf_copath<std::uint8_t> & copath,
    const dpf::prg::aes128::block_type & seed)
{
    for (const auto & node : copath.nodes)
    {
        if (eq_block(node.seed, seed))
            return true;
    }
    return false;
}

}  // namespace

TEST(Pprf, PathBitHelpersAreConstexpr)
{
    using input_t = std::uint8_t;
    static_assert(dpf::detail::pprf_impl::path_bit<input_t>(0b1010'0000u, 0));
    static_assert(!dpf::detail::pprf_impl::path_bit<input_t>(0b0010'0000u, 0));
    static_assert(dpf::detail::pprf_impl::path_prefix<input_t>(0b1010'1100u, 4)
        == input_t{0b1010u});
    static_assert(dpf::detail::pprf_impl::prefix_matches<input_t>(
        0b1010'1100u, 4, input_t{0b1010u}));
    static_assert(!dpf::detail::pprf_impl::prefix_matches<input_t>(
        0b1010'1100u, 4, input_t{0b1011u}));
    SUCCEED();
}

TEST(Pprf, Uint32)
{
    check_pprf_spread<std::uint32_t>(0x00c0ffeeu);
}

TEST(Pprf, Uint128)
{
    using input_t = simde_uint128;
    const input_t alpha = (input_t{1} << 120) | input_t{0xdeadbeefull};
    check_pprf_spread<input_t>(alpha);
}

TEST(PprfCopath, OneHiddenAgreesWithOnePoint)
{
    using input_t = std::uint8_t;
    constexpr std::size_t n = dpf::utils::bitlength_of_v<input_t>;
    auto master = dpf::make_pprf_master<input_t>();
    const input_t alpha = 0x2au;

    auto one = dpf::puncture(master, alpha, /*program_alpha=*/true);
    const input_t hidden[] = {alpha};
    auto copath = dpf::puncture(master, std::begin(hidden), std::end(hidden),
        /*program_hidden=*/true);

    ASSERT_EQ(copath.nodes.size(), n);
    ASSERT_EQ(copath.programmed.size(), 1u);
    EXPECT_TRUE(eq_block(copath.programmed[0], *one.programmed));

    for (std::size_t i = 0; i < n; ++i)
    {
        EXPECT_EQ(copath.nodes[i].level, i + 1);
        EXPECT_TRUE(eq_block(copath.nodes[i].seed, one.siblings[i]));
        EXPECT_EQ(copath.nodes[i].prefix,
            dpf::detail::pprf_impl::path_prefix(
                static_cast<input_t>(
                    alpha ^ static_cast<input_t>(input_t{1} << (n - 1 - i))),
                i + 1));
    }

    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xi = static_cast<input_t>(x);
        EXPECT_TRUE(eq_block(dpf::pprf_eval(copath, xi), dpf::pprf_eval(one, xi)));
        if (xi != alpha)
            EXPECT_TRUE(
                eq_block(dpf::pprf_eval(copath, xi), dpf::pprf_eval(master, xi)));
    }

    auto bare = dpf::puncture(master, std::begin(hidden), std::end(hidden),
        /*program_hidden=*/false);
    EXPECT_TRUE(bare.programmed.empty());
    EXPECT_THROW((void)dpf::pprf_eval(bare, alpha), std::invalid_argument);
    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xi = static_cast<input_t>(x);
        if (xi == alpha)
            continue;
        EXPECT_TRUE(
            eq_block(dpf::pprf_eval(bare, xi), dpf::pprf_eval(master, xi)));
    }
}

TEST(PprfCopath, AdjacentAndSeparatedSets)
{
    using input_t = std::uint8_t;
    constexpr std::size_t n = dpf::utils::bitlength_of_v<input_t>;
    auto master = dpf::make_pprf_master<input_t>();

    const input_t adjacent[] = {2, 3};
    auto adj = dpf::puncture(master, std::begin(adjacent), std::end(adjacent));
    // Share a 7-bit prefix: one sibling per shared level, none at the leaves.
    EXPECT_EQ(adj.nodes.size(), n - 1);
    EXPECT_LE(adj.nodes.size(), n * adj.hidden.size());
    EXPECT_LT(adj.nodes.size(), 2 * n);

    const input_t separated[] = {2, 5};
    auto sep = dpf::puncture(master, std::begin(separated), std::end(separated));
    EXPECT_LE(sep.nodes.size(), n * sep.hidden.size());
    EXPECT_GT(sep.nodes.size(), adj.nodes.size());

    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xi = static_cast<input_t>(x);
        const bool in_adj = (xi == 2 || xi == 3);
        const bool in_sep = (xi == 2 || xi == 5);
        if (in_adj)
            EXPECT_THROW((void)dpf::pprf_eval(adj, xi), std::invalid_argument);
        else
            EXPECT_TRUE(
                eq_block(dpf::pprf_eval(adj, xi), dpf::pprf_eval(master, xi)));
        if (in_sep)
            EXPECT_THROW((void)dpf::pprf_eval(sep, xi), std::invalid_argument);
        else
            EXPECT_TRUE(
                eq_block(dpf::pprf_eval(sep, xi), dpf::pprf_eval(master, xi)));
    }
}

TEST(PprfCopath, HiddenLeafNotPublishedWhenSiblingHidden)
{
    using input_t = std::uint8_t;
    auto master = dpf::make_pprf_master<input_t>();
    const input_t leaf = 2;
    const auto leaf_seed = dpf::pprf_eval(master, leaf);

    const input_t alone[] = {leaf};
    auto solo = dpf::puncture(master, std::begin(alone), std::end(alone),
        /*program_hidden=*/true);
    EXPECT_TRUE(node_has_seed(solo, dpf::pprf_eval(master, static_cast<input_t>(3))));

    const input_t both[] = {2, 3};
    auto pair = dpf::puncture(master, std::begin(both), std::end(both),
        /*program_hidden=*/true);
    EXPECT_FALSE(node_has_seed(pair, leaf_seed));
    EXPECT_FALSE(node_has_seed(pair, dpf::pprf_eval(master, static_cast<input_t>(3))));
    ASSERT_EQ(pair.programmed.size(), 2u);
    EXPECT_TRUE(eq_block(pair.programmed[0], leaf_seed));
    EXPECT_TRUE(eq_block(pair.programmed[1],
        dpf::pprf_eval(master, static_cast<input_t>(3))));
    EXPECT_TRUE(eq_block(dpf::pprf_eval(pair, leaf), leaf_seed));
    EXPECT_TRUE(eq_block(dpf::pprf_eval(pair, static_cast<input_t>(3)),
        dpf::pprf_eval(master, static_cast<input_t>(3))));
}

TEST(PprfCopath, EmptyAndDuplicates)
{
    using input_t = std::uint8_t;
    auto master = dpf::make_pprf_master<input_t>();

    const std::vector<input_t> none{};
    auto empty = dpf::puncture(master, none.begin(), none.end());
    ASSERT_EQ(empty.nodes.size(), 1u);
    EXPECT_EQ(empty.nodes[0].level, 0u);
    EXPECT_EQ(empty.nodes[0].prefix, input_t{});
    EXPECT_TRUE(eq_block(empty.nodes[0].seed, master.root));
    EXPECT_TRUE(empty.hidden.empty());
    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xi = static_cast<input_t>(x);
        EXPECT_TRUE(
            eq_block(dpf::pprf_eval(empty, xi), dpf::pprf_eval(master, xi)));
    }

    const input_t dups[] = {5, 2, 5, 2, 5};
    auto once = dpf::puncture(master, std::begin(dups), std::end(dups));
    const input_t unique[] = {2, 5};
    auto clean = dpf::puncture(master, std::begin(unique), std::end(unique));
    ASSERT_EQ(once.hidden.size(), 2u);
    EXPECT_EQ(once.hidden, clean.hidden);
    ASSERT_EQ(once.nodes.size(), clean.nodes.size());
    for (std::size_t i = 0; i < once.nodes.size(); ++i)
    {
        EXPECT_EQ(once.nodes[i].level, clean.nodes[i].level);
        EXPECT_EQ(once.nodes[i].prefix, clean.nodes[i].prefix);
        EXPECT_TRUE(eq_block(once.nodes[i].seed, clean.nodes[i].seed));
    }
}

TEST(PprfCopath, FullDomainPublishesNothing)
{
    using input_t = std::uint8_t;
    auto master = dpf::make_pprf_master<input_t>();
    std::vector<input_t> all(256);
    for (unsigned x = 0; x < 256; ++x)
        all[x] = static_cast<input_t>(x);

    auto bare = dpf::puncture(master, all.begin(), all.end(),
        /*program_hidden=*/false);
    EXPECT_TRUE(bare.nodes.empty());
    ASSERT_EQ(bare.hidden.size(), 256u);
    EXPECT_TRUE(bare.programmed.empty());
    for (unsigned x = 0; x < 256; ++x)
        EXPECT_THROW((void)dpf::pprf_eval(bare, static_cast<input_t>(x)),
            std::invalid_argument);

    auto programmed = dpf::puncture(master, all.begin(), all.end(),
        /*program_hidden=*/true);
    EXPECT_TRUE(programmed.nodes.empty());
    ASSERT_EQ(programmed.programmed.size(), 256u);
    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xi = static_cast<input_t>(x);
        EXPECT_TRUE(eq_block(dpf::pprf_eval(programmed, xi),
            dpf::pprf_eval(master, xi)));
    }
}

TEST(PprfCopath, NodePrefixesAreConsistent)
{
    using input_t = std::uint8_t;
    auto master = dpf::make_pprf_master<input_t>();
    const input_t hidden[] = {2, 5, 200};
    auto copath = dpf::puncture(master, std::begin(hidden), std::end(hidden));

    EXPECT_LE(copath.nodes.size(),
        dpf::utils::bitlength_of_v<input_t> * copath.hidden.size());

    for (const auto & node : copath.nodes)
    {
        ASSERT_GE(node.level, 1u);
        ASSERT_LE(node.level, dpf::utils::bitlength_of_v<input_t>);
        // Right-aligned prefix fits in `level` bits.
        if (node.level < dpf::utils::bitlength_of_v<input_t>)
        {
            EXPECT_LT(static_cast<unsigned>(node.prefix),
                1u << node.level);
        }
    }

    for (unsigned x = 0; x < 256; ++x)
    {
        const auto xi = static_cast<input_t>(x);
        if (std::binary_search(std::begin(hidden), std::end(hidden), xi))
            continue;
        EXPECT_TRUE(
            eq_block(dpf::pprf_eval(copath, xi), dpf::pprf_eval(master, xi)));
    }
}

TEST(PprfCopath, Uint32SetMatchesMaster)
{
    using input_t = std::uint32_t;
    auto master = dpf::make_pprf_master<input_t>();
    const input_t hidden[] = {0u, 0x00c0ffeeu, 0xffffffffu};
    auto copath = dpf::puncture(master, std::begin(hidden), std::end(hidden),
        /*program_hidden=*/true);

    ASSERT_EQ(copath.programmed.size(), 3u);
    for (input_t h : hidden)
        EXPECT_TRUE(eq_block(dpf::pprf_eval(copath, h), dpf::pprf_eval(master, h)));

    const input_t samples[] = {1u, 2u, 0x00c0ffefu, 0x80000000u, 0xfffffffeu};
    for (input_t x : samples)
    {
        if (std::find(std::begin(hidden), std::end(hidden), x) != std::end(hidden))
            continue;
        EXPECT_TRUE(
            eq_block(dpf::pprf_eval(copath, x), dpf::pprf_eval(master, x)));
    }
}