/// @file party/run.cpp
/// @brief Shared (2+1) party case implementations.

#include "cases.hpp"
#include "dist_ds.hpp"
#include "iknp_deal.hpp"
#include "flow_util.hpp"
#include "key_io.hpp"
#include "registry.hpp"

#include <algorithm>
#include <array>
#include <chrono>
#include <cstdint>
#include <cstring>
#include <fstream>
#include <iostream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "asio.hpp"
#include "dpf.hpp"
#include "dpf/beaver.hpp"
#include "dpf/net/party_tape_io.hpp"
#include "dpf/net/trio.hpp"
#include "grotto/carry.hpp"

bool do_quickack = false;

namespace dpf
{
namespace party
{
namespace
{

using net::role;
using net::trio;
using u64 = std::uint64_t;
using util::evaluate_online;
using util::evaluate_online_auth;

struct Counter
{
    int draws = 0;
    u64 operator()()
    {
        ++draws;
        return 0x9e3779b97f4a7c15ull * static_cast<u64>(draws);
    }
};

std::pair<u64, u64> split_u64(u64 secret, unsigned tag = 1)
{
    u64 p0 = 0x9e3779b97f4a7c15ull * (tag + 1u);
    return {p0, secret - p0};
}

template <typename Ring>
std::pair<Ring, Ring> split_ring(Ring secret, unsigned tag = 1)
{
    using traits = beavers::ring_traits<Ring>;
    Ring p0 = Ring{static_cast<u64>(0x9e3779b97f4a7c15ull * (tag + 1u))};
    return {p0, traits::sub(secret, p0)};
}

void require(bool cond, const char * msg)
{
    if (!cond)
        throw std::runtime_error(msg);
}

template <typename Ring>
Ring open_additive(trio & net, role self, Ring mine)
{
    role peer = self == role::p0 ? role::p1 : role::p0;
    return net.open_with(peer, mine);
}

template <typename T>
T open_subtractive(trio & net, role self, T mine)
{
    role peer = self == role::p0 ? role::p1 : role::p0;
    T theirs = net.exchange_with(peer, mine);
    return self == role::p0 ? static_cast<T>(mine - theirs)
                            : static_cast<T>(theirs - mine);
}

template <typename T>
constexpr T share_bits(T s) noexcept
{
    return s;
}

template <typename T, std::size_t Party, sharing Scheme>
constexpr T share_bits(const secret_share<T, Party, Scheme> & s) noexcept
{
    return s.raw();
}

template <typename Ring>
void install_and_bind(beavers::session<Ring> & s, trio & net, role self,
    const std::vector<typename beavers::session<Ring>::wire> & inputs,
    const std::vector<Ring> & secrets)
{
    auto tape = net::accept_session<Ring>(net);
    s.install_party(self == role::p0 ? 0u : 1u, tape);
    for (std::size_t i = 0; i < inputs.size(); ++i)
    {
        auto [p0, p1] = split_ring<Ring>(secrets[i], static_cast<unsigned>(i + 1));
        s.bind_party(inputs[i], self == role::p0 ? p0 : p1);
    }
}

void install_and_bind_u64(beavers::session<u64> & s, trio & net, role self,
    const std::vector<beavers::session<u64>::wire> & inputs,
    const std::vector<u64> & secrets)
{
    auto tape = net::accept_session<u64>(net);
    s.install_party(self == role::p0 ? 0u : 1u, tape);
    for (std::size_t i = 0; i < inputs.size(); ++i)
    {
        auto [p0, p1] = split_u64(secrets[i], static_cast<unsigned>(i + 1));
        s.bind_party(inputs[i], self == role::p0 ? p0 : p1);
    }
}

// ---------------------------------------------------------------------------
// Beaver
// ---------------------------------------------------------------------------

int beaver_product_two(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * y);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        require(rng.draws == 5, "draws");
        require(s.preprocessing_count() == 3u, "prep");
        require(s.monomial({{x, 1u}, {y, 1u}}).open()
                == s.lambda(x).open() * s.lambda(y).open(),
            "mono");
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x, y}, {7u, 9u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 63u, "product");
    return 0;
}

int beaver_product_three_square(role self, trio & net)
{
    beavers::session<u64> s;
    auto a = s.input();
    auto b = s.input();
    auto c = s.input();
    auto p = s.product(a, b, c);
    auto sq = s(b * b);
    if (self == role::p2)
    {
        require(s.monomial_count() == 5u, "monos");
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {a, b, c}, {2u, 3u, 5u});
    evaluate_online(s, net, self);
    auto op = open_additive(net, self, s.value_party(p));
    auto os = open_additive(net, self, s.value_party(sq));
    if (self == role::p0)
    {
        require(op == 30u, "p");
        require(os == 9u, "sq");
    }
    return 0;
}

int beaver_mul_square(role self, trio & net)
{
    beavers::session<u64> s;
    auto a = s.input();
    auto x = s.input();
    auto z = s(a * x * x);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        require(s.round_of(z) == 1, "round");
        require(s.preprocessing_count() == 5u, "prep");
        require(rng.draws == 7, "draws");
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {a, x}, {4u, 5u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 100u, "mul_square");
    return 0;
}

int beaver_chained_square(role self, trio & net)
{
    beavers::session<u64> s;
    auto a = s.input();
    auto x = s.input();
    auto x2 = s(x * x);
    auto z = s(a * x2);
    if (self == role::p2)
    {
        require(s.round_of(x2) == 1 && s.round_of(z) == 2, "rounds");
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {a, x}, {4u, 5u});
    evaluate_online(s, net, self);
    auto o2 = open_additive(net, self, s.value_party(x2));
    auto oz = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
    {
        require(o2 == 25u, "x2");
        require(oz == 100u, "z");
    }
    return 0;
}

int beaver_dot(role self, trio & net)
{
    beavers::session<u64> s;
    auto x0 = s.input();
    auto x1 = s.input();
    auto x2 = s.input();
    auto y0 = s.input();
    auto y1 = s.input();
    auto y2 = s.input();
    auto z = s.dot({x0, x1, x2}, {y0, y1, y2});
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        require(s.preprocessing_count() == 7u, "prep");
        require(rng.draws == 13, "draws");
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x0, x1, x2, y0, y1, y2},
        {1u, 2u, 3u, 4u, 5u, 6u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 32u, "dot");
    return 0;
}

int beaver_dot_reuse(role self, trio & net)
{
    beavers::session<u64> s;
    auto a = s.input();
    auto b = s.input();
    auto z = s.dot({a, a}, {b, b});
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {a, b}, {3u, 4u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 24u, "dot_reuse");
    return 0;
}

int beaver_inner_product(role self, trio & net)
{
    using wire = beavers::session<u64>::wire;
    beavers::session<u64> s;
    std::vector<wire> fx, fy;
    for (int i = 0; i < 3; ++i)
    {
        fx.push_back(s.input());
        fy.push_back(s.input());
    }
    auto z = s.dot(fx, fy);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    std::vector<wire> wires = fx;
    wires.insert(wires.end(), fy.begin(), fy.end());
    install_and_bind_u64(s, net, self, wires, {0u, 2u, 5u, 4u, 0u, 7u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 35u, "inner");
    return 0;
}

int beaver_scale(role self, trio & net)
{
    beavers::session<u64> s;
    auto sc = s.input();
    auto l0 = s.input();
    auto l1 = s.input();
    auto l2 = s.input();
    auto out = s.scale(sc, {l0, l1, l2});
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {sc, l0, l1, l2}, {3u, 1u, 2u, 4u});
    evaluate_online(s, net, self);
    const u64 expect[] = {3u, 6u, 12u};
    for (std::size_t i = 0; i < out.size(); ++i)
    {
        auto o = open_additive(net, self, s.value_party(out[i]));
        if (self == role::p0)
            require(o == expect[i], "scale");
    }
    return 0;
}

int beaver_bit_mul_mux(role self, trio & net)
{
    beavers::session<u64> s;
    auto b = s.bit();
    auto x = s.input();
    auto y = s.input();
    auto bm = s.bit_mul(b, x);
    auto mx = s.mux(b, x, y);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {b, x, y}, {1u, 10u, 3u});
    evaluate_online(s, net, self);
    auto ob = open_additive(net, self, s.value_party(bm));
    auto om = open_additive(net, self, s.value_party(mx));
    if (self == role::p0)
    {
        require(ob == 10u, "bit_mul");
        require(om == 10u, "mux");
    }
    return 0;
}

int beaver_xor_mux(role self, trio & net)
{
    using ring = xor_wrapper<std::uint64_t>;
    beavers::session<ring> s;
    auto b = s.bit();
    auto x = s.input();
    auto y = s.input();
    auto mx = s.mux(b, x, y);
    if (self == role::p2)
    {
        struct XorCounter
        {
            Counter c;
            ring operator()() { return ring{c()}; }
        } rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind(s, net, self, {b, x, y},
        {ring{0}, ring{0xaaaau}, ring{0x1111u}});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(mx));
    if (self == role::p0)
        require(o == ring{0x1111u}, "xor_mux");
    return 0;
}

int beaver_modint_product(role self, trio & net)
{
    using ring = modint<61>;
    beavers::session<ring> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * y);
    if (self == role::p2)
    {
        s.sample();
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind(s, net, self, {x, y}, {ring{7}, ring{9}});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == ring{63}, "modint");
    return 0;
}

int beaver_oracle_stream(role self, trio & net)
{
    using block = prg::aes128::block_type;
    block seed = simde_mm_set_epi64x(0x9, 0x9);
    constexpr std::size_t n = 32;
    if (self == role::p2)
    {
        beavers::oracle<u64> src(seed, 4);
        beavers::beaver2<u64> bulk[n];
        beavers::fill_beaver2(src, 0, bulk, n);
        auto at31 = beavers::beaver2_at(src, 31);
        require(bulk[31].a == at31.a && bulk[31].ab == at31.ab, "stream");
        require(bulk[0].ab.open() == bulk[0].a.open() * bulk[0].b.open(), "o0");
        std::vector<u64> p0, p1;
        for (std::size_t i = 0; i < n; ++i)
        {
            p0.insert(p0.end(),
                {bulk[i].a.p0, bulk[i].b.p0, bulk[i].ab.p0, bulk[i].out.p0});
            p1.insert(p1.end(),
                {bulk[i].a.p1, bulk[i].b.p1, bulk[i].ab.p1, bulk[i].out.p1});
        }
        net.to(role::p0).send_vec(p0);
        net.to(role::p1).send_vec(p1);
        return 0;
    }
    auto shares = net.to(role::p2).recv_vec<u64>();
    require(shares.size() == n * 4, "size");
    const std::size_t i = 7;
    u64 a = shares[i * 4 + 0];
    u64 b = shares[i * 4 + 1];
    u64 ab = shares[i * 4 + 2];
    auto [x0, x1] = split_u64(6u, 1);
    auto [y0, y1] = split_u64(7u, 2);
    u64 xs = self == role::p0 ? x0 : x1;
    u64 ys = self == role::p0 ? y0 : y1;
    role peer = self == role::p0 ? role::p1 : role::p0;
    u64 dx_mine = xs + a;
    u64 dy_mine = ys + b;
    u64 dx = dx_mine + net.exchange_with(peer, dx_mine);
    u64 dy = dy_mine + net.exchange_with(peer, dy_mine);
    u64 z = ab;
    z -= dx * b;
    z -= dy * a;
    if (self == role::p0)
        z += dx * dy;
    auto open = open_additive(net, self, z);
    if (self == role::p0)
        require(open == 42u, "oracle online");
    return 0;
}

int beaver_scheduled_replay(role self, trio & net)
{
    using block = prg::aes128::block_type;
    block seed = simde_mm_set_epi64x(0x55, 0xaa);
    beavers::session<u64> left;
    auto x = left.input();
    auto y = left.input();
    auto z = left(3 * x * x + 2 * y + 5);
    if (self == role::p2)
    {
        beavers::oracle<u64> src(seed, 4);
        left.sample_from(src, 3);
        beavers::session<u64> right;
        auto rx = right.input();
        auto ry = right.input();
        (void)right(3 * rx * rx + 2 * ry + 5);
        right.sample_from(src, 3);
        require(left.lambda(x) == right.lambda(rx), "replay");
        net::deal_session(net, left);
        return 0;
    }
    install_and_bind_u64(left, net, self, {x, y}, {4u, 6u});
    evaluate_online(left, net, self);
    auto o = open_additive(net, self, left.value_party(z));
    if (self == role::p0)
        require(o == 3u * 16u + 2u * 6u + 5u, "scheduled");
    return 0;
}

int beaver_grotto_appendix_e(role self, trio & net)
{
    beavers::session<u64> linear;
    auto x = linear.input();
    auto sgn = linear.input();
    auto a0 = linear.input();
    auto a1 = linear.input();
    auto lin = linear(sgn * (a1 * x + a0));
    if (self == role::p2)
    {
        require(linear.round_of(lin) == 2, "round");
        require(linear.preprocessing_count() == 6u, "prep");
        Counter rng;
        linear.sample(rng);
        net::deal_session(net, linear);
        return 0;
    }
    install_and_bind_u64(linear, net, self, {x, sgn, a0, a1}, {2u, 3u, 4u, 5u});
    evaluate_online(linear, net, self);
    auto o = open_additive(net, self, linear.value_party(lin));
    if (self == role::p0)
        require(o == 3u * (5u * 2u + 4u), "grotto_e");
    return 0;
}

int beaver_polynomial_share_powers(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * x * y + x * y);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x, y}, {3u, 4u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 3u * 3u * 4u + 3u * 4u, "poly");
    return 0;
}

int beaver_rejects_bad_use(role self, trio & net)
{
    if (self == role::p2)
    {
        bool threw = false;
        try
        {
            beavers::session<u64> s;
            auto x = s.input();
            auto y = s.input();
            (void)s.dot({x}, {y, y});
        }
        catch (const std::invalid_argument &)
        {
            threw = true;
        }
        require(threw, "reject");
        net.to(role::p0).send(net::msg::case_ok, std::uint8_t{1});
        net.to(role::p1).send(net::msg::case_ok, std::uint8_t{1});
        return 0;
    }
    (void)net.to(role::p2).recv<std::uint8_t>(net::msg::case_ok);
    return 0;
}

int beaver_one_shot_triples(role self, trio & net)
{
    if (self == role::p2)
    {
        Counter rng;
        auto t2 = beavers::sample_beaver2<u64>(rng);
        require(t2.ab.open() == t2.a.open() * t2.b.open(), "b2");
        auto t3 = beavers::sample_beaver3<u64>(rng);
        require(t3.abc.open() == t3.a.open() * t3.b.open() * t3.c.open(), "b3");
        net.to(role::p0).send(net::msg::case_ok, std::uint8_t{1});
        net.to(role::p1).send(net::msg::case_ok, std::uint8_t{1});
        beavers::session<u64> s;
        auto x = s.input();
        auto y = s.input();
        (void)s(x * y);
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    (void)net.to(role::p2).recv<std::uint8_t>(net::msg::case_ok);
    beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x * y);
    install_and_bind_u64(s, net, self, {x, y}, {8u, 9u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 72u, "one_shot");
    return 0;
}

int beaver_batch_one_round(role self, trio & net)
{
    beavers::session<u64> s;
    auto a = s.input();
    auto b = s.input();
    auto c = s.input();
    auto p = s(a * b);
    auto q = s(a * c);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {a, b, c}, {2u, 3u, 5u});
    evaluate_online(s, net, self);
    auto op = open_additive(net, self, s.value_party(p));
    auto oq = open_additive(net, self, s.value_party(q));
    if (self == role::p0)
    {
        require(op == 6u && oq == 10u, "batch");
    }
    return 0;
}

int beaver_blind_survives(role self, trio & net)
{
    beavers::session<u64> s;
    auto a = s.input();
    auto b = s.input();
    auto p = s(a * b);
    auto q = s(p * a);
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        auto lam = s.lambda(a);
        net::deal_session(net, s);
        require(s.lambda(a) == lam, "survived");
        return 0;
    }
    install_and_bind_u64(s, net, self, {a, b}, {3u, 4u});
    evaluate_online(s, net, self);
    auto op = open_additive(net, self, s.value_party(p));
    auto oq = open_additive(net, self, s.value_party(q));
    if (self == role::p0)
        require(op == 12u && oq == 36u, "blind");
    return 0;
}

int beaver_sums_cancellations(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto y = s.input();
    auto z = s(x + y - x);
    if (self == role::p2)
    {
        require(s.monomial_count() == 0u, "no products");
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x, y}, {11u, 22u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 22u, "sum");
    return 0;
}

// ---------------------------------------------------------------------------
// Wildcard: parties assign over the p0–p1 link, then evaluate.
// ---------------------------------------------------------------------------

template <typename Key, typename Share>
void assign_wildcard_over_link(Key & key, Share my_share, trio & net, role self)
{
    role peer = self == role::p0 ? role::p1 : role::p0;
    auto & wrap = std::get<0>(key.leaf_nodes);
    if (wrap.is_ready())
        wrap.begin_update();
    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);
}

int wildcard_single_leaf(role self, trio & net)
{
    using input_type = std::uint8_t;
    using concrete_type = std::uint32_t;
    using output_type = wildcard_value<concrete_type>;
    input_type x = 0xAA;
    input_type x0 = 0x31;
    input_type x1 = static_cast<input_type>(x ^ x0);
    concrete_type y_exp = 0xAAAAAAAA;
    concrete_type y_shr0 = 0x12345678;
    concrete_type y_shr1 = y_exp - y_shr0;

    auto on = [&](auto key) {
        const concrete_type my_share =
            self == role::p0 ? y_shr0 : y_shr1;
        assign_wildcard_over_link(key, my_share, net, self);

        auto y = *eval_point(key, x);
        auto open = open_subtractive(net, self, share_bits(y));
        if (self == role::p0)
            require(open == y_exp, "wildcard on");
        auto yz = *eval_point(key, static_cast<input_type>(x ^ 1));
        auto oz = open_subtractive(net, self, share_bits(yz));
        if (self == role::p0)
            require(oz == 0u, "wildcard off");
    };
    const auto encoded = encoded_xor_point(x0, x1);
    require_tree_prefix(
        dist_with_point_key<prg::aes128, prg::aes128, true>(
            net, self, x0, x1, output_type{}, on, on),
        self,
        verifiable_tree_prefix<prg::aes128, prg::aes128, output_type>(x0, x1),
        verifiable_lane<prg::aes128, prg::aes128, output_type>(encoded));
    dist_with_point_key(net, self, x0, x1, output_type{}, on, on);
    return 0;
}

int wildcard_fail_before_assign(role self, trio & net)
{
    using input_type = std::uint8_t;
    using output_type = wildcard_value<std::uint32_t>;
    const input_type alpha = 0xAA;
    const input_type x0 = 0x31;
    const input_type x1 = static_cast<input_type>(alpha ^ x0);
    auto on = [&](const auto & key) {
        bool threw = false;
        try
        {
            (void)eval_point(key, alpha);
        }
        catch (const std::runtime_error &)
        {
            threw = true;
        }
        require(threw, "fail before");
    };
    const auto encoded = encoded_xor_point(x0, x1);
    require_tree_prefix(
        dist_with_point_key<prg::aes128, prg::aes128, true>(
            net, self, x0, x1, output_type{}, on, on),
        self,
        verifiable_tree_prefix<prg::aes128, prg::aes128, output_type>(x0, x1),
        verifiable_lane<prg::aes128, prg::aes128, output_type>(encoded));
    return 0;
}

// ---------------------------------------------------------------------------
// Geneval / DS — p0 and p1 run Doerner–Shelat. p2 deals pads only.
// ---------------------------------------------------------------------------

int geneval_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const std::uint64_t beta = 7;
    require_tree_prefix(dist_with_point_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, beta, [&](const auto & key) {
        const auto y = *eval_point(key, alpha);
        const auto open = open_subtractive(net, self, share_bits(y));
        if (self == role::p0)
            require(open == beta, "geneval");
    }, [&](const auto & key) {
        const auto y = *eval_point(key, alpha);
        (void)open_subtractive(net, self, share_bits(y));
    }), self, verifiable_tree_prefix<prg::aes128, prg::aes128, std::uint64_t>(x0, x1));
    return 0;
}

int geneval_arith_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha - x0);
    const std::uint64_t beta = 11;
    require_tree_prefix(dist_with_point_key<prg::aes128, prg::aes128, true>(
        net, self, x0, x1, beta, [&](const auto & key) {
        const auto y = *eval_point(key, alpha);
        const auto open = open_subtractive(net, self, share_bits(y));
        if (self == role::p0)
            require(open == beta, "arith");
    }, [&](const auto & key) {
        const auto y = *eval_point(key, alpha);
        (void)open_subtractive(net, self, share_bits(y));
    }, false, true), self,
        verifiable_tree_prefix<prg::aes128, prg::aes128, std::uint64_t>(
            alpha, Input{}, true));
    return 0;
}

int ds_key_agrees(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x44;
    const Input x0 = 0x12;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const std::uint64_t beta = 5;
    require_tree_prefix(dist_with_point_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, beta, [&](const auto & key) {
        const auto y = *eval_point(key, alpha);
        const auto open = open_subtractive(net, self, share_bits(y));
        if (self == role::p0)
            require(open == beta, "ds");
    }, [&](const auto & key) {
        const auto y = *eval_point(key, alpha);
        (void)open_subtractive(net, self, share_bits(y));
    }), self, verifiable_tree_prefix<prg::aes128, prg::aes128, std::uint64_t>(x0, x1));
    return 0;
}

// ---------------------------------------------------------------------------
// Verifiable
// ---------------------------------------------------------------------------

int verifiable_honest_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const std::uint64_t beta = 7;
    auto on = [&](const auto & key) {
        proof_token pi{};
        auto y = *eval_point(key, alpha, prove(pi));
        role peer = self == role::p0 ? role::p1 : role::p0;
        auto send_pi = [&](const proof_token & t) {
            net.to(peer).send_bytes(net::msg::proof_token,
                reinterpret_cast<const std::uint8_t *>(&t), sizeof(t));
        };
        auto recv_pi = [&]() {
            proof_token t{};
            auto b = net.to(peer).recv_bytes(net::msg::proof_token);
            require(b.size() == sizeof(t), "pi size");
            std::memcpy(&t, b.data(), sizeof(t));
            return t;
        };
        proof_token peer_pi{};
        if (self == role::p0)
        {
            send_pi(pi);
            peer_pi = recv_pi();
        }
        else
        {
            peer_pi = recv_pi();
            send_pi(pi);
        }
        bool ok = verify(pi, peer_pi);
        auto open = open_subtractive(net, self, share_bits(y));
        if (self == role::p0)
        {
            require(ok, "verify");
            require(open == beta, "vpoint");
        }
    };
    require_tree_prefix(dist_with_point_key<prg::aes128, prg::aes128, true>(
        net, self, x0, x1, beta, on, on), self,
        verifiable_tree_prefix<prg::aes128, prg::aes128, std::uint64_t>(x0, x1));
    return 0;
}

int verifiable_tampered_rejects(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const std::uint64_t beta = 7;
    auto on = [&](const auto & key) {
        proof_token pi{};
        (void)*eval_point(key, alpha, prove(pi));
        if (self == role::p0)
            pi[0] = simde_mm_xor_si128(pi[0], simde_mm_set1_epi8(0xff));
        role peer = self == role::p0 ? role::p1 : role::p0;
        auto send_pi = [&](const proof_token & t) {
            net.to(peer).send_bytes(net::msg::proof_token,
                reinterpret_cast<const std::uint8_t *>(&t), sizeof(t));
        };
        auto recv_pi = [&]() {
            proof_token t{};
            auto b = net.to(peer).recv_bytes(net::msg::proof_token);
            std::memcpy(&t, b.data(), sizeof(t));
            return t;
        };
        proof_token peer_pi{};
        if (self == role::p0)
        {
            send_pi(pi);
            peer_pi = recv_pi();
        }
        else
        {
            peer_pi = recv_pi();
            send_pi(pi);
        }
        bool ok = verify(pi, peer_pi);
        if (self == role::p0)
            require(!ok, "tampered token");
    };
    require_tree_prefix(dist_with_point_key<prg::aes128, prg::aes128, true>(
        net, self, x0, x1, beta, on, on), self,
        verifiable_tree_prefix<prg::aes128, prg::aes128, std::uint64_t>(x0, x1));
    return 0;
}

// ---------------------------------------------------------------------------
// Grotto / DCF
// ---------------------------------------------------------------------------

int blocked_dcf_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x40;
    const Input x0 = 0x14;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const std::uint64_t beta = 9;
    auto on = [&](const auto & key) {
        auto y = eval_point(cmp, key, static_cast<Input>(alpha - 1));
        auto open =
            open_additive(net, self, share_bits(y)) & key.cmp().mask;
        if (self == role::p0)
            require(open == beta, "dcf lt");
        auto y2 = eval_point(cmp, key, alpha);
        auto o2 =
            open_additive(net, self, share_bits(y2)) & key.cmp().mask;
        if (self == role::p0)
            require(o2 == 0u, "dcf on");
    };
    require_opened(dist_with_cmp_key<prg::aes128, prg::aes128, true>(net, self, x0, x1, lt(beta), on, on),
        self, encoded_xor_point(x0, x1));
    return 0;
}

int grotto_prefix_horner(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto z = s.horner(x, {u64{1}, u64{2}, u64{3}});
    if (self == role::p2)
    {
        Counter rng;
        s.sample(rng);
        net::deal_session(net, s);
        return 0;
    }
    install_and_bind_u64(s, net, self, {x}, {5u});
    evaluate_online(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    if (self == role::p0)
        require(o == 1u + 2u * 5u + 3u * 25u, "horner");
    return 0;
}

/// @brief Opt-in MAC path: dealer sets Δ, parties open with auth openings.
int beaver_auth_horner(role self, trio & net)
{
    beavers::session<u64> s;
    auto x = s.input();
    auto z = s.horner(x, {u64{1}, u64{2}, u64{3}});
    const u64 secret = 5u;
    if (self == role::p2)
    {
        auto key = sample_mac_key<u64>();
        s.set_mac_key(key);
        Counter rng;
        s.sample(rng);
        s.bind(x, secret, rng);
        auto xv = s.value_auth(x);
        net::deal_session(net, s);
        net.to(role::p0).send(net::msg::mac_key, key.delta);
        net.to(role::p1).send(net::msg::mac_key, key.delta);
        net.to(role::p0).send(net::msg::mac_share, xv.party(0));
        net.to(role::p1).send(net::msg::mac_share, xv.party(1));
        return 0;
    }
    auto tape = net::accept_session<u64>(net);
    s.install_party(self == role::p0 ? 0u : 1u, tape);
    auto delta = net.to(role::p2).recv<u64>(net::msg::mac_key);
    s.set_mac_key(mac_key<u64>{delta});
    auto xin = net.to(role::p2).recv<mac_share<u64>>(net::msg::mac_share);
    s.bind_party(x, xin);
    evaluate_online_auth(s, net, self);
    auto o = open_additive(net, self, s.value_party(z));
    require(s.verify_delta(x, mac_key<u64>{delta}), "auth delta x");
    if (self == role::p0)
        require(o == 1u + 2u * 5u + 3u * 25u, "auth horner");
    return 0;
}

/// @brief Carry with verifiable DCF tokens and an output MAC.
int carry_verifiable_mac(role self, trio & net)
{
    using Key0 = grotto::carry_detail::lt_v_pair::first_type;
    using Key1 = grotto::carry_detail::lt_v_pair::second_type;
    constexpr unsigned n = 8;
    constexpr unsigned s = 3;
    const std::uint64_t opened = 0x3cu;
    if (self == role::p2)
    {
        grotto::carry_auth auth{};
        auth.verifiable = true;
        auth.output_mac = true;
        auto keys = grotto::make_carry_in_keys(n, s, auth);
        require(keys.low_lt_v.has_value(), "carry lt v");
        require(keys.has_mac, "carry mac");
        auto send_one = [&](role r) {
            send_key(net.to(r), keys.low_lt_v->first);
            send_key(net.to(r), keys.low_lt_v->second);
            net.to(r).send(net::msg::mac_key, keys.mac.delta);
            net.to(r).send(net::msg::bytes, keys.rin);
            net.to(r).send(net::msg::bytes, keys.rout0);
            net.to(r).send(net::msg::bytes, keys.rout1);
        };
        send_one(role::p0);
        send_one(role::p1);
        return 0;
    }
    auto k0 = recv_key<Key0>(net.to(role::p2));
    auto k1 = recv_key<Key1>(net.to(role::p2));
    auto delta = net.to(role::p2).recv<std::uint64_t>(net::msg::mac_key);
    auto rin = net.to(role::p2).recv<std::uint64_t>(net::msg::bytes);
    auto rout0 = net.to(role::p2).recv<std::uint64_t>(net::msg::bytes);
    auto rout1 = net.to(role::p2).recv<std::uint64_t>(net::msg::bytes);

    grotto::carry_detail::carry_key_pair keys{};
    keys.recipe = grotto::plan_carry_in(n, s);
    keys.auth.verifiable = true;
    keys.auth.output_mac = true;
    keys.rin = rin;
    keys.rout0 = rout0;
    keys.rout1 = rout1;
    keys.low_lt_v.emplace(std::move(k0), std::move(k1));
    keys.mac = mac_key<std::uint64_t>{delta};
    keys.has_mac = true;

    const unsigned party = self == role::p0 ? 0u : 1u;
    dpf::proof_token toks[4]{};
    const auto nt = grotto::prove_carry_keys(keys, party, opened, toks, 4);
    require(nt >= 1u, "carry tokens");

    role peer = self == role::p0 ? role::p1 : role::p0;
    auto send_pi = [&](const dpf::proof_token & t) {
        net.to(peer).send_bytes(net::msg::proof_token,
            reinterpret_cast<const std::uint8_t *>(&t), sizeof(t));
    };
    auto recv_pi = [&]() {
        dpf::proof_token t{};
        auto b = net.to(peer).recv_bytes(net::msg::proof_token);
        require(b.size() == sizeof(t), "pi size");
        std::memcpy(&t, b.data(), sizeof(t));
        return t;
    };
    std::vector<dpf::proof_token> mine(toks, toks + nt);
    std::vector<dpf::proof_token> theirs(nt);
    if (self == role::p0)
    {
        for (std::size_t i = 0; i < nt; ++i)
        {
            send_pi(mine[i]);
            theirs[i] = recv_pi();
        }
    }
    else
    {
        for (std::size_t i = 0; i < nt; ++i)
        {
            theirs[i] = recv_pi();
            send_pi(mine[i]);
        }
    }
    require(dpf::verify_batch(mine, theirs), "carry verify_batch");

    auto y = grotto::eval_carry_in(keys, party, opened);
    auto y_peer = open_additive(net, self, y.value);
    // Reconstruct shares for MAC: party 0 holds y.value, peer opened sum is
    // y0+y1; split as (y.value, sum - y.value) on p0 for the check.
    std::uint64_t y0 = 0;
    std::uint64_t y1 = 0;
    if (self == role::p0)
    {
        y0 = y.value;
        y1 = y_peer - y.value;
        auto [m0, m1] = grotto::mac_carry_result(keys, y0, y1);
        require(dpf::mac_verify(m0, m1, keys.mac, mine[0], theirs[0]),
            "carry mac");
        require(!dpf::mac_verify(m0, m1, keys.mac, detail::vdpf::zero_proof(),
            theirs[0]), "carry mac zero token");
        (void)m0;
        (void)m1;
    }
    return 0;
}

/// @brief Extractable sketch: fold opened payloads, exchange, verify.
int extractable_sketch_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const dpf::fp61 beta{7};
    auto on = [&](const auto & key) {
        const std::array<dpf::fp61, 1> rs{dpf::fp61{3}};
        sketch_share local{};
        auto sk = sketch(local, rs);
        (void)*eval_point(key, alpha, sk);
        role peer = self == role::p0 ? role::p1 : role::p0;
        sketch_share theirs =
            net.exchange_with(peer, local, net::msg::sketch_share);
        bool ok = self == role::p0 ? sketch_verify(local, theirs)
                                   : sketch_verify(theirs, local);
        if (self == role::p0)
            require(ok, "extractable sketch");
    };
    dist_with_extractable_point_key(
        net, self, x0, x1, beta, on, on);
    return 0;
}

/// @brief Two-party geneval point. Same check as `geneval_point`, pads from IKNP.
int iknp_geneval_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const u64 beta = 7;
    require_tree_prefix(dist_with_point_key_iknp<prg::aes128, prg::aes128, true>(
        net, self, x0, x1, beta, [&](const auto & key) {
            const auto open = open_subtractive(net, self, share_bits(*eval_point(key, alpha)));
            if (self == role::p0)
                require(open == beta, "iknp geneval");
        }, [&](const auto & key) {
            (void)open_subtractive(net, self, share_bits(*eval_point(key, alpha)));
        }), self, verifiable_tree_prefix<prg::aes128, prg::aes128, u64>(x0, x1));
    return 0;
}

/// @brief Two-party additive-share input. Same check as `geneval_arith_point`.
int iknp_arith_point(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha - x0);
    const u64 beta = 11;
    require_tree_prefix(dist_with_point_key_iknp<prg::aes128, prg::aes128, true>(
        net, self, x0, x1, beta, [&](const auto & key) {
            const auto open = open_subtractive(net, self, share_bits(*eval_point(key, alpha)));
            if (self == role::p0)
                require(open == beta, "iknp arith");
        }, [&](const auto & key) {
            (void)open_subtractive(net, self, share_bits(*eval_point(key, alpha)));
        }, false, true), self,
        verifiable_tree_prefix<prg::aes128, prg::aes128, u64>(alpha, Input{}, true));
    return 0;
}

/// @brief Two-party wildcard: keygen hides the lane, then β is assigned on the link.
int iknp_wildcard_leaf(role self, trio & net)
{
    using input_type = std::uint8_t;
    using concrete_type = std::uint32_t;
    using output_type = wildcard_value<concrete_type>;
    const input_type x = 0xAA;
    const input_type x0 = 0x31;
    const input_type x1 = static_cast<input_type>(x ^ x0);
    const concrete_type y_exp = 0xAAAAAAAA;
    const concrete_type y_shr0 = 0x12345678;
    const concrete_type y_shr1 = y_exp - y_shr0;
    auto on = [&](auto key) {
        const concrete_type my_share = self == role::p0 ? y_shr0 : y_shr1;
        assign_wildcard_over_link(key, my_share, net, self);
        auto open = open_subtractive(net, self, share_bits(*eval_point(key, x)));
        if (self == role::p0)
            require(open == y_exp, "iknp wildcard on");
        auto oz = open_subtractive(net, self,
            share_bits(*eval_point(key, static_cast<input_type>(x ^ 1))));
        if (self == role::p0)
            require(oz == 0u, "iknp wildcard off");
    };
    dist_with_point_key_iknp(net, self, x0, x1, output_type{}, on, on);
    return 0;
}

/// @brief Two-party extractable sketch. Same check as `extractable_sketch_point`.
int iknp_extractable_sketch(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x2a;
    const Input x0 = 0x10;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const dpf::fp61 beta{7};
    auto on = [&](const auto & key) {
        const std::array<dpf::fp61, 1> rs{dpf::fp61{3}};
        sketch_share local{};
        auto sk = sketch(local, rs);
        (void)*eval_point(key, alpha, sk);
        role peer = self == role::p0 ? role::p1 : role::p0;
        sketch_share theirs = net.exchange_with(peer, local, net::msg::sketch_share);
        bool ok = self == role::p0 ? sketch_verify(local, theirs)
                                   : sketch_verify(theirs, local);
        if (self == role::p0)
            require(ok, "iknp extractable sketch");
    };
    dist_with_extractable_point_key_iknp(net, self, x0, x1, beta, on, on);
    return 0;
}

/// @brief Two-party comparison, point stays shared.
int iknp_cmp_oblivious(role self, trio & net)
{
    using Input = std::uint8_t;
    const Input alpha = 0x40;
    const Input x0 = 0x14;
    const Input x1 = static_cast<Input>(alpha ^ x0);
    const u64 beta = 9;
    auto on = [&](const auto & key) {
        const u64 mask = key.cmp().mask;
        auto below = open_additive(net, self, share_bits(eval_point(cmp, key,
            static_cast<Input>(alpha - 1)))) & mask;
        auto at = open_additive(net, self, share_bits(eval_point(cmp, key, alpha))) & mask;
        if (self == role::p0)
        {
            require(below == beta, "iknp cmp below");
            require(at == 0u, "iknp cmp at");
        }
    };
    dist_with_cmp_key_iknp(net, self, x0, x1, lt(beta), on, on);
    return 0;
}

void register_core_flows_impl()
{
    register_flow({"beaver_product_two", "beaver product bench", beaver_product_two, true});
    register_flow({"beaver_product_three_square", "beaver product bench", beaver_product_three_square, true});
    register_flow({"beaver_mul_square", "beaver product bench", beaver_mul_square, true});
    register_flow({"beaver_chained_square", "beaver product bench", beaver_chained_square, true});
    register_flow({"beaver_dot", "beaver dot bench", beaver_dot, true});
    register_flow({"beaver_dot_reuse", "beaver dot bench", beaver_dot_reuse, true});
    register_flow({"beaver_inner_product", "beaver dot bench", beaver_inner_product, true});
    register_flow({"beaver_scale", "beaver scale bench", beaver_scale, true});
    register_flow({"beaver_bit_mul_mux", "beaver mux bench", beaver_bit_mul_mux, true});
    register_flow({"beaver_xor_mux", "beaver mux bench", beaver_xor_mux, true});
    register_flow({"beaver_modint_product", "beaver product bench", beaver_modint_product, true});
    register_flow({"beaver_oracle_stream", "beaver stream bench", beaver_oracle_stream, true});
    register_flow({"beaver_scheduled_replay", "beaver stream bench", beaver_scheduled_replay, true});
    register_flow({"beaver_grotto_appendix_e", "beaver grotto poly bench", beaver_grotto_appendix_e, true});
    register_flow({"beaver_polynomial_share_powers", "beaver poly bench", beaver_polynomial_share_powers, true});
    register_flow({"beaver_rejects_bad_use", "beaver smoke", beaver_rejects_bad_use, false});
    register_flow({"beaver_one_shot_triples", "beaver product bench", beaver_one_shot_triples, true});
    register_flow({"beaver_batch_one_round", "beaver product bench", beaver_batch_one_round, true});
    register_flow({"beaver_blind_survives", "beaver product bench", beaver_blind_survives, true});
    register_flow({"beaver_sums_cancellations", "beaver poly bench", beaver_sums_cancellations, true});
    register_flow({"wildcard_single_leaf", "wildcard dpf bench", wildcard_single_leaf, true});
    register_flow({"wildcard_fail_before_assign", "wildcard smoke", wildcard_fail_before_assign, false});
    register_flow({"geneval_point", "geneval dpf bench", geneval_point, true});
    register_flow({"geneval_arith_point", "geneval dpf bench", geneval_arith_point, true});
    register_flow({"ds_key_agrees", "ds geneval dpf bench", ds_key_agrees, true});
    register_flow({"verifiable_honest_point", "verifiable dpf bench", verifiable_honest_point, true});
    register_flow({"verifiable_tampered_rejects", "verifiable smoke", verifiable_tampered_rejects, false});
    register_flow({"blocked_dcf_point", "dcf grotto bench", blocked_dcf_point, true});
    register_flow({"grotto_prefix_horner", "beaver horner grotto bench", grotto_prefix_horner, true});
    register_flow({"beaver_auth_horner", "beaver auth mac", beaver_auth_horner, false});
    register_flow({"carry_verifiable_mac", "carry vdpf mac", carry_verifiable_mac, false});
    register_flow({"extractable_sketch_point", "extractable sketch", extractable_sketch_point, false});
    // Tag is `iknp`, not `bench`: party_bench spawns these with p0/p1 only.
    register_flow({"iknp_geneval_point", "iknp geneval dpf", iknp_geneval_point, true});
    register_flow({"iknp_arith_point", "iknp geneval dpf", iknp_arith_point, true});
    register_flow({"iknp_wildcard_leaf", "iknp wildcard dpf", iknp_wildcard_leaf, true});
    register_flow({"iknp_extractable_sketch", "iknp extractable", iknp_extractable_sketch, true});
    register_flow({"iknp_cmp_oblivious", "iknp dcf", iknp_cmp_oblivious, true});
}

}  // namespace

void register_core_flows()
{
    register_core_flows_impl();
}

namespace
{

struct run_opts
{
    std::string dir;
    std::string case_name;
    std::string tag;
    std::uint64_t repeat = 1;
    std::uint64_t warmup = 0;
    bool metrics = false;
    bool list_only = false;
    bool pair = false;
};

run_opts parse_opts(int argc, char ** argv)
{
    run_opts o;
    for (int i = 1; i < argc; ++i)
    {
        std::string a = argv[i];
        auto need = [&](const char * flag) -> std::string {
            if (i + 1 >= argc)
                throw std::invalid_argument(std::string("missing value for ") + flag);
            return argv[++i];
        };
        if (a == "--dir")
            o.dir = need("--dir");
        else if (a == "--case")
            o.case_name = need("--case");
        else if (a == "--tag")
            o.tag = need("--tag");
        else if (a == "--repeat")
            o.repeat = std::stoull(need("--repeat"));
        else if (a == "--warmup")
            o.warmup = std::stoull(need("--warmup"));
        else if (a == "--metrics")
            o.metrics = true;
        else if (a == "--list")
            o.list_only = true;
        else if (a == "--pair")
            o.pair = true;
    }
    return o;
}

int run_one_flow(role self, const flow & f, trio & net, const run_opts & o)
{
    using clock = std::chrono::steady_clock;
    std::uint64_t total_ns = 0;
    std::uint64_t min_ns = ~std::uint64_t{0};
    std::uint64_t max_ns = 0;
    std::uint64_t bytes_sent = 0;
    std::uint64_t bytes_recv = 0;
    std::uint64_t frames_sent = 0;
    std::uint64_t frames_recv = 0;
    std::uint64_t bytes_recv_from_p2 = 0;
    std::uint64_t bytes_recv_from_peer = 0;
    std::uint64_t bytes_sent_to_p2 = 0;
    std::uint64_t bytes_sent_to_peer = 0;
    std::uint64_t payload_sent = 0;
    std::uint64_t payload_recv = 0;
    std::uint64_t rounds = 0;
    std::uint64_t prg_evals = 0;
    const std::uint64_t iters = o.warmup + o.repeat;
    for (std::uint64_t i = 0; i < iters; ++i)
    {
        // Align parties before each iteration (including first).
        if (o.pair)
        {
            if (self == role::p2)
                throw std::invalid_argument("pair flows have no p2");
            const role peer = self == role::p0 ? role::p1 : role::p0;
            (void)net.exchange_with(peer, i, net::msg::case_ok);
        }
        else if (self == role::p2)
        {
            net.to(role::p0).send(net::msg::case_ok, i);
            net.to(role::p1).send(net::msg::case_ok, i);
        }
        else
        {
            auto got = net.to(role::p2).recv<std::uint64_t>(net::msg::case_ok);
            if (got != i)
                throw std::runtime_error("repeat barrier mismatch");
        }

        // The barrier is harness traffic. The tally below is the flow itself.
        net.reset_tally();
        dpf::prg::reset_eval_count();
        const auto t0 = clock::now();
        int rc = f.run(self, net);
        const auto t1 = clock::now();
        if (rc != 0)
            return rc;
        if (i < o.warmup)
            continue;
        const auto ns = static_cast<std::uint64_t>(
            std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count());
        total_ns += ns;
        min_ns = std::min(min_ns, ns);
        max_ns = std::max(max_ns, ns);
        const auto tally = net.tally();
        const auto from_p2 = net.tally_from_p2();
        const auto from_peer = net.tally_from_peer();
        bytes_sent += tally.bytes_sent;
        bytes_recv += tally.bytes_recv;
        frames_sent += tally.frames_sent;
        frames_recv += tally.frames_recv;
        bytes_recv_from_p2 += from_p2.bytes_recv;
        bytes_recv_from_peer += from_peer.bytes_recv;
        bytes_sent_to_p2 += from_p2.bytes_sent;
        bytes_sent_to_peer += from_peer.bytes_sent;
        payload_sent += tally.payload_sent;
        payload_recv += tally.payload_recv;
        rounds += tally.exchanges;
        prg_evals += dpf::prg::eval_count();
    }
    if (o.metrics)
    {
        const std::uint64_t reps = o.repeat == 0 ? 1 : o.repeat;
        const std::uint64_t avg = total_ns / reps;
        const std::uint64_t wire_overhead =
            (bytes_sent + bytes_recv) - (payload_sent + payload_recv);
        std::string path = o.dir + "/metrics." + role_name(self);
        std::ofstream out(path, std::ios::app);
        out << "flow=" << f.name
            << " role=" << role_name(self)
            << " repeat=" << o.repeat
            << " warmup=" << o.warmup
            << " avg_ns=" << avg
            << " min_ns=" << (o.repeat ? min_ns : 0)
            << " max_ns=" << max_ns
            << " total_ns=" << total_ns
            << " bytes_sent=" << bytes_sent
            << " bytes_recv=" << bytes_recv
            << " frames_sent=" << frames_sent
            << " frames_recv=" << frames_recv
            << " bytes_recv_from_p2=" << bytes_recv_from_p2
            << " bytes_recv_from_peer=" << bytes_recv_from_peer
            << " bytes_sent_to_p2=" << bytes_sent_to_p2
            << " bytes_sent_to_peer=" << bytes_sent_to_peer
            << " payload_sent=" << payload_sent
            << " payload_recv=" << payload_recv
            << " wire_overhead=" << wire_overhead
            << " rounds=" << rounds
            << " prg_evals=" << prg_evals
            << " avg_bytes_sent=" << (bytes_sent / reps)
            << " avg_bytes_recv=" << (bytes_recv / reps)
            << " avg_frames_sent=" << (frames_sent / reps)
            << " avg_frames_recv=" << (frames_recv / reps)
            << '\n';
    }
    return 0;
}

}  // namespace

int run(role self, int argc, char ** argv)
{
    register_all_flows();
    try
    {
        auto o = parse_opts(argc, argv);
        if (o.list_only)
        {
            if (self == role::p0)
                list_flows(std::cout, o.tag);
            return 0;
        }
        if (o.dir.empty())
        {
            std::cerr << role_name(self) << ": need --dir\n";
            return 2;
        }
        if (o.case_name.empty() && o.tag.empty())
        {
            std::cerr << role_name(self) << ": need --case or --tag\n";
            return 2;
        }
        if (o.repeat == 0)
            o.repeat = 1;

        std::vector<const flow *> flows;
        if (!o.case_name.empty())
        {
            auto * f = find_flow(o.case_name);
            if (!f)
                throw std::invalid_argument("unknown flow: " + o.case_name);
            flows.push_back(f);
        }
        else
        {
            flows = select_flows(o.tag);
            if (flows.empty())
                throw std::invalid_argument("no flows match --tag " + o.tag);
        }

        auto net = o.pair ? trio::connect_pair(self, o.dir)
                          : trio::connect_local(self, o.dir);
        for (const flow * f : flows)
        {
            int rc = run_one_flow(self, *f, net, o);
            if (rc != 0)
                return rc;
        }
        return 0;
    }
    catch (const std::exception & ex)
    {
        std::cerr << role_name(self) << " failed: " << ex.what() << "\n";
        return 1;
    }
}

}  // namespace party
}  // namespace dpf

// Amalgamate bulk + extreme flows into this TU: libdpf headers define
// non-inline symbols (utils/AES/SHA256), so a second party .cpp would ODR-collide.
#include "flows_bulk.cpp"
#include "flows_extreme.cpp"
#include "flows_recent.cpp"
#include "flows_coverage.cpp"
#include "flows_gadget.cpp"