/// @file party/dist_dpf3.hpp
/// @brief Networked (2,3) Shamir DPF keygen (three key holders).
/// @details A pure two-party variant is not meaningful: after keygen each of
///          three parties must hold a Shamir share of the evaluation. Use
///          `dist_with_*_iknp` in `iknp_deal.hpp` for two-party VDPF keygen
///          without a pad dealer.
/// @brief Distributed dual-spine Doerner–Shelat keygen for (2,3) point keys.
/// @details Two `deal_point` / `point_party` runs (spines A and B) with Fig-3
///          `τ` payloads. p0 samples `τ` and `π`; overlapping halves are shipped
///          so p0 holds party-1, p2 holds party-2, p1 holds party-3.
///          Produced keys are always verifiable. Default spines keep `α` as XOR
///          shares (F_DPF3DS): `π` is peeled from the leaf seed of the shared
///          path, not from an opened point. p0 sends p1 only its `τ` halves.
///          Pass `RevealPoint=true` only when the caller wants the tree prefix
///          (and packed lane on updatable keys). Pass `dpf::updatable` for
///          beaver leaves and a later networked Fig-10 update via
///          `dist_update_payload`.

#ifndef LIBDPF_PARTY_DIST_DPF3_HPP__
#define LIBDPF_PARTY_DIST_DPF3_HPP__

#include <cstring>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include "hedley/hedley.h"

#include "dist_ds.hpp"
#include "key_io.hpp"

#include "dpf/dpf3.hpp"
#include "dpf/dpf3_ds.hpp"
#include "dpf/fp61.hpp"
#include "dpf/shamir3.hpp"
#include "dpf/wildcard.hpp"

namespace dpf
{
namespace party
{

/// @brief How trio roles map onto Shamir party indices `{1,2,3}`.
/// @details **dealer**: p0→1, p1→2, p2→3 (trusted keygen on p2, then ship).
///          **dist**: p0→1, p2→2, p1→3 (`dist_with_dpf3_key` assembly).
enum class dpf3_role_map : unsigned
{
    dealer = 0,
    dist = 1,
};

/// @brief Shamir party index for a trio role under `map`.
HEDLEY_CONST
HEDLEY_NO_THROW
constexpr int dpf3_party_of(role r, dpf3_role_map map) noexcept
{
    if (map == dpf3_role_map::dist)
    {
        if (r == role::p0)
            return 1;
        if (r == role::p2)
            return 2;
        return 3;  // p1
    }
    // dealer
    if (r == role::p0)
        return 1;
    if (r == role::p1)
        return 2;
    return 3;  // p2
}

/// @brief Public VDPF+ offsets for one (2,3) keygen.
struct dpf3_pi_msg
{
    shamir3::xor61 pi_a{};
    shamir3::xor61 pi_b{};
};

/// @brief Fig-10 patch broadcast (leaf patches so p2 need not learn `α`).
template <typename Leaf>
struct dpf3_fig10_msg
{
    Leaf patch_a{};
    Leaf patch_b{};
    shamir3::xor61 pi_a{};
    shamir3::xor61 pi_b{};
};

/// @brief p1's τ halves for spines A and B (party-3 strings). Does not reveal β.
struct dpf3_tau_share
{
    shamir3::xor61 t1{};
    shamir3::xor61 t3{};
};

namespace detail_dist_dpf3
{

template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
    typename OutputT>
using spine_out0 = decltype(dist::point_party<InteriorPRG, ExteriorPRG, InputT,
    OutputT, role::p0>(std::declval<trio &>(), std::declval<InputT>(),
    std::declval<OutputT>()));

template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
    typename OutputT>
using spine_key0 = spine_out0<InteriorPRG, ExteriorPRG, InputT, OutputT>;

template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
    typename OutputT>
using spine_out1 = decltype(dist::point_party<InteriorPRG, ExteriorPRG, InputT,
    OutputT, role::p1>(std::declval<trio &>(), std::declval<InputT>(),
    std::declval<OutputT>()));

template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
    typename OutputT>
using spine_key1 = spine_out1<InteriorPRG, ExteriorPRG, InputT, OutputT>;

template <typename Key, typename Share>
void assign_wildcard_over_link(Key & key, Share my_share, trio & net, role self)
{
    const role peer = self == role::p0 ? role::p1 : role::p0;
    auto & wrap = std::get<0>(key.leaf_nodes);
    auto blinded = wrap.compute_and_get_blinded_output_share(my_share);
    auto peer_blinded = net.exchange_with(peer, blinded);
    auto leaf = wrap.compute_and_get_leaf_share(peer_blinded);
    auto peer_leaf = net.exchange_with(peer, leaf);
    wrap.reconstruct_correction_word(peer_leaf);
}

/// @brief Assign XOR payload from each party's τ half (payload never opened).
template <typename Key>
void assign_spine_tau_halves(trio & net, role self, Key & key,
    shamir3::xor61 my_half)
{
    assign_wildcard_over_link(key, my_half, net, self);
}

template <typename Key, typename Leaf>
void apply_fig10_leaves(Key & key, const dpf3_fig10_msg<Leaf> & msg)
{
    detail::dpf3_impl::apply_leaf_patch(key.a.dpf_key, msg.patch_a);
    detail::dpf3_impl::apply_leaf_patch(key.b.dpf_key, msg.patch_b);
}

template <typename Key, typename Leaf>
void apply_fig10(Key & key, const dpf3_fig10_msg<Leaf> & msg)
{
    apply_fig10_leaves(key, msg);
    key.a.offset = msg.pi_a;
    key.b.offset = msg.pi_b;
}

template <typename Key, typename Node>
HEDLEY_WARN_UNUSED_RESULT
shamir3::xor61 peel_from_seed(const Key & key, const Node & seed)
{
    const auto y = key.template traverse_exterior<0>(seed);
    using Y = std::decay_t<decltype(y)>;
    if constexpr (is_secret_share_v<Y>)
        return shamir3::xor61{y.raw()};
    else if constexpr (std::is_integral_v<Y>
        || std::is_convertible_v<Y, std::uint64_t>)
        return shamir3::xor61{static_cast<std::uint64_t>(y)};
    else
    {
        std::uint64_t w = 0;
        static_assert(sizeof(Y) >= sizeof(w),
            "peel_from_seed: leaf narrower than xor61");
        std::memcpy(&w, &y, sizeof(w));
        return shamir3::xor61{w};
    }
}

template <typename InputT>
struct dpf3_opened
{
    InputT opened_prefix{};
    /// Meaningful when the spine payload is a packed wildcard.
    unsigned opened_lane = 0;
    bool lane_opened = false;
};

template <typename T, typename = void>
struct has_nested_key : std::false_type {};
template <typename T>
struct has_nested_key<T, std::void_t<decltype(std::declval<T &>().key)>>
    : std::true_type {};
template <typename T>
inline constexpr bool has_nested_key_v = has_nested_key<T>::value;

template <typename T, typename = void>
struct has_opened_prefix : std::false_type {};
template <typename T>
struct has_opened_prefix<T,
    std::void_t<decltype(std::declval<const T &>().opened_prefix)>>
    : std::true_type {};
template <typename T>
inline constexpr bool has_opened_prefix_v = has_opened_prefix<T>::value;

template <typename T, typename = void>
struct has_opened_lane : std::false_type {};
template <typename T>
struct has_opened_lane<T,
    std::void_t<decltype(std::declval<const T &>().opened_lane)>>
    : std::true_type {};
template <typename T>
inline constexpr bool has_opened_lane_v = has_opened_lane<T>::value;

template <typename T, typename = void>
struct has_member_dpf_key : std::false_type {};
template <typename T>
struct has_member_dpf_key<T, std::void_t<decltype(std::declval<T &>().dpf_key)>>
    : std::true_type {};

template <typename Held>
auto spine_key_of(Held && held)
{
    using H = std::decay_t<Held>;
    if constexpr (has_member_dpf_key<H>::value)
        return std::move(held.dpf_key);
    else if constexpr (has_nested_key_v<H>)
        return std::move(held.key);
    else
        return std::move(held);
}

template <typename Held, typename InputT>
void note_opened(dpf3_opened<InputT> & learned, const Held & held)
{
    if constexpr (has_opened_prefix_v<std::decay_t<Held>>)
    {
        learned.opened_prefix = held.opened_prefix;
        if constexpr (has_opened_lane_v<std::decay_t<Held>>)
        {
            learned.opened_lane = held.opened_lane;
            learned.lane_opened = true;
        }
    }
}

template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
    typename OutputT, bool Updatable, bool RevealPoint, typename Fn1,
    typename Fn2, typename Fn3>
std::optional<dpf3_opened<InputT>> dist_with_dpf3_key_impl(trio & net, role self,
    InputT x0, InputT x1,
    const fp61 beta, Fn1 && on1, Fn2 && on2, Fn3 && on3)
{
    using X = shamir3::xor61;
    using tau_quad = detail::dpf3_impl::tau_quad;
    using Key0 = spine_key0<InteriorPRG, ExteriorPRG, InputT, OutputT>;
    using Key1 = spine_key1<InteriorPRG, ExteriorPRG, InputT, OutputT>;
    using node = typename dpf::tree_traits<InteriorPRG>::node;
    constexpr bool V = true;  // dist point_party always opens correction seeds

    if (self == role::p2)
    {
        dist::deal_point<InteriorPRG, InputT, OutputT, !RevealPoint>(net);
        dist::deal_point<InteriorPRG, InputT, OutputT, !RevealPoint>(net);
        const auto pi = net.recv_from<dpf3_pi_msg>(role::p0, net::msg::delta);
        auto key_b0 = recv_key<Key0>(net, role::p0);
        auto key_a1 = recv_key<Key1>(net, role::p1);
        detail::dpf3_impl::vdpf_plus_key<Key1> plus_a1{std::move(key_a1),
            pi.pi_a};
        detail::dpf3_impl::vdpf_plus_key<Key0> plus_b0{std::move(key_b0),
            pi.pi_b};
        dpf3_key<2, decltype(plus_a1), decltype(plus_b0)> k2{
            std::move(plus_a1), std::move(plus_b0), V, false, Updatable};
        std::forward<Fn2>(on2)(std::move(k2));
        return std::nullopt;
    }

    // p0 samples τ from a Shamir split of β; p1 receives only (t1, t3).
    tau_quad t{};
    X my_a{};
    X my_b{};
    if (self == role::p0)
    {
        t = detail::dpf3_impl::sample_taus(beta);
        my_a = t.t0;
        my_b = t.t2;
        net.send_to(role::p1, net::msg::delta, dpf3_tau_share{t.t1, t.t3});
    }
    else
    {
        const auto sh = net.recv_from<dpf3_tau_share>(role::p0, net::msg::delta);
        my_a = sh.t1;
        my_b = sh.t3;
        t.t1 = sh.t1;
        t.t3 = sh.t3;
    }

    dpf3_opened<InputT> learned{};
    node seed_a{};
    node seed_b{};

    // Non-updatable and updatable both plant beaver leaves; τ halves are
    // assigned without opening the spine payload. The outer Updatable flag
    // only gates Fig-10.
    // Spines still use per-message `exchange_with`. Wiring RoundSink here
    // needs a round budget that covers wild+hash leaf mux; undersizing hangs
    // the peer on `msg::round_batch`.
    OutputT wild{};
    if (self == role::p0)
    {
        auto out_a0 = dist::point_party<InteriorPRG, ExteriorPRG, InputT,
            OutputT, role::p0, false, RevealPoint>(net, x0, wild, false,
            &seed_a);
        note_opened(learned, out_a0);
        auto key_a0 = spine_key_of(std::move(out_a0));
        auto out_b0 = dist::point_party<InteriorPRG, ExteriorPRG, InputT,
            OutputT, role::p0, false, RevealPoint>(net, x0, wild, false,
            &seed_b);
        auto key_b0 = spine_key_of(std::move(out_b0));
        assign_spine_tau_halves(net, self, key_a0, my_a);
        assign_spine_tau_halves(net, self, key_b0, my_b);
        dpf3_pi_msg pi{};
        pi.pi_a = t.t0 + peel_from_seed(key_a0, seed_a);
        pi.pi_b = t.t2 + peel_from_seed(key_b0, seed_b);
        net.send_to(role::p1, net::msg::delta, pi);
        net.send_to(role::p2, net::msg::delta, pi);
        send_key(net, role::p2, key_b0);
        detail::dpf3_impl::vdpf_plus_key<Key0> plus_a0{std::move(key_a0),
            pi.pi_a};
        detail::dpf3_impl::vdpf_plus_key<Key0> plus_b0{std::move(key_b0),
            pi.pi_b};
        dpf3_key<1, decltype(plus_a0), decltype(plus_b0)> k1{
            std::move(plus_a0), std::move(plus_b0), V, false, Updatable};
        std::forward<Fn1>(on1)(std::move(k1));
        if constexpr (RevealPoint)
            return learned;
        return std::nullopt;
    }

    auto out_a1 = dist::point_party<InteriorPRG, ExteriorPRG, InputT,
        OutputT, role::p1, false, RevealPoint>(net, x1, wild, false,
        &seed_a);
    note_opened(learned, out_a1);
    auto key_a1 = spine_key_of(std::move(out_a1));
    auto out_b1 = dist::point_party<InteriorPRG, ExteriorPRG, InputT,
        OutputT, role::p1, false, RevealPoint>(net, x1, wild, false,
        &seed_b);
    auto key_b1 = spine_key_of(std::move(out_b1));
    assign_spine_tau_halves(net, self, key_a1, my_a);
    assign_spine_tau_halves(net, self, key_b1, my_b);
    const auto pi = net.recv_from<dpf3_pi_msg>(role::p0, net::msg::delta);
    send_key(net, role::p2, key_a1);
    detail::dpf3_impl::vdpf_plus_key<Key1> plus_a1{std::move(key_a1),
        pi.pi_a};
    detail::dpf3_impl::vdpf_plus_key<Key1> plus_b1{std::move(key_b1),
        pi.pi_b};
    dpf3_key<3, decltype(plus_a1), decltype(plus_b1)> k3{
        std::move(plus_a1), std::move(plus_b1), V, false, Updatable};
    std::forward<Fn3>(on3)(std::move(k3));
    if constexpr (RevealPoint)
        return learned;
    return std::nullopt;
}

}  // namespace detail_dist_dpf3

/// @brief Distributed dual-spine (2,3) keygen over the trio.
/// @details Role map after assembly: **p0 → party 1**, **p2 → party 2**,
///          **p1 → party 3**. Keys are always verifiable. Each computing
///          party's view of `α` is only its XOR share; p1's view of `β` is
///          only its τ halves (not a clear payload). Default return is empty.
/// @tparam RevealPoint when true, p0/p1 also reconstruct the tree prefix
/// @return Opened prefix when `RevealPoint`, else empty. Empty on p2.
/// @throws std::runtime_error if a frame is truncated or tagged wrong
/// \complexity Two `point_party` spines, so the local work is two O(n) walks, plus O(1) τ arithmetic.
/// \rounds The rounds of two `point_party` calls, then one τ-share send (p0 to p1), one `dpf3_pi_msg` to p1 and p2, and one key send to p2. p2 runs two `deal_point` calls first. Counted in `dist_with_dpf3_key_impl`. Fig-10 update is `dist_update_payload`, not this function.
/// \communication Two dealer tapes (see `deal_point` / `point_party`), one `dpf3_tau_share`, one `dpf3_pi_msg` (two `xor61` values), and one key blob (`send_key`, `sizeof` of the spine key).
/// \preprocessing p2's `deal_point` pads for both spines. p0 samples the four τ strings locally from a Shamir split of β.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          bool RevealPoint = false,
          typename InputT,
          typename Fn1,
          typename Fn2,
          typename Fn3>
[[nodiscard]] std::optional<InputT> dist_with_dpf3_key(trio & net, role self,
    InputT x0, InputT x1, const fp61 beta, Fn1 && on1, Fn2 && on2, Fn3 && on3)
{
    auto opened = detail_dist_dpf3::dist_with_dpf3_key_impl<InteriorPRG,
        ExteriorPRG, InputT, dpf::wildcard_value<shamir3::xor61>, false,
        RevealPoint>(net, self, x0, x1, beta, std::forward<Fn1>(on1),
        std::forward<Fn2>(on2), std::forward<Fn3>(on3));
    if (!opened)
        return std::nullopt;
    return opened->opened_prefix;
}

/// @brief What an updatable (2,3) keygen reconstructs when `RevealPoint`.
template <typename InputT>
struct dpf3_updatable_opened
{
    InputT opened_prefix{};
    unsigned opened_lane = 0;
};

/// @brief Distributed dual-spine (2,3) keygen with beaver leaves (Fig-10-ready).
/// @tparam RevealPoint when true, return prefix and packed lane on p0/p1
/// @return Opened values when `RevealPoint`, else empty. Empty on p2.
/// \complexity Two `point_party` spines, so the local work is two O(n) walks, plus O(1) τ arithmetic.
/// \rounds The rounds of two `point_party` calls, then one τ-share send (p0 to p1), one `dpf3_pi_msg` to p1 and p2, and one key send to p2. p2 runs two `deal_point` calls first. Counted in `dist_with_dpf3_key_impl`. Fig-10 update is `dist_update_payload`, not this function.
/// \communication Two dealer tapes (see `deal_point` / `point_party`), one `dpf3_tau_share`, one `dpf3_pi_msg` (two `xor61` values), and one key blob (`send_key`, `sizeof` of the spine key).
/// \preprocessing p2's `deal_point` pads for both spines. p0 samples the four τ strings locally from a Shamir split of β.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          bool RevealPoint = false,
          typename InputT,
          typename Fn1,
          typename Fn2,
          typename Fn3>
[[nodiscard]] std::optional<dpf3_updatable_opened<InputT>> dist_with_dpf3_key(
    trio & net, role self, InputT x0, InputT x1, const fp61 beta, updatable,
    Fn1 && on1, Fn2 && on2, Fn3 && on3)
{
    auto opened = detail_dist_dpf3::dist_with_dpf3_key_impl<InteriorPRG,
        ExteriorPRG, InputT, dpf::wildcard_value<shamir3::xor61>, true,
        RevealPoint>(net, self, x0, x1, beta, std::forward<Fn1>(on1),
        std::forward<Fn2>(on2), std::forward<Fn3>(on3));
    if (!opened)
        return std::nullopt;
    return dpf3_updatable_opened<InputT>{opened->opened_prefix,
        opened->opened_lane};
}

/// @brief Networked Fig-10 payload update for keys from `dist_with_dpf3_key`.
/// @details p0 and p1 (who are given `α`) exchange peels, p0 samples fresh `τ`
///          and broadcasts leaf patches + new public `π`. p2 applies patches
///          without learning `α`. Requires `key.updatable`.
/// @throws std::invalid_argument if the key is not updatable
template <typename Key, typename InputT>
void dist_update_payload(trio & net, role self, Key & key, InputT alpha,
    fp61 beta_new)
{
    static_assert(Key::is_dpf3, "dist_update_payload: dpf3 key");
    if (!key.updatable)
        throw std::invalid_argument(
            "dist_update_payload: key was not generated with dpf::updatable");

    using X = shamir3::xor61;
    using InnerA = typename Key::plus_a_type::inner_type;
    using Leaf = decltype(detail::dpf3_impl::make_leaf_patch<InnerA>(alpha,
        X{}));
    using Msg = dpf3_fig10_msg<Leaf>;

    struct peel_pair
    {
        X peel_a{};
        X peel_b{};
    };

    Msg msg{};
    if (self == role::p0 || self == role::p1)
    {
        peel_pair mine{detail::dpf3_impl::peel(key.a.dpf_key, alpha),
            detail::dpf3_impl::peel(key.b.dpf_key, alpha)};
        const role peer = self == role::p0 ? role::p1 : role::p0;
        const peel_pair theirs = net.exchange_with(peer, mine);

        if (self == role::p0)
        {
            // p0 = party 1 (A0,B0); p1 = party 3 (A1,B1).
            // Leaf CW patches only move the half whose path control bit is set;
            // compute π from a post-patch peel (same as dealer refresh_offset).
            detail::dpf3_impl::tau_quad told{};
            told.t0 = mine.peel_a + key.a.offset;
            told.t1 = theirs.peel_a + key.a.offset;
            told.t2 = mine.peel_b + key.b.offset;
            told.t3 = theirs.peel_b + key.b.offset;
            const auto tnew = detail::dpf3_impl::sample_taus(beta_new);
            const X dA = (tnew.t0 + tnew.t1) + (told.t0 + told.t1);
            const X dB = (tnew.t2 + tnew.t3) + (told.t2 + told.t3);
            msg.patch_a =
                detail::dpf3_impl::make_leaf_patch<InnerA>(alpha, dA);
            msg.patch_b =
                detail::dpf3_impl::make_leaf_patch<InnerA>(alpha, dB);
            detail_dist_dpf3::apply_fig10_leaves(key, msg);
            msg.pi_a = tnew.t0 + detail::dpf3_impl::peel(key.a.dpf_key, alpha);
            msg.pi_b = tnew.t2 + detail::dpf3_impl::peel(key.b.dpf_key, alpha);
            key.a.offset = msg.pi_a;
            key.b.offset = msg.pi_b;
            net.send_to(role::p1, net::msg::delta, msg);
            net.send_to(role::p2, net::msg::delta, msg);
            return;
        }
        msg = net.recv_from<Msg>(role::p0, net::msg::delta);
        detail_dist_dpf3::apply_fig10(key, msg);
        return;
    }

    msg = net.recv_from<Msg>(role::p0, net::msg::delta);
    detail_dist_dpf3::apply_fig10(key, msg);
}

}  // namespace party
}  // namespace dpf

#endif  // LIBDPF_PARTY_DIST_DPF3_HPP__