/// @file dpf/dpf3_cmp.hpp
/// @brief Three-evaluator comparison, blocked comparison, and interval keys.
/// @details Each evaluator holds one two-party DCF share of a shared tree
///          (party 1 and 3 hold the party-0 half; party 2 holds the party-1
///          half). Evaluating a single half yields that party's complementary
///          share of the predicate; open with `reconstruct_cmp_halves` on any
///          authorized pair (1+2 or 2+3). The threshold is never assembled by
///          locally reconstructing a full two-party key. Interval containment
///          reuses the three-party comparison; `lo`/`hi` stay public as in F_IC.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_DPF3_CMP_HPP__
#define LIBDPF_INCLUDE_DPF_DPF3_CMP_HPP__

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

#include "hedley/hedley.h"

#include "dpf/dcf.hpp"
#include "dpf/eval_full.hpp"
#include "dpf/eval_interval.hpp"
#include "dpf/eval_sequence.hpp"
#include "dpf/eval_unified.hpp"
#include "dpf/fp61.hpp"
#include "dpf/incremental.hpp"
#include "dpf/interval.hpp"
#include "dpf/path_memoizer.hpp"
#include "dpf/placement.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/shamir3.hpp"
#include "dpf/wildcard.hpp"

namespace dpf
{

/// @brief One party's comparison key: a single DCF share plus a Shamir payload tip.
/// @tparam Party party index in `{1, 2, 3}`
/// @tparam Key two-party DCF key type (party-0 or party-1 half)
template <int Party, typename Key>
struct dpf3_cmp_key
{
    static_assert(Party >= 1 && Party <= 3, "dpf3_cmp party is 1, 2, or 3");
    static constexpr int party = Party;
    static constexpr bool is_dpf3_cmp = true;
    static constexpr bool is_dpf3 = false;

    using input_type = typename Key::input_type;
    using key_type = Key;

    /// @brief Inner comparison half. Eval that accepts a `dpf_key` also accepts
    ///        this object and reads `dpf_key`.
    Key dpf_key{};
    shamir3::share beta_share{};
};

namespace detail
{
namespace dpf3_cmp_impl
{

template <typename K0, typename K1>
auto assign_share_pair(K0 k0, K1 k1, shamir3::share sh, uint64_t if_false_u)
{
    const uint64_t mask = k0.cmp().mask;
    const uint64_t di = sh.value.raw() & mask;
    const uint64_t target =
        detail::incr::cmp_assign_target(k0.cmp(), di, if_false_u & mask);
    uint64_t add0 = 0, add1 = 0;
    const uint64_t r = detail::dcf_impl::sample_addend_blind(mask,
        [] { return dpf::uniform_sample<typename K0::interior_node>(); });
    detail::incr::split_cmp_addend(target, mask, r, add0, add1);
    assign_cmp_local(k0, di, add0);
    assign_cmp_local(k1, di, add1);
    return std::make_tuple(std::move(k0), std::move(k1), sh);
}

template <typename K0, typename K1>
auto wrap_halves(K0 half0, K1 half1, const std::array<shamir3::share, 3> & shares)
{
    K0 half0_copy = half0;  // party 3 holds the same half as party 1
    dpf3_cmp_key<1, K0> out1{std::move(half0), shares[0]};
    dpf3_cmp_key<2, K1> out2{std::move(half1), shares[1]};
    dpf3_cmp_key<3, K0> out3{std::move(half0_copy), shares[2]};
    return std::make_tuple(std::move(out1), std::move(out2), std::move(out3));
}

/// @brief Build three keys that each hold one DCF half of a shared tree.
/// @details One wild comparison tree; δ is planted once into the public value
///          words. Absorb is split once. Each evaluator keeps a single half
///          (Fig-3 style overlap): party 1 → k0, party 2 → k1, party 3 → k0.
///          Eval of one half is that party's complementary share; open with
///          `reconstruct_cmp_halves`. A Shamir tip tracks the payload for
///          bookkeeping / updates.
template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
    typename Spec>
auto make_from_spec(InputT thresh, Spec spec, uint64_t if_true_u,
    uint64_t if_false_u)
{
    using Concrete = uint64_t;
    auto wild = dpf::make_dpf<InteriorPRG, ExteriorPRG>(thresh, spec);
    auto base0 = std::move(wild.first);
    auto base1 = std::move(wild.second);
    using K0 = decltype(base0);

    const uint64_t mask = base0.cmp().mask;
    const uint64_t delta =
        detail::dcf_impl::beta_delta_u64(
            detail::dcf_impl::u64_to_beta<Concrete>(if_true_u),
            detail::dcf_impl::u64_to_beta<Concrete>(if_false_u), mask);
    const auto shares = shamir3::share_secret(fp61{delta});

    // One shared tree; public δ is the clear payload difference. Absorb is
    // split once. Each evaluator keeps a single half (Fig-3 style overlap):
    // party 1 → k0, party 2 → k1, party 3 → k0.
    uint64_t add0 = 0, add1 = 0;
    const uint64_t target =
        detail::incr::cmp_assign_target(base0.cmp(), delta, if_false_u & mask);
    const uint64_t r = detail::dcf_impl::sample_addend_blind(mask,
        [] { return dpf::uniform_sample<typename K0::interior_node>(); });
    detail::incr::split_cmp_addend(target, mask, r, add0, add1);
    assign_cmp_local(base0, delta, add0);
    assign_cmp_local(base1, delta, add1);
    return wrap_halves(std::move(base0), std::move(base1), shares);
}

template <typename KeyT, typename Query, typename PathMemoizer>
std::uint64_t eval_one(const KeyT & key, Query && x, PathMemoizer && path)
{
    const auto y = dpf::eval_point(dpf::cmp, key.dpf_key, std::forward<Query>(x),
        std::forward<PathMemoizer>(path));
    if constexpr (is_secret_share_v<std::decay_t<decltype(y)>>)
        return static_cast<std::uint64_t>(y.raw());
    else
        return static_cast<std::uint64_t>(y);
}

/// @brief Open complementary DCF halves (additive uint64 shares → fp61).
inline fp61 reconstruct_cmp_halves(std::uint64_t k0_share, std::uint64_t k1_share)
{
    return fp61{k0_share + k1_share};
}

}  // namespace dpf3_cmp_impl
}  // namespace detail

/// @brief Open complementary DCF halves (k0-holder with k1-holder).
HEDLEY_WARN_UNUSED_RESULT
inline fp61 reconstruct_cmp_halves(std::uint64_t k0_share, std::uint64_t k1_share)
{
    return detail::dpf3_cmp_impl::reconstruct_cmp_halves(k0_share, k1_share);
}

/// @brief Generate three `lt` comparison keys with a Shamir-shared payload.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          typename InputT>
HEDLEY_WARN_UNUSED_RESULT
auto make_dpf3_cmp(InputT thresh, uint64_t if_true, uint64_t if_false = 0)
{
    return detail::dpf3_cmp_impl::make_from_spec<InteriorPRG, ExteriorPRG>(
        thresh, dpf::lt(dpf::wildcard_value<uint64_t>{}), if_true, if_false);
}

/// @brief Generate three comparison keys from any comparison spec.
/// @details Accepts the same packs as two-party `make_dpf`: `lt`/`leq`/`gt`/`geq`,
///          `*_at`, `idcf`, `block_width`, and path paints. The shape is kept;
///          the payload is installed on a wildcard channel so `update_payload_cmp`
///          can rewrite it. A spec that is already wildcard is left unassigned.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          typename InputT,
          typename Spec,
          typename = std::enable_if_t<is_cmp_spec_v<std::decay_t<Spec>>>>
HEDLEY_WARN_UNUSED_RESULT
auto make_dpf3_cmp(InputT thresh, Spec spec)
{
    using S = std::decay_t<Spec>;
    using Beta = typename S::beta_type;
    using Concrete = dpf::concrete_type_t<Beta>;
    if constexpr (is_wildcard_v<Beta>)
    {
        auto keys = dpf::make_dpf<InteriorPRG, ExteriorPRG>(
            thresh, std::move(spec));
        const auto shares = shamir3::share_secret(fp61{0});
        return detail::dpf3_cmp_impl::wrap_halves(
            std::move(keys.first), std::move(keys.second), shares);
    }
    else
    {
        static_assert(utils::bitlength_of_v<Concrete> <= 64
                && !detail::has_from_seed<Concrete>::value,
            "dpf3 comparison payloads use the uint64 ring");
        const uint64_t t = detail::dcf_impl::beta_to_u64_simple(spec.if_true,
            ~uint64_t{0});
        const uint64_t f = detail::dcf_impl::beta_to_u64_simple(spec.if_false,
            ~uint64_t{0});
        return detail::dpf3_cmp_impl::make_from_spec<InteriorPRG, ExteriorPRG>(
            thresh, cmp_spec_as_wildcard(std::move(spec)), t, f);
    }
}

/// @brief Blocked comparison with Shamir-shared value words.
template <std::size_t B,
          typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          typename InputT>
HEDLEY_WARN_UNUSED_RESULT
auto make_dpf3_cmp_blocked(InputT thresh, uint64_t if_true,
    uint64_t if_false = 0)
{
    return detail::dpf3_cmp_impl::make_from_spec<InteriorPRG, ExteriorPRG>(
        thresh, dpf::block_width<B>(dpf::lt(dpf::wildcard_value<uint64_t>{})),
        if_true, if_false);
}

/// @brief Evaluate a three-party comparison key at `x` (one DCF share).
/// @details Returns an unreduced additive `uint64` share. Open a complementary
///          pair with `reconstruct_cmp_halves` (sum, then reduce into `fp61`).
template <typename KeyT, typename Query,
          typename PathMemoizer = basic_path_memoizer<typename KeyT::key_type>,
          std::enable_if_t<std::decay_t<KeyT>::is_dpf3_cmp, int> = 0>
HEDLEY_WARN_UNUSED_RESULT
std::uint64_t eval_dpf3_cmp(const KeyT & key, Query && x,
    PathMemoizer && path = PathMemoizer{})
{
    return detail::dpf3_cmp_impl::eval_one(key, std::forward<Query>(x),
        std::forward<PathMemoizer>(path));
}

/// @brief Full-domain comparison into an additive-share buffer (one DCF expand).
/// \complexity Same expansion as `eval_interval` on the whole domain. L = 2^{n - lg(outputs_per_leaf)} leaf nodes, n = `depth`. Time Θ(L) interior traversals. The output buffer stores one slot per domain point (2^n).
template <typename KeyT,
          std::enable_if_t<std::decay_t<KeyT>::is_dpf3_cmp, int> = 0>
HEDLEY_WARN_UNUSED_RESULT
std::vector<std::uint64_t> eval_full(const KeyT & key)
{
    using in = typename KeyT::input_type;
    const auto from = std::numeric_limits<in>::min();
    const auto to = std::numeric_limits<in>::max();
    auto buf = dpf::make_output_buffer(dpf::cmp, key.dpf_key, from, to);
    dpf::eval_interval(dpf::cmp, key.dpf_key, from, to, buf);
    std::vector<std::uint64_t> out(buf.size());
    for (std::size_t i = 0; i < buf.size(); ++i)
    {
        if constexpr (is_secret_share_v<std::decay_t<decltype(buf[i])>>)
            out[i] = static_cast<std::uint64_t>(buf[i].raw());
        else
            out[i] = static_cast<std::uint64_t>(buf[i]);
    }
    return out;
}

/// @brief Interval comparison into an additive-share buffer.
/// \complexity O(L) interior traversals and O(L) workspace in the basic memoizer. L is the number of leaf nodes covering the closed interval (`get_nodes_at_level` at `depth`). Level k expands `(to >> (n-k)) - (from >> (n-k)) + 1` nodes; those counts sum to Θ(L). The output buffer holds one slot per input in the interval. n is `depth`.
template <typename KeyT, typename LaneT,
          std::enable_if_t<std::decay_t<KeyT>::is_dpf3_cmp, int> = 0>
HEDLEY_WARN_UNUSED_RESULT
std::vector<std::uint64_t> eval_interval(const KeyT & key, LaneT from, LaneT to)
{
    auto buf = dpf::make_output_buffer(dpf::cmp, key.dpf_key, from, to);
    dpf::eval_interval(dpf::cmp, key.dpf_key, from, to, buf);
    std::vector<std::uint64_t> out(buf.size());
    for (std::size_t i = 0; i < buf.size(); ++i)
    {
        if constexpr (is_secret_share_v<std::decay_t<decltype(buf[i])>>)
            out[i] = static_cast<std::uint64_t>(buf[i].raw());
        else
            out[i] = static_cast<std::uint64_t>(buf[i]);
    }
    return out;
}

/// @brief In-place comparison payload update via the linear value-CW channel.
/// @details Applies `assign_cmp_local` with each party's Shamir share of
///          `δ' − δ`, adding that increment into the existing value CWs and
///          absorb addends. The threshold / spine stays. `if_false` is fixed
///          at 0 (the common `lt(β)` case); changing `if_false` needs a remake.
template <typename K1, typename K2, typename K3>
void update_payload_cmp(K1 & k1, K2 & k2, K3 & k3, uint64_t if_true_old,
    uint64_t if_true_new)
{
    static_assert(K1::is_dpf3_cmp && K2::is_dpf3_cmp && K3::is_dpf3_cmp,
        "dpf3_cmp keys");
    const uint64_t mask = k1.dpf_key.cmp().mask;
    // Clear δ on the shared tree; difference is taken in fp61.
    const uint64_t d =
        (fp61{if_true_new} - fp61{if_true_old}).raw() & mask;
    const uint64_t target = detail::incr::cmp_assign_target(k1.dpf_key.cmp(), d, 0);
    uint64_t add0 = 0, add1 = 0;
    const uint64_t r = detail::dcf_impl::sample_addend_blind(mask,
        [] {
            return dpf::uniform_sample<
                typename std::decay_t<decltype(k1.dpf_key)>::interior_node>();
        });
    detail::incr::split_cmp_addend(target, mask, r, add0, add1);
    auto bump = [&](auto & key, uint64_t add) {
        const uint64_t old = [&] {
            if constexpr (is_party_key_v<std::decay_t<decltype(key.dpf_key)>>)
                return static_cast<uint64_t>(key.dpf_key.cmp_addend().raw());
            else if constexpr (std::is_integral_v<
                                   std::decay_t<decltype(key.dpf_key.cmp_addend())>>)
                return static_cast<uint64_t>(key.dpf_key.cmp_addend());
            else
                return static_cast<uint64_t>(key.dpf_key.cmp_addend().raw());
        }();
        assign_cmp_local(key.dpf_key, d, (old + add) & mask);
    };
    bump(k1, add0);
    bump(k2, add1);
    bump(k3, add0);
    const auto shares = shamir3::share_secret(
        fp61{if_true_new} - fp61{if_true_old});
    k1.beta_share = shamir3::share{k1.party, k1.beta_share.value + shares[0].value};
    k2.beta_share = shamir3::share{k2.party, k2.beta_share.value + shares[1].value};
    k3.beta_share = shamir3::share{k3.party, k3.beta_share.value + shares[2].value};
}

/// @brief Fresh comparison keys for a new payload (new spines — not an update).
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          typename InputT>
HEDLEY_WARN_UNUSED_RESULT
auto remake_dpf3_cmp(InputT thresh, uint64_t if_true_new,
    uint64_t if_false_new = 0)
{
    return make_dpf3_cmp<InteriorPRG, ExteriorPRG>(thresh, if_true_new,
        if_false_new);
}

/// @brief One party's interval key on a three-party comparison.
/// @tparam Party party index in `{1, 2, 3}`
/// @tparam CmpKey a `dpf3_cmp_key`
template <int Party, typename CmpKey>
struct dpf3_ic_key
{
    static constexpr int party = Party;
    static constexpr bool is_dpf3_ic = true;
    static constexpr bool is_dpf3 = false;
    using input_type = typename CmpKey::input_type;
    using cmp_key_type = CmpKey;

    /// @brief Inner three-party comparison. Its `dpf_key` is the DPF half.
    CmpKey dpf_key{};
    uint64_t lo = 0;  // public bounds (F_IC)
    uint64_t hi = 0;
    uint64_t input_mask = 0;
    uint64_t group_mask = 0;
    /// Additive half of δ (k0-holder / k1-holder), opened with `reconstruct_cmp_halves`.
    std::uint64_t delta_share = 0;
    std::uint64_t cr_share = 0;
};

/// @brief Generate three interval-containment keys.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          typename InputT>
HEDLEY_WARN_UNUSED_RESULT
auto make_dpf3_ic(InputT r, InputT p, InputT q, uint64_t if_true,
    uint64_t if_false = 0)
{
    auto two = dpf::make_dpf(r, dpf::ic(p, q, if_true, if_false));
    const uint64_t nmask = two.first.input_mask;
    const uint64_t gmask = two.first.group_mask;
    const uint64_t lo = static_cast<uint64_t>(p);
    const uint64_t hi = static_cast<uint64_t>(q);
    const InputT gamma = detail::ic_impl::gamma_of(r);

    const uint64_t delta = (if_true - if_false) & gmask;
    auto cmp_keys = make_dpf3_cmp<InteriorPRG, ExteriorPRG>(gamma, delta, 0);

    uint64_t cr0 = 0, cr1 = 0;
    if constexpr (is_secret_share_v<decltype(two.first.cr_share)>)
    {
        cr0 = static_cast<uint64_t>(two.first.cr_share.raw());
        cr1 = static_cast<uint64_t>(two.second.cr_share.raw());
    }
    else
    {
        cr0 = static_cast<uint64_t>(two.first.cr_share);
        cr1 = static_cast<uint64_t>(two.second.cr_share);
    }
    const uint64_t cr_clear = (cr0 + cr1) & gmask;
    // Additive halves so reconstruct_cmp_halves (sum) recovers the clear words.
    const uint64_t d0 = dpf::uniform_sample<std::uint64_t>() & gmask;
    const uint64_t d1 = (delta - d0) & gmask;
    const uint64_t c0 = dpf::uniform_sample<std::uint64_t>() & gmask;
    const uint64_t c1 = (cr_clear - c0) & gmask;

    using Cmp1 = std::decay_t<decltype(std::get<0>(cmp_keys))>;
    using Cmp2 = std::decay_t<decltype(std::get<1>(cmp_keys))>;
    using Cmp3 = std::decay_t<decltype(std::get<2>(cmp_keys))>;

    dpf3_ic_key<1, Cmp1> k1{std::move(std::get<0>(cmp_keys)), lo, hi, nmask,
        gmask, d0, c0};
    dpf3_ic_key<2, Cmp2> k2{std::move(std::get<1>(cmp_keys)), lo, hi, nmask,
        gmask, d1, c1};
    dpf3_ic_key<3, Cmp3> k3{std::move(std::get<2>(cmp_keys)), lo, hi, nmask,
        gmask, d0, c0};
    return std::make_tuple(std::move(k1), std::move(k2), std::move(k3));
}

/// @brief Evaluate a three-party interval key at `x`.
template <typename KeyT, typename Query,
          std::enable_if_t<std::decay_t<KeyT>::is_dpf3_ic, int> = 0>
HEDLEY_WARN_UNUSED_RESULT
std::uint64_t eval_dpf3_ic(const KeyT & key, Query && x)
{
    using in_type = typename KeyT::input_type;
    const uint64_t xu = detail::ic_impl::bits_of(in_type(x));
    const uint64_t xp = detail::ic_impl::shift_p(xu, key.lo, key.input_mask);
    const uint64_t xq = detail::ic_impl::shift_q0(xu, key.hi, key.input_mask);
    basic_path_memoizer<typename KeyT::cmp_key_type::key_type> path;
    const uint64_t a = eval_dpf3_cmp(key.dpf_key,
        detail::ic_impl::input_from_bits<in_type>(xp), path);
    const uint64_t b = eval_dpf3_cmp(key.dpf_key,
        detail::ic_impl::input_from_bits<in_type>(xq), path);
    const int cx = detail::ic_impl::public_cx(xu, key.lo, key.hi, key.input_mask);
    uint64_t scaled = 0;
    if (cx == 1)
        scaled = key.delta_share;
    else if (cx == -1)
        scaled = static_cast<uint64_t>(0) - key.delta_share;
    return (static_cast<uint64_t>(0) - a) + b + key.cr_share + scaled;
}

/// @brief `eval_point` overload for three-party comparison keys.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <typename KeyT, typename Query,
          typename PathMemoizer = basic_path_memoizer<typename KeyT::key_type>,
          std::enable_if_t<std::decay_t<KeyT>::is_dpf3_cmp, int> = 0>
HEDLEY_WARN_UNUSED_RESULT
std::uint64_t eval_point(const KeyT & key, Query && x,
    PathMemoizer && path = PathMemoizer{})
{
    return eval_dpf3_cmp(key, std::forward<Query>(x),
        std::forward<PathMemoizer>(path));
}

/// @brief `eval_point` overload for three-party interval keys.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <typename KeyT, typename Query,
          std::enable_if_t<std::decay_t<KeyT>::is_dpf3_ic, int> = 0>
HEDLEY_WARN_UNUSED_RESULT
std::uint64_t eval_point(const KeyT & key, Query && x)
{
    return eval_dpf3_ic(key, std::forward<Query>(x));
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_DPF3_CMP_HPP__