/// @file dpf/pprf.hpp
/// @brief GGM puncturable PRF on the library AES-128 PRG.
/// @details A master seed evaluates at every domain point. `puncture` at a
///          single `α` returns one sibling seed per level so evaluation
///          everywhere except `α` costs O(n) PRG calls. `puncture` over a
///          set `H` returns the copath of that set: every node whose parent
///          lies on a path to `H` and which itself does not. Shared prefixes
///          are stored once. The programmed leaf values are stored beside the
///          puncture when the holder of the master asks for them; an audit
///          opening of a replica-seed pool leaves them out (`program_hidden =
///          false`). Complexity of the set walk is O(|H| · n) PRG expansions
///          and at most |H| · n published nodes; the domain is never
///          materialized. One-point and set forms agree on `{α}`.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license.

#ifndef LIBDPF_INCLUDE_DPF_PPRF_HPP__
#define LIBDPF_INCLUDE_DPF_PPRF_HPP__

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <vector>

#include "hedley/hedley.h"

#include "dpf/prg_aes.hpp"
#include "dpf/random.hpp"
#include "dpf/utils.hpp"

namespace dpf
{

/// @brief Master key for a puncturable PRF over `InputT`.
/// @tparam InputT domain type (`uint32_t`, `simde_uint128`, ...)
/// @tparam PRG exterior/interior PRG; defaults to AES-128
template <typename InputT, typename PRG = dpf::prg::aes128>
struct pprf_master
{
    using input_type = InputT;
    using prg = PRG;
    using block_type = typename PRG::block_type;
    static constexpr std::size_t bitlength = utils::bitlength_of_v<InputT>;

    block_type root{};
};

/// @brief Punctured key: sibling seed per level, optional programmed leaf.
template <typename InputT, typename PRG = dpf::prg::aes128>
struct pprf_punctured
{
    using input_type = InputT;
    using prg = PRG;
    using block_type = typename PRG::block_type;
    static constexpr std::size_t bitlength = utils::bitlength_of_v<InputT>;

    input_type alpha{};
    /// @brief `siblings[i]` is the seed for the sibling of the path bit at
    ///        level `i` (MSB = 0).
    std::array<block_type, bitlength> siblings{};
    /// @brief Path bit taken toward `alpha` at each level (1 = right).
    std::array<std::uint8_t, bitlength> path_bits{};
    /// @brief When set, `pprf_eval` at `alpha` returns this leaf.
    std::optional<block_type> programmed{};
};

/// @brief One published node on the copath of a hidden set.
/// @details `level` is the number of path bits from the root (`0` only for the
///          empty-`H` root). `prefix` holds those bits right-aligned, so the
///          top `level` bits of a domain point `x` match when
///          `(x >> (n - level)) == prefix`.
template <typename InputT, typename PRG = dpf::prg::aes128>
struct pprf_copath_node
{
    using input_type = InputT;
    using prg = PRG;
    using block_type = typename PRG::block_type;

    std::size_t level{};
    input_type prefix{};
    block_type seed{};
};

/// @brief Copath of a hidden set: sorted unique points plus published nodes.
/// @details Empty `hidden` publishes the root so every domain point evaluates.
///          A full-domain `hidden` publishes nothing. When `programmed` is
///          non-empty it has one leaf per hidden point, in the same order as
///          `hidden`. A live seed is never among `nodes`.
template <typename InputT, typename PRG = dpf::prg::aes128>
struct pprf_copath
{
    using input_type = InputT;
    using prg = PRG;
    using block_type = typename PRG::block_type;
    static constexpr std::size_t bitlength = utils::bitlength_of_v<InputT>;

    std::vector<input_type> hidden{};
    std::vector<pprf_copath_node<InputT, PRG>> nodes{};
    std::vector<block_type> programmed{};
};

/// @brief Sample a fresh master seed.
template <typename InputT, typename PRG = dpf::prg::aes128>
HEDLEY_WARN_UNUSED_RESULT
HEDLEY_NO_THROW
pprf_master<InputT, PRG> make_pprf_master() noexcept
{
    pprf_master<InputT, PRG> m;
    m.root = dpf::uniform_sample<typename PRG::block_type>();
    return m;
}

namespace detail
{
namespace pprf_impl
{

/// @brief Path bit of `x` at `level` (0 = MSB).
template <typename InputT>
HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr bool path_bit(InputT x, std::size_t level) noexcept
{
    constexpr std::size_t n = utils::bitlength_of_v<InputT>;
    constexpr auto to_int = utils::to_integral_type<InputT>{};
    using integral_t = typename utils::to_integral_type<InputT>::integral_type;
    const integral_t xi = static_cast<integral_t>(to_int(x));
    const std::size_t shift = n - 1 - level;
    return static_cast<bool>(utils::shift_right(xi, shift) & integral_t{1});
}

/// @brief Split `seed` into left (pos 0) and right (pos 1) children.
template <typename PRG>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
void expand_children(typename PRG::block_type seed,
    typename PRG::block_type & left, typename PRG::block_type & right) noexcept
{
    // Match DPF interior: pos 0 = left child, pos 1 = right child.
    left = PRG::eval(seed, 0);
    right = PRG::eval(seed, 1);
}

/// @brief Top `level` path bits of `x`, right-aligned.
template <typename InputT>
HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr InputT path_prefix(InputT x, std::size_t level) noexcept
{
    if (level == 0)
        return InputT{};
    constexpr std::size_t n = utils::bitlength_of_v<InputT>;
    constexpr auto to_int = utils::to_integral_type<InputT>{};
    using integral_t = typename utils::to_integral_type<InputT>::integral_type;
    const integral_t xi = static_cast<integral_t>(to_int(x));
    return static_cast<InputT>(utils::shift_right(xi, n - level));
}

/// @brief Whether `x` lies under the node `(level, prefix)`.
template <typename InputT>
HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr bool prefix_matches(InputT x, std::size_t level,
    InputT prefix) noexcept
{
    return path_prefix(x, level) == prefix;
}

/// @brief Recursively publish the copath of `hidden[lo, hi)` under `seed`.
/// @details `hidden` must be sorted unique. Emits at most one node per
///          off-path sibling; shared prefixes are visited once. When
///          `program_hidden` is set, appends leaf seeds in sorted `hidden`
///          order.
template <typename InputT, typename PRG>
void collect_copath(typename PRG::block_type seed, std::size_t level,
    InputT prefix, const std::vector<InputT> & hidden, std::size_t lo,
    std::size_t hi, bool program_hidden,
    std::vector<pprf_copath_node<InputT, PRG>> & nodes,
    std::vector<typename PRG::block_type> & programmed)
{
    constexpr std::size_t n = utils::bitlength_of_v<InputT>;
    if (lo == hi)
        return;
    if (level == n)
    {
        if (program_hidden)
            programmed.push_back(seed);
        return;
    }

    // `hidden` is sorted by domain value, so MSB-first path order matches.
    std::size_t mid = lo;
    while (mid < hi && !path_bit(hidden[mid], level))
        ++mid;

    typename PRG::block_type left{}, right{};
    expand_children<PRG>(seed, left, right);

    const InputT left_prefix =
        static_cast<InputT>((static_cast<InputT>(prefix) << 1) | InputT{0});
    const InputT right_prefix =
        static_cast<InputT>((static_cast<InputT>(prefix) << 1) | InputT{1});

    const bool left_on_path = mid > lo;
    const bool right_on_path = hi > mid;

    // Match the one-point walk: publish the off-path sibling, then descend.
    if (left_on_path && right_on_path)
    {
        collect_copath<InputT, PRG>(left, level + 1, left_prefix, hidden, lo,
            mid, program_hidden, nodes, programmed);
        collect_copath<InputT, PRG>(right, level + 1, right_prefix, hidden, mid,
            hi, program_hidden, nodes, programmed);
    }
    else if (left_on_path)
    {
        nodes.push_back(
            pprf_copath_node<InputT, PRG>{level + 1, right_prefix, right});
        collect_copath<InputT, PRG>(left, level + 1, left_prefix, hidden, lo,
            mid, program_hidden, nodes, programmed);
    }
    else
    {
        nodes.push_back(
            pprf_copath_node<InputT, PRG>{level + 1, left_prefix, left});
        collect_copath<InputT, PRG>(right, level + 1, right_prefix, hidden, mid,
            hi, program_hidden, nodes, programmed);
    }
}

}  // namespace pprf_impl
}  // namespace detail

/// @brief Evaluate the master at `x`. Returns one AES block (the leaf seed).
template <typename InputT, typename PRG = dpf::prg::aes128>
HEDLEY_WARN_UNUSED_RESULT
HEDLEY_NO_THROW
typename PRG::block_type pprf_eval(const pprf_master<InputT, PRG> & master,
    InputT x) noexcept
{
    using block = typename PRG::block_type;
    constexpr std::size_t n = pprf_master<InputT, PRG>::bitlength;
    block cur = master.root;
    for (std::size_t level = 0; level < n; ++level)
    {
        block left{}, right{};
        detail::pprf_impl::expand_children<PRG>(cur, left, right);
        cur = detail::pprf_impl::path_bit(x, level) ? right : left;
    }
    return cur;
}

/// @brief Puncture `master` at `alpha`. Optionally program `F(alpha)`.
/// @param program_alpha when true (default), store `F(α)` beside the siblings
template <typename InputT, typename PRG = dpf::prg::aes128>
HEDLEY_WARN_UNUSED_RESULT
HEDLEY_NO_THROW
pprf_punctured<InputT, PRG> puncture(const pprf_master<InputT, PRG> & master,
    InputT alpha, bool program_alpha = true) noexcept
{
    using block = typename PRG::block_type;
    constexpr std::size_t n = pprf_master<InputT, PRG>::bitlength;
    pprf_punctured<InputT, PRG> out{};
    out.alpha = alpha;
    block cur = master.root;
    for (std::size_t level = 0; level < n; ++level)
    {
        block left{}, right{};
        detail::pprf_impl::expand_children<PRG>(cur, left, right);
        const bool bit = detail::pprf_impl::path_bit(alpha, level);
        out.path_bits[level] = static_cast<std::uint8_t>(bit);
        out.siblings[level] = bit ? left : right;  // keep the off-path child
        cur = bit ? right : left;
    }
    if (program_alpha)
        out.programmed = cur;
    return out;
}

/// @brief Evaluate a punctured key. At `alpha` returns the programmed leaf
///        when present; otherwise throws `std::invalid_argument`.
template <typename InputT, typename PRG = dpf::prg::aes128>
HEDLEY_WARN_UNUSED_RESULT
typename PRG::block_type pprf_eval(const pprf_punctured<InputT, PRG> & key,
    InputT x)
{
    using block = typename PRG::block_type;
    constexpr std::size_t n = pprf_punctured<InputT, PRG>::bitlength;
    if (x == key.alpha)
    {
        if (!key.programmed.has_value())
            throw std::invalid_argument(
                "pprf_eval: punctured point has no programmed value");
        return *key.programmed;
    }

    // First level where x diverges from alpha; expand from that sibling.
    for (std::size_t level = 0; level < n; ++level)
    {
        const bool xbit = detail::pprf_impl::path_bit(x, level);
        const bool abit = static_cast<bool>(key.path_bits[level]);
        if (xbit == abit)
            continue;
        block cur = key.siblings[level];
        for (std::size_t j = level + 1; j < n; ++j)
        {
            block left{}, right{};
            detail::pprf_impl::expand_children<PRG>(cur, left, right);
            cur = detail::pprf_impl::path_bit(x, j) ? right : left;
        }
        return cur;
    }
    throw std::logic_error("pprf_eval: path didn't diverge from alpha");
}

/// @brief Puncture `master` at every point in `[first, last)`.
/// @details Duplicate points are collapsed. Empty `H` publishes the root.
///          Full-domain `H` publishes nothing. Shared path prefixes are
///          stored once. Visits O(|H| · n) nodes and emits at most |H| · n.
/// @param program_hidden defaults to `false` so an audit opening does not
///        carry the live seeds; set `true` to store `F(h)` for each `h ∈ H`
/// @return sorted unique `hidden`, the published copath, and optional leaves
template <typename InputT, typename PRG = dpf::prg::aes128, typename ForwardIt>
HEDLEY_WARN_UNUSED_RESULT
pprf_copath<InputT, PRG> puncture(const pprf_master<InputT, PRG> & master,
    ForwardIt first, ForwardIt last, bool program_hidden = false)
{
    pprf_copath<InputT, PRG> out{};
    out.hidden.assign(first, last);
    std::sort(out.hidden.begin(), out.hidden.end());
    out.hidden.erase(std::unique(out.hidden.begin(), out.hidden.end()),
        out.hidden.end());

    if (out.hidden.empty())
    {
        out.nodes.push_back(
            pprf_copath_node<InputT, PRG>{0, InputT{}, master.root});
        return out;
    }

    if (program_hidden)
        out.programmed.reserve(out.hidden.size());

    detail::pprf_impl::collect_copath<InputT, PRG>(master.root, 0, InputT{},
        out.hidden, 0, out.hidden.size(), program_hidden, out.nodes,
        out.programmed);
    return out;
}

/// @brief Evaluate a set-punctured key at `x`.
/// @details Expands from the deepest copath node whose prefix matches `x`,
///          using the same child split as the one-point eval. A hidden `x`
///          returns the programmed leaf when one was stored, and throws
///          `std::invalid_argument` otherwise. Throws the same when no
///          published node covers `x` (full-domain puncture).
template <typename InputT, typename PRG = dpf::prg::aes128>
HEDLEY_WARN_UNUSED_RESULT
typename PRG::block_type pprf_eval(const pprf_copath<InputT, PRG> & key,
    InputT x)
{
    using block = typename PRG::block_type;
    constexpr std::size_t n = pprf_copath<InputT, PRG>::bitlength;

    const auto hit = std::lower_bound(key.hidden.begin(), key.hidden.end(), x);
    if (hit != key.hidden.end() && *hit == x)
    {
        if (key.programmed.empty())
            throw std::invalid_argument(
                "pprf_eval: hidden point has no programmed value");
        const auto idx = static_cast<std::size_t>(
            std::distance(key.hidden.begin(), hit));
        return key.programmed[idx];
    }

    const pprf_copath_node<InputT, PRG> * best = nullptr;
    for (const auto & node : key.nodes)
    {
        if (node.level == 0
            || detail::pprf_impl::prefix_matches(x, node.level, node.prefix))
        {
            if (best == nullptr || node.level > best->level)
                best = &node;
        }
    }
    if (best == nullptr)
        throw std::invalid_argument("pprf_eval: no copath node covers point");

    block cur = best->seed;
    for (std::size_t level = best->level; level < n; ++level)
    {
        block left{}, right{};
        detail::pprf_impl::expand_children<PRG>(cur, left, right);
        cur = detail::pprf_impl::path_bit(x, level) ? right : left;
    }
    return cur;
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PPRF_HPP__