#include <gtest/gtest.h>

#include <cstdint>
#include <utility>
#include <vector>

#include "dpf.hpp"

namespace
{

// Full-block payload ⇒ outputs_per_leaf = 1, so a single public F matches the
// leaf correction word share-for-share.
using input_t = std::uint8_t;
using output_t = simde_uint128;
constexpr std::size_t n = 256;

output_t leaf_cw_as_output(const dpf::utils::dpf_type_t<dpf::prg::aes128,
    dpf::prg::aes128, input_t, output_t> & key)
{
    using exterior = typename std::decay_t<decltype(key)>::exterior_node;
    return dpf::extract_leaf<exterior, output_t>(key.template leaf<0>(), 0);
}

// party_key wrappers from make_dpf
template <typename Key>
output_t leaf_cw_as_output(const Key & key)
{
    using exterior = typename Key::exterior_node;
    return dpf::extract_leaf<exterior, output_t>(key.template leaf<0>(), 0);
}

output_t make_payload(std::uint64_t lo)
{
    return output_t{lo};
}

}  // namespace

TEST(LeafLater, MatchesCorrectedImmediately)
{
    const input_t alpha = 42;
    const output_t beta = make_payload(7);

    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    // Public leaf CW (same on both keys); applying it recovers the immediate path.
    const output_t F = leaf_cw_as_output(k0);
    EXPECT_EQ(F, leaf_cw_as_output(k1));

    auto full0 = dpf::eval_full(k0);
    auto full1 = dpf::eval_full(k1);

    std::vector<output_t> buf0(n), buf1(n);
    std::vector<std::uint8_t> c0(n), c1(n);
    dpf::eval_full(buf0, c0, k0, dpf::leaf_later{});
    dpf::eval_full(buf1, c1, k1, dpf::leaf_later{});
    dpf::apply_leaf_correction(buf0, c0, F);
    dpf::apply_leaf_correction(buf1, c1, F);

    auto it0 = std::begin(full0.second);
    auto it1 = std::begin(full1.second);
    for (std::size_t i = 0; i < n; ++i, ++it0, ++it1)
    {
        const auto imm0 = dpf::detail_walk::group_value(*it0);
        const auto imm1 = dpf::detail_walk::group_value(*it1);
        EXPECT_EQ(buf0[i], imm0) << "party0 @" << i;
        EXPECT_EQ(buf1[i], imm1) << "party1 @" << i;
        EXPECT_EQ(buf0[i] - buf1[i], imm0 - imm1) << i;
        const auto got = buf0[i] - buf1[i];
        if (i == alpha)
            EXPECT_EQ(got, beta);
        else
            EXPECT_EQ(got, output_t{0});
    }
}

TEST(LeafLater, ApplyAfterRotateMatchesCorrectThenRotate)
{
    const input_t alpha = 10;
    const output_t beta = make_payload(99);
    const std::size_t s = 32;

    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    const output_t F = leaf_cw_as_output(k0);

    auto full0 = dpf::eval_full(k0);
    auto full1 = dpf::eval_full(k1);
    std::vector<output_t> early0(n), early1(n);
    {
        auto a = std::begin(full0.second);
        auto b = std::begin(full1.second);
        for (std::size_t i = 0; i < n; ++i, ++a, ++b)
        {
            early0[i] = dpf::detail_walk::group_value(*a);
            early1[i] = dpf::detail_walk::group_value(*b);
        }
    }
    early0 = dpf::cyclic_shift(early0, s);
    early1 = dpf::cyclic_shift(early1, s);

    std::vector<output_t> late0(n), late1(n);
    std::vector<std::uint8_t> c0(n), c1(n);
    dpf::eval_full(late0, c0, k0, dpf::leaf_later{});
    dpf::eval_full(late1, c1, k1, dpf::leaf_later{});
    dpf::cyclic_shift_pair(late0, c0, s);
    dpf::cyclic_shift_pair(late1, c1, s);
    dpf::apply_leaf_correction(late0, c0, F);
    dpf::apply_leaf_correction(late1, c1, F);

    for (std::size_t i = 0; i < n; ++i)
    {
        EXPECT_EQ(late0[i], early0[i]) << i;
        EXPECT_EQ(late1[i], early1[i]) << i;
    }
}

TEST(LeafLater, DuoramShapedUpdate)
{
    // Unit at r, rotate by i*−r, then apply a public F built from the leaf CW
    // of a key keyed at the online message (same as correcting immediately for
    // that message). Neighbours stay unchanged.
    const input_t r = 10;
    const input_t i_star = 42;
    const std::size_t shift = static_cast<std::size_t>(i_star - r);
    const output_t message = make_payload(7);

    auto [k0, k1] = dpf::make_dpf(r, message);
    const output_t F = leaf_cw_as_output(k0);

    std::vector<output_t> mem0(n), mem1(n);
    mem0[i_star] = make_payload(100);

    std::vector<output_t> v0(n), v1(n);
    std::vector<std::uint8_t> t0(n), t1(n);
    dpf::eval_full_add_into(v0, t0, k0, dpf::leaf_later{}, dpf::rotate{shift});
    dpf::eval_full_add_into(v1, t1, k1, dpf::leaf_later{}, dpf::rotate{shift});
    dpf::apply_leaf_correction(v0, t0, F);
    dpf::apply_leaf_correction(v1, t1, F);

    for (std::size_t i = 0; i < n; ++i)
    {
        mem0[i] = mem0[i] + v0[i];
        mem1[i] = mem1[i] + v1[i];
    }

    EXPECT_EQ(mem0[i_star] - mem1[i_star], make_payload(100) + message);
    EXPECT_EQ(mem0[r] - mem1[r], output_t{0});
    EXPECT_EQ(mem0[0] - mem1[0], output_t{0});
}