/// @file party/iknp_deal.hpp
/// @brief Two-party IKNP setup for every tape `dist_ds.hpp` takes from p2.
/// @details p0 and p1 sample the same pad correlations the dealer would, write
///          them into a local inbox, and then run the existing walk. p2 is not
///          connected. Covered tapes: additive carry, point (reveal and
///          oblivious hash, wildcard, half-tree, extractable), comparison
///          (edge, value-word, blocked suffix, verifiable hash), and interval
///          containment.

#ifndef LIBDPF_PARTY_IKNP_DEAL_HPP__
#define LIBDPF_PARTY_IKNP_DEAL_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

#include "hedley/hedley.h"

#include "dist_ds.hpp"
#include "dpf/iknp.hpp"

namespace dpf
{
namespace party
{
namespace iknp_deal
{

using net::role;
using net::trio;

struct tape
{
    std::vector<std::uint8_t> bytes;

    template <typename T>
    void put(net::msg tag, const T & value)
    {
        static_assert(std::is_trivially_copyable_v<T>);
        append(tag, &value, sizeof(T));
    }

    template <typename T>
    void put_vec(net::msg tag, const T * data, std::size_t n)
    {
        static_assert(std::is_trivially_copyable_v<T>);
        append(tag, data, n * sizeof(T));
    }

    void append(net::msg tag, const void * payload, std::size_t n)
    {
        if (n > 0xffffffffu)
            throw std::invalid_argument("iknp tape frame");
        const auto len = static_cast<std::uint32_t>(n);
        const auto t = net::to_u16(tag);
        std::uint8_t hdr[6];
        std::memcpy(hdr, &len, 4);
        std::memcpy(hdr + 4, &t, 2);
        bytes.insert(bytes.end(), hdr, hdr + 6);
        if (n != 0)
        {
            const auto * p = static_cast<const std::uint8_t *>(payload);
            bytes.insert(bytes.end(), p, p + n);
        }
    }
};

struct cursor
{
    const iknp::material * mat = nullptr;
    std::size_t block = 0;
    std::size_t bit = 0;
    std::size_t b2a = 0;
    std::size_t cw = 0;
    std::size_t zero = 0;
    std::vector<dist::ring_zero_share_msg> zeros;

    dist::and_share_msg take_block()
    {
        const auto & s = mat->blocks.at(block++);
        dist::and_share_msg m;
        m.a = s.a;
        m.b = s.b;
        m.c = s.c;
        return m;
    }

    dist::bit_and_pad_msg take_bit()
    {
        const auto & s = mat->bits.at(bit++);
        return dist::bit_and_pad_msg{s.a, s.b, s.c};
    }

    dist::b2a_pad_msg take_b2a()
    {
        const auto & s = mat->b2a.at(b2a++);
        return dist::b2a_pad_msg{s.r, s.add};
    }

    dist::cw_pad_msg take_cw()
    {
        const auto & s = mat->cws.at(cw++);
        dist::cw_pad_msg m;
        m.rand = s.rand;
        m.gamma = s.gamma;
        m.bit = s.bit;
        return m;
    }

    dist::ring_zero_share_msg take_zero()
    {
        return zeros.at(zero++);
    }

    void done() const
    {
        if (mat == nullptr)
            return;
        if (block != mat->blocks.size() || bit != mat->bits.size()
            || b2a != mat->b2a.size() || cw != mat->cws.size()
            || zero != zeros.size())
            throw std::runtime_error("iknp tape did not consume every pad");
    }
};

inline int me_of(role self)
{
    if (self == role::p0)
        return 0;
    if (self == role::p1)
        return 1;
    throw std::invalid_argument("iknp parties are p0 and p1");
}

inline void exchange_zeros(trio & net, int me,
    const std::vector<std::uint64_t> & masks, cursor & cur)
{
    cur.zeros.resize(masks.size());
    const role peer = me == 0 ? role::p1 : role::p0;
    for (std::size_t i = 0; i < masks.size(); ++i)
    {
        if (me == 0)
        {
            const auto r = dist::random_ring_word(masks[i]);
            cur.zeros[i].share = r;
            net.send_to(peer, net::msg::ring_vector, detail::dcf_impl::neg_m(r, masks[i]));
        }
        else
        {
            cur.zeros[i].share = net.recv_from<std::uint64_t>(peer, net::msg::ring_vector);
        }
    }
}

inline void barrier(trio & net, role self)
{
    const role peer = self == role::p0 ? role::p1 : role::p0;
    const std::uint8_t ready = 1;
    (void)net.exchange_with(peer, ready, net::msg::delta);
}

inline simde__m128i sample_root(int me)
{
    const auto r = dpf::uniform_sample<simde__m128i>();
    return me == 0 ? dpf::unset_lo_bit(r) : dpf::set_lo_bit(r);
}

inline dist::level_pad_msg take_level(cursor & cur, int me)
{
    const auto cw = cur.take_cw();
    const auto first = cur.take_block();
    const auto second = cur.take_block();
    dist::level_pad_msg m;
    m.cw = cw;
    if (me == 0)
    {
        m.mine = first;
        m.theirs = second;
    }
    else
    {
        m.mine = second;
        m.theirs = first;
    }
    return m;
}

struct demand
{
    std::size_t blocks = 0;
    std::size_t bits = 0;
    std::size_t b2a = 0;
    std::size_t cws = 0;
    std::vector<std::uint64_t> zero_masks;
};

inline void add_point(demand & d, std::size_t depth, std::size_t n_leaf,
    bool oblivious_hash, bool leaf_b2a)
{
    d.cws += depth;
    d.blocks += depth * 2 + n_leaf;
    if (oblivious_hash)
        d.bits += depth * dist::hash_level_and_count();
    if (leaf_b2a)
        d.b2a += n_leaf * 128;
}

inline void add_comparison(demand & d, std::size_t depth, std::size_t prefix,
    std::size_t cmp_block, std::size_t cmp_q, std::size_t cmp_h,
    bool verifiable, bool oblivious, bool edge, std::uint64_t mask)
{
    d.cws += depth;
    d.blocks += depth * 2;
    if (oblivious && edge && prefix > 1)
        d.bits += prefix - 1;
    if (oblivious)
    {
        if (cmp_block == 0)
        {
            const std::size_t nlev = prefix < depth ? prefix : depth;
            d.bits += nlev * 128;
            d.b2a += nlev * 129;
        }
        if (cmp_block > 0 && cmp_q > 0 && cmp_h > 0 && cmp_h <= depth)
        {
            d.bits += 1;
            d.b2a += 3;
        }
        if (verifiable)
            d.bits += depth * dist::hash_level_and_count();
    }
    d.zero_masks.push_back(mask);
}

template <std::size_t NLeaf>
void emit_point(tape & out, cursor & cur, int me, std::size_t depth,
    bool half, bool oblivious_hash, bool leaf_b2a)
{
    static_assert(NLeaf >= 1, "leaf mux sends at least one pad");
    if (half)
        out.put(net::msg::dpf_key, sample_root(me));
    for (std::size_t level = 0; level < depth; ++level)
    {
        out.put(net::msg::beaver_tape, take_level(cur, me));
        if (oblivious_hash)
        {
            const std::size_t n = dist::hash_level_and_count();
            std::vector<dist::bit_and_pad_msg> hash(n);
            for (std::size_t i = 0; i < n; ++i)
                hash[i] = cur.take_bit();
            out.put_vec(net::msg::beaver_tape, hash.data(), hash.size());
        }
    }
    dist::leaf_pad_msg<NLeaf> lanes{};
    for (std::size_t i = 0; i < NLeaf; ++i)
        lanes.lanes[i] = cur.take_block();
    out.put(net::msg::beaver_tape, lanes);
    if (leaf_b2a)
    {
        std::vector<dist::b2a_pad_msg> pads(NLeaf * 128);
        for (auto & p : pads)
            p = cur.take_b2a();
        out.put_vec(net::msg::beaver_tape, pads.data(), pads.size());
    }
}

inline void emit_comparison(tape & out, cursor & cur, int me, std::size_t depth,
    std::size_t prefix, std::size_t cmp_block, std::size_t cmp_q,
    std::size_t cmp_h, bool verifiable, bool oblivious, bool edge, bool half)
{
    if (half)
        out.put(net::msg::dpf_key, sample_root(me));
    if (oblivious && edge && prefix > 1)
    {
        std::vector<dist::bit_and_pad_msg> edge_pads(prefix - 1);
        for (auto & p : edge_pads)
            p = cur.take_bit();
        out.put_vec(net::msg::beaver_tape, edge_pads.data(), edge_pads.size());
    }
    for (std::size_t level = 0; level < depth; ++level)
    {
        out.put(net::msg::beaver_tape, take_level(cur, me));
        if (oblivious)
        {
            if (cmp_block == 0 && level < prefix)
            {
                dist::cmp_level_obliv_msg msg{};
                for (std::size_t i = 0; i < 128; ++i)
                {
                    const auto bit = cur.take_bit();
                    const auto b2 = cur.take_b2a();
                    msg.bits[i] = dist::word_bit_for(bit, b2);
                }
                msg.ai = cur.take_b2a();
                out.put(net::msg::beaver_tape, msg);
            }
            if (cmp_block > 0 && cmp_q > 0 && level + 1 == cmp_h)
            {
                dist::suffix_obliv_msg msg{};
                msg.and_pad = cur.take_bit();
                for (int i = 0; i < 3; ++i)
                    msg.b2a[i] = cur.take_b2a();
                out.put(net::msg::beaver_tape, msg);
            }
            if (verifiable)
            {
                const std::size_t n = dist::hash_level_and_count();
                std::vector<dist::bit_and_pad_msg> hash(n);
                for (std::size_t i = 0; i < n; ++i)
                    hash[i] = cur.take_bit();
                out.put_vec(net::msg::beaver_tape, hash.data(), hash.size());
            }
        }
    }
    out.put(net::msg::beaver_tape, cur.take_zero());
}

template <typename Concrete, typename Leaf>
void exchange_wildcard(trio & net, int me, dist::wildcard_leaf_pad_msg<Concrete, Leaf> & pad)
{
    const role peer = me == 0 ? role::p1 : role::p0;
    Concrete mine_out = dpf::uniform_sample<Concrete>();
    Leaf mine_vec = dist::random_leaf<Leaf>();
    Leaf peer_vec = net.exchange_with(peer, mine_vec, net::msg::ring_vector);
    Leaf cross_for_peer = dpf::multiply_leaf(peer_vec, mine_out);
    Leaf cross = net.exchange_with(peer, cross_for_peer, net::msg::ring_vector);
    Leaf zero{};
    if (me == 0)
    {
        zero = dist::random_leaf<Leaf>();
        const Leaf neg = dpf::subtract_leaf<Concrete>(Leaf{}, zero);
        net.send_to(peer, net::msg::ring_vector, neg);
    }
    else
    {
        zero = net.recv_from<Leaf>(peer, net::msg::ring_vector);
    }
    pad.output_blind = mine_out;
    pad.vector_blind = mine_vec;
    pad.peer_vector_blind = peer_vec;
    pad.cross = cross;
    pad.zero_leaf = zero;
}

inline void install(trio & net, role self, tape & out, cursor & cur)
{
    cur.done();
    barrier(net, self);
    net.install_inbox(std::move(out.bytes));
}

template <typename InteriorPRG, typename InputT, typename OutputT,
          bool ObliviousHash, bool Additive>
void prepare_point(trio & net, role self)
{
    const int me = me_of(self);
    using dpf_type = utils::dpf_type_t<InteriorPRG, InteriorPRG, InputT, OutputT,
        verifiable>;
    using tree = dpf::tree_traits<InteriorPRG>;
    constexpr std::size_t depth = dpf_type::depth;
    constexpr std::size_t lg = dpf_type::lg_outputs_per_leaf;
    constexpr std::size_t n_leaf = dist::leaf_and_slots(lg);
    constexpr bool leaf_b2a = lg > 0
        && !utils::has_characteristic_two_v<dpf::concrete_type_t<OutputT>>
        && (!dpf::is_wildcard_v<OutputT> || ObliviousHash);
    demand d;
    if constexpr (Additive)
        d.bits += utils::bitlength_of_v<InputT> - 1;
    add_point(d, depth, n_leaf, ObliviousHash, leaf_b2a);
    const role peer = self == role::p0 ? role::p1 : role::p0;
    auto mat = iknp::sample(net.to(peer), me, d.blocks, d.bits, d.b2a, d.cws);
    cursor cur;
    cur.mat = &mat;
    tape out;
    if constexpr (Additive)
    {
        constexpr std::size_t nbits = utils::bitlength_of_v<InputT>;
        for (std::size_t i = 0; i + 1 < nbits; ++i)
            out.put(net::msg::beaver_tape, cur.take_bit());
    }
    emit_point<n_leaf>(out, cur, me, depth, tree::is_half_tree, ObliviousHash,
        leaf_b2a);
    if constexpr (dpf::is_wildcard_v<OutputT>)
    {
        using concrete = dpf::concrete_type_t<OutputT>;
        using leaf_type = dpf::leaf_node_t<typename dpf_type::exterior_node, concrete>;
        dist::wildcard_leaf_pad_msg<concrete, leaf_type> pad{};
        exchange_wildcard<concrete, leaf_type>(net, me, pad);
        out.put(net::msg::beaver_tape, pad);
    }
    install(net, self, out, cur);
}

template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
          bool Oblivious, typename Spec, typename ...Tags>
void prepare_comparison(trio & net, role self, const Spec & spec,
    const Tags & ...tags)
{
    const int me = me_of(self);
    using pair_type = dist::comparison_pair_t<InteriorPRG, ExteriorPRG, InputT,
        Spec, Tags...>;
    using key_type = typename pair_type::first_type::key_type;
    using tree = dpf::tree_traits<InteriorPRG>;
    (void)std::initializer_list<int>{((void)tags, 0)...};
    dcf_runtime_spec runtime{};
    bool found = false;
    detail::incr::collect_cmp_one(runtime, found, spec);
    if (!found)
        throw std::invalid_argument("dist comparison keygen needs one comparison");
    if (runtime.prefix == 0)
        runtime.prefix = utils::bitlength_of_v<InputT>;
    if constexpr (Oblivious)
    {
        if (is_paint_kind(runtime.kind))
            throw std::invalid_argument(
                "oblivious comparison does not support paint kinds");
    }
    const bool edge = runtime.kind == cmp_kind::leq || runtime.kind == cmp_kind::gt;
    demand d;
    add_comparison(d, key_type::depth, runtime.prefix, key_type::cmp_block,
        key_type::cmp_q, key_type::cmp_h, key_type::is_verifiable, Oblivious,
        edge, runtime.mask);
    const role peer = self == role::p0 ? role::p1 : role::p0;
    auto mat = iknp::sample(net.to(peer), me, d.blocks, d.bits, d.b2a, d.cws);
    cursor cur;
    cur.mat = &mat;
    exchange_zeros(net, me, d.zero_masks, cur);
    tape out;
    emit_comparison(out, cur, me, key_type::depth, runtime.prefix,
        key_type::cmp_block, key_type::cmp_q, key_type::cmp_h,
        key_type::is_verifiable, Oblivious, edge, tree::is_half_tree);
    install(net, self, out, cur);
}

template <typename InteriorPRG, typename ExteriorPRG, typename InputT,
          bool Reveal, typename Beta>
void prepare_ic(trio & net, role self, const ic_pack<Beta> & spec)
{
    const int me = me_of(self);
    using input_type = std::decay_t<InputT>;
    const auto inner = detail::ic_impl::inner_lt(spec);
    const std::uint64_t gmask = detail::ic_impl::group_mask_of<Beta>();
    constexpr std::size_t nbits = utils::bitlength_of_v<input_type>;
    using pair_type = dist::comparison_pair_t<InteriorPRG, ExteriorPRG,
        input_type, decltype(inner)>;
    using key_type = typename pair_type::first_type::key_type;
    using tree = dpf::tree_traits<InteriorPRG>;
    dcf_runtime_spec runtime{};
    bool found = false;
    detail::incr::collect_cmp_one(runtime, found, inner);
    if (!found)
        throw std::invalid_argument("dist comparison keygen needs one comparison");
    if (runtime.prefix == 0)
        runtime.prefix = nbits;
    constexpr bool oblivious = !Reveal;
    if constexpr (oblivious)
    {
        if (is_paint_kind(runtime.kind))
            throw std::invalid_argument(
                "oblivious comparison does not support paint kinds");
    }
    demand d;
    if constexpr (oblivious)
    {
        d.bits += dist::ic_gamma_and_count(nbits);
        d.bits += dist::ic_correction_and_count(nbits);
        d.b2a += 4;
    }
    const bool edge = runtime.kind == cmp_kind::leq || runtime.kind == cmp_kind::gt;
    add_comparison(d, key_type::depth, runtime.prefix, key_type::cmp_block,
        key_type::cmp_q, key_type::cmp_h, key_type::is_verifiable, oblivious,
        edge, runtime.mask);
    if constexpr (Reveal)
    {
        d.zero_masks.push_back(gmask);
        d.zero_masks.push_back(gmask);
    }
    else
    {
        d.zero_masks.push_back(gmask);
    }
    const role peer = self == role::p0 ? role::p1 : role::p0;
    auto mat = iknp::sample(net.to(peer), me, d.blocks, d.bits, d.b2a, d.cws);
    cursor cur;
    cur.mat = &mat;
    exchange_zeros(net, me, d.zero_masks, cur);
    tape out;
    if constexpr (oblivious)
    {
        const std::size_t gamma_n = dist::ic_gamma_and_count(nbits);
        const std::size_t corr_n = dist::ic_correction_and_count(nbits);
        std::vector<dist::bit_and_pad_msg> gamma(gamma_n), corr(corr_n);
        for (auto & p : gamma)
            p = cur.take_bit();
        for (auto & p : corr)
            p = cur.take_bit();
        std::vector<dist::b2a_pad_msg> b2(4);
        for (auto & p : b2)
            p = cur.take_b2a();
        out.put_vec(net::msg::beaver_tape, gamma.data(), gamma.size());
        out.put_vec(net::msg::beaver_tape, corr.data(), corr.size());
        out.put_vec(net::msg::beaver_tape, b2.data(), b2.size());
    }
    emit_comparison(out, cur, me, key_type::depth, runtime.prefix,
        key_type::cmp_block, key_type::cmp_q, key_type::cmp_h,
        key_type::is_verifiable, oblivious, edge, tree::is_half_tree);
    out.put(net::msg::beaver_tape, cur.take_zero());
    if constexpr (Reveal)
        out.put(net::msg::beaver_tape, cur.take_zero());
    install(net, self, out, cur);
}

}  // namespace iknp_deal

/// @brief Two-party point keygen. Same walk as `dist_with_point_key`; pads
///        come from IKNP instead of p2.
/// @details Ideal (semi-honest): p0/p1 input XOR or additive shares of α and
///          payload β; each outputs its verifiable DPF key. With
///          `RevealPoint`, the tree prefix (and packed wildcard lane when
///          applicable) may be opened. Without it, α, peer seeds, peer
///          payload shares, and path-opening pad bits stay hidden.
/// \complexity Same tree walk as `dist_with_point_key`, plus one IKNP
///             `sample`. With tape length `T = Θ(n)` on a reveal point
///             (or `T = Θ(n · hash_level_and_count())` when oblivious),
///             `sample` does two Chou–Orlandi base sessions of `κ = 128`
///             OTs on P-256 and two OT-extension directions of cost
///             `Θ(κ T)` bits, then the usual `n` Doerner–Shelat opens.
///             Dealer `make_dpf` is still 0 rounds / 0 bytes; Half-Tree
///             §5.2 (ePrint 2022/1431) is a different `n+3`-round COT/OLE
///             hybrid this path does not implement.
/// \rounds Base OT + extension (constant rounds) then the usual per-level
///         DS exchanges; no p2.
/// \communication IKNP seed OT and extension rows (`Θ(κ T)` bits), then
///                the same peer messages as the dealer walk (no dealer
///                frames).
/// \preprocessing IKNP replaces the dealer tape; there is no offline p2.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          bool RevealPoint = false,
          typename InputT,
          typename OutputT,
          typename Fn0,
          typename Fn1>
HEDLEY_WARN_UNUSED_RESULT
auto dist_with_point_key_iknp(
    trio & net, role self, InputT x0, InputT x1, OutputT beta, Fn0 && on0,
    Fn1 && on1, bool already_encoded = false, bool additive = false)
{
    using opened_type = point_key_opened<InputT>;
    using dpf_type = utils::dpf_type_t<InteriorPRG, ExteriorPRG, InputT,
        OutputT, verifiable>;
    constexpr bool wild_lane = RevealPoint && dpf::is_wildcard_v<OutputT>
        && dpf_type::lg_outputs_per_leaf > 0;
    if (additive && already_encoded)
        throw std::invalid_argument(
            "additive DS input is converted to XOR shares; it is not pre-opened");
    if (self == role::p2)
        throw std::invalid_argument("iknp point keygen has no dealer");
    if (additive)
        iknp_deal::prepare_point<InteriorPRG, InputT, OutputT, !RevealPoint, true>(
            net, self);
    else
        iknp_deal::prepare_point<InteriorPRG, InputT, OutputT, !RevealPoint, false>(
            net, self);
    if (additive)
    {
        if (self == role::p0)
            x0 = dist::additive_to_xor_share<role::p0>(net, x0);
        else
            x1 = dist::additive_to_xor_share<role::p1>(net, x1);
    }
    if constexpr (RevealPoint)
    {
        auto finish = [](auto result) {
            opened_type out;
            out.opened_prefix = result.opened_prefix;
            if constexpr (wild_lane)
                out.opened_lane = result.opened_lane;
            return out;
        };
        if (self == role::p0)
        {
            auto result = dist::point_party<InteriorPRG, ExteriorPRG, InputT,
                OutputT, role::p0, false, true>(net, x0, beta, already_encoded);
            std::forward<Fn0>(on0)(result.dpf_key);
            return std::optional<opened_type>(finish(std::move(result)));
        }
        auto result = dist::point_party<InteriorPRG, ExteriorPRG, InputT,
            OutputT, role::p1, false, true>(net, x1, beta, already_encoded);
        std::forward<Fn1>(on1)(result.dpf_key);
        return std::optional<opened_type>(finish(std::move(result)));
    }
    else
    {
        if (self == role::p0)
            std::forward<Fn0>(on0)(dist::point_party<InteriorPRG, ExteriorPRG,
                InputT, OutputT, role::p0, false, false>(
                net, x0, beta, already_encoded));
        else
            std::forward<Fn1>(on1)(dist::point_party<InteriorPRG, ExteriorPRG,
                InputT, OutputT, role::p1, false, false>(
                net, x1, beta, already_encoded));
    }
}

/// @brief Two-party extractable point keygen.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          typename InputT,
          typename OutputT,
          typename Fn0,
          typename Fn1>
void dist_with_extractable_point_key_iknp(trio & net, role self, InputT x0,
    InputT x1, OutputT beta, Fn0 && on0, Fn1 && on1)
{
    static_assert(!dpf::is_wildcard_v<OutputT>,
        "extractable socket gen does not reconstruct a wildcard lane");
    if (self == role::p2)
        throw std::invalid_argument("iknp point keygen has no dealer");
    iknp_deal::prepare_point<InteriorPRG, InputT, OutputT, false, false>(
        net, self);
    if (self == role::p0)
    {
        on0(dist::point_party<InteriorPRG, ExteriorPRG, InputT, OutputT,
            role::p0, true>(net, x0, beta));
    }
    else
    {
        on1(dist::point_party<InteriorPRG, ExteriorPRG, InputT, OutputT,
            role::p1, true>(net, x1, beta));
    }
}

/// @brief Two-party comparison keygen.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          bool Reveal = false,
          typename InputT,
          typename Spec,
          typename Fn0,
          typename Fn1,
          typename ...Tags>
auto dist_with_cmp_key_iknp(trio & net, role self, InputT x0, InputT x1,
    const Spec & spec, Fn0 && on0, Fn1 && on1, const Tags & ...tags)
    -> std::conditional_t<Reveal, std::optional<InputT>, suppressed_point>
{
    if (self == role::p2)
        throw std::invalid_argument("iknp comparison keygen has no dealer");
    iknp_deal::prepare_comparison<InteriorPRG, ExteriorPRG, InputT, !Reveal>(
        net, self, spec, tags...);
    if constexpr (Reveal)
    {
        if (self == role::p0)
        {
            auto result = dist::comparison_party<role::p0, InteriorPRG,
                ExteriorPRG, false, true>(net, x0, spec, tags...);
            std::forward<Fn0>(on0)(result.dpf_key);
            return result.opened_point;
        }
        auto result = dist::comparison_party<role::p1, InteriorPRG,
            ExteriorPRG, false, true>(net, x1, spec, tags...);
        std::forward<Fn1>(on1)(result.dpf_key);
        return result.opened_point;
    }
    else
    {
        if (self == role::p0)
        {
            std::forward<Fn0>(on0)(dist::comparison_party<role::p0, InteriorPRG,
                ExteriorPRG, false, false>(net, x0, spec, tags...));
        }
        else
        {
            std::forward<Fn1>(on1)(dist::comparison_party<role::p1, InteriorPRG,
                ExteriorPRG, false, false>(net, x1, spec, tags...));
        }
        return {};
    }
}

/// @brief Two-party interval keygen.
template <typename InteriorPRG = dpf::prg::aes128,
          typename ExteriorPRG = InteriorPRG,
          bool Reveal = false,
          typename InputT,
          typename Beta,
          typename Fn0,
          typename Fn1>
auto dist_with_ic_key_iknp(trio & net, role self, InputT r0, InputT r1,
    const ic_pack<Beta> & spec, Fn0 && on0, Fn1 && on1)
    -> std::conditional_t<Reveal, std::optional<std::decay_t<InputT>>,
          suppressed_point>
{
    using input_type = std::decay_t<InputT>;
    using out_beta = concrete_type_t<Beta>;
    detail::ic_impl::check_input<input_type>();
    detail::ic_impl::check_bounds<input_type>(spec);
    if constexpr (detail::cmp_group_info<Beta>::custom)
    {
        throw std::invalid_argument(
            "dist interval keygen currently supports <=64-bit ring payloads");
        if constexpr (Reveal)
            return std::nullopt;
        else
            return {};
    }
    else
    {
        if (self == role::p2)
            throw std::invalid_argument("iknp interval keygen has no dealer");
        const auto inner_spec = detail::ic_impl::inner_lt(spec);
        const std::uint64_t gmask = detail::ic_impl::group_mask_of<Beta>();
        const std::uint64_t nmask = detail::ic_impl::input_mask_of<input_type>();
        iknp_deal::prepare_ic<InteriorPRG, ExteriorPRG, input_type, Reveal>(
            net, self, spec);
        const input_type mine = self == role::p0 ? r0 : r1;
        std::uint64_t delta = 0;
        std::uint64_t fval = 0;
        if constexpr (!is_wildcard_v<Beta>)
        {
            delta = detail::dcf_impl::beta_delta_u64(
                spec.if_true, spec.if_false, gmask);
            fval = detail::dcf_impl::beta_to_u64_simple(spec.if_false, gmask);
        }
        if constexpr (Reveal)
        {
            const role peer = self == role::p0 ? role::p1 : role::p0;
            const input_type theirs =
                net.exchange_with(peer, mine, net::msg::delta);
            const input_type r = utils::xor_input_shares(mine, theirs);
            const std::uint64_t r_bits = detail::ic_impl::bits_of(r);
            input_type gamma = detail::ic_impl::gamma_of(r);
            utils::flip_msb_if_signed_integral(gamma);
            const std::uint64_t cr = detail::ic_impl::correction(
                r_bits, spec.lo, spec.hi, nmask, gmask);
            const std::uint64_t absorb =
                (detail::ic_impl::mul_mask(delta, cr, gmask) + fval) & gmask;
            const auto recv_side_shares = [&]() {
                const auto z0 =
                    net.recv_from<dist::ring_zero_share_msg>(role::p2, net::msg::beaver_tape);
                const auto z1 =
                    net.recv_from<dist::ring_zero_share_msg>(role::p2, net::msg::beaver_tape);
                const std::uint64_t side = self == role::p1 ? 1ULL : 0ULL;
                const std::uint64_t dshare = is_wildcard_v<Beta>
                    ? 0ULL : (z0.share + side * delta) & gmask;
                const std::uint64_t cshare = is_wildcard_v<Beta>
                    ? 0ULL : (z1.share + side * absorb) & gmask;
                const std::uint64_t dcoeff = is_wildcard_v<Beta>
                    ? (z0.share + side) & gmask : 0ULL;
                const std::uint64_t ccoeff = is_wildcard_v<Beta>
                    ? (z1.share + side * cr) & gmask : 0ULL;
                return std::array<std::uint64_t, 4>{
                    dshare, cshare, dcoeff, ccoeff};
            };
            if (self == role::p0)
            {
                auto inner = dist::comparison_party<role::p0, InteriorPRG,
                    ExteriorPRG, true>(net, gamma, inner_spec);
                const auto shares = recv_side_shares();
                using raw_key = typename decltype(inner)::key_type;
                auto key = detail::ic_impl::make_side<0, raw_key, input_type,
                    out_beta>(std::move(inner), spec.lo, spec.hi, nmask, gmask,
                    shares[0], shares[1], shares[2], shares[3]);
                on0(std::move(key));
            }
            else
            {
                auto inner = dist::comparison_party<role::p1, InteriorPRG,
                    ExteriorPRG, true>(net, gamma, inner_spec);
                const auto shares = recv_side_shares();
                using raw_key = typename decltype(inner)::key_type;
                auto key = detail::ic_impl::make_side<1, raw_key, input_type,
                    out_beta>(std::move(inner), spec.lo, spec.hi, nmask, gmask,
                    shares[0], shares[1], shares[2], shares[3]);
                on1(std::move(key));
            }
            return r;
        }
        else
        {
            input_type gamma_share{};
            std::uint64_t cr_share = 0;
            if (self == role::p0)
            {
                auto got = dist::ic_gamma_and_correction<0, input_type>(net,
                    mine, spec.lo, spec.hi, nmask, gmask);
                gamma_share = got.first;
                cr_share = got.second;
            }
            else
            {
                auto got = dist::ic_gamma_and_correction<1, input_type>(net,
                    mine, spec.lo, spec.hi, nmask, gmask);
                gamma_share = got.first;
                cr_share = got.second;
            }
            const std::uint64_t absorb_share = is_wildcard_v<Beta>
                ? 0ULL
                : (detail::ic_impl::mul_mask(delta, cr_share, gmask)
                    + (self == role::p1 ? fval : 0ULL)) & gmask;
            const std::uint64_t cr_coeff_share =
                is_wildcard_v<Beta> ? cr_share : 0ULL;
            const auto recv_delta = [&]() {
                const auto z0 =
                    net.recv_from<dist::ring_zero_share_msg>(role::p2, net::msg::beaver_tape);
                const std::uint64_t side = self == role::p1 ? 1ULL : 0ULL;
                const std::uint64_t dshare = is_wildcard_v<Beta>
                    ? 0ULL : (z0.share + side * delta) & gmask;
                const std::uint64_t dcoeff = is_wildcard_v<Beta>
                    ? (z0.share + side) & gmask : 0ULL;
                return std::array<std::uint64_t, 2>{dshare, dcoeff};
            };
            if (self == role::p0)
            {
                auto inner = dist::comparison_party<role::p0, InteriorPRG,
                    ExteriorPRG, false, false>(net, gamma_share, inner_spec);
                const auto d = recv_delta();
                using raw_key = typename decltype(inner)::key_type;
                auto key = detail::ic_impl::make_side<0, raw_key, input_type,
                    out_beta>(std::move(inner), spec.lo, spec.hi, nmask, gmask,
                    d[0], absorb_share, d[1], cr_coeff_share);
                on0(std::move(key));
            }
            else
            {
                auto inner = dist::comparison_party<role::p1, InteriorPRG,
                    ExteriorPRG, false, false>(net, gamma_share, inner_spec);
                const auto d = recv_delta();
                using raw_key = typename decltype(inner)::key_type;
                auto key = detail::ic_impl::make_side<1, raw_key, input_type,
                    out_beta>(std::move(inner), spec.lo, spec.hi, nmask, gmask,
                    d[0], absorb_share, d[1], cr_coeff_share);
                on1(std::move(key));
            }
            return {};
        }
    }
}

}  // namespace party
}  // namespace dpf

#endif  // LIBDPF_PARTY_IKNP_DEAL_HPP__