#include <gtest/gtest.h>

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <numeric>
#include <random>
#include <string>
#include <sys/wait.h>
#include <unistd.h>
#include <vector>

#include "dpf/beaver.hpp"
#include "dpf/buffered_prg.hpp"
#include "dpf/net/dealer_cursor.hpp"
#include "dpf/net/memory_sink.hpp"
#include "dpf/net/mux_sink.hpp"
#include "dpf/net/party_tape_io.hpp"
#include "dpf/net/sink_exchange.hpp"
#include "dpf/net/stream_sink.hpp"
#include "dpf/net/trio.hpp"
#include "dpf/protocol.hpp"
#include "flow_util.hpp"
#include "simde/simde/x86/avx2.h"

namespace
{

using dpf::protocol::empty_pad;
using dpf::net::dealer_cursor;
using dpf::net::make_memory_sink_pair;
using dpf::net::role;

struct Round1
{
    std::pair<int, int> operator()(int input, std::uint64_t blind) const
    {
        const int b = static_cast<int>(blind & 0xffffu);
        return {input - b, b};
    }
};

struct Round2
{
    int operator()(int forward, int peer_swap, std::uint64_t blind,
        empty_pad) const
    {
        (void)blind;
        return forward + peer_swap;
    }
};

struct Round1Empty
{
    std::pair<int, int> operator()(int input, empty_pad) const
    {
        return {input, input};
    }
};

struct Round2Empty
{
    int operator()(int forward, int peer_swap, empty_pad, empty_pad) const
    {
        return forward + peer_swap;
    }
};

TEST(ProtocolBatch, MemoryTwoMoveShuffled)
{
    constexpr std::size_t N = 100;
    auto [sink0, sink1] = make_memory_sink_pair(N, {sizeof(int)});
    dealer_cursor dealer(N, {});
    using Prg = dpf::randomness::aes_buffered_prg<std::uint64_t, std::uint64_t>;
    const auto seed = dpf::randomness::sample_master_seed<>();
    Prg prg0(seed);
    Prg prg1(seed);

    auto s0 = dpf::protocol::make_two_move_session<int, int, std::uint64_t,
        empty_pad, int>(N, sink0, dealer, prg0, Round1{}, Round2{});
    auto s1 = dpf::protocol::make_two_move_session<int, int, std::uint64_t,
        empty_pad, int>(N, sink1, dealer, prg1, Round1{}, Round2{});

    std::vector<std::size_t> order(N);
    std::iota(order.begin(), order.end(), 0);
    std::mt19937 rng(0xC0FFEE);
    std::shuffle(order.begin(), order.end(), rng);

    for (std::size_t i = 0; i < N; ++i)
    {
        s0.submit(order[i], static_cast<int>(order[i] + 1));
        s1.submit(i, static_cast<int>(i + 7));
    }
    s0.drive();
    s1.drive();

    for (std::size_t i = 0; i < N; ++i)
    {
        // Each party reconstructs input0 + input1 from shares of the blind.
        // Round1: fwd = in - b, swap = b. Peer sees the other party's b.
        // Round2: fwd + peer_swap = (in_me - b_me) + b_peer.
        // Sum of both outputs is not simply in0+in1; check local oracle:
        const auto blind = prg0.at<0>(static_cast<std::uint64_t>(i));
        const int b = static_cast<int>(blind & 0xffffu);
        // Party 0 submitted order-permuted; rebuild expectation for index i
        // as party 0's view after drive.
        (void)b;
        EXPECT_EQ(s0.take(i), s0.take(i));
        EXPECT_EQ(s1.take(i), s1.take(i));
    }

    // Both parties used the same PRG seed, so for a fixed index the blinds
    // match. Party 0 input at index j is j+1 when submitted as order mapping:
    // s0.submit(order[k], order[k]+1) => input[order[k]] = order[k]+1
    // so input0[i] = i+1. Party 1 input[i] = i+7.
    // out0 = (i+1 - b) + b = i+1
    // out1 = (i+7 - b) + b = i+7
    for (std::size_t i = 0; i < N; ++i)
    {
        EXPECT_EQ(s0.take(i), static_cast<int>(i + 1));
        EXPECT_EQ(s1.take(i), static_cast<int>(i + 7));
    }
}

TEST(ProtocolBatch, MemoryEmptyBlindAndCorrection)
{
    constexpr std::size_t N = 16;
    auto [sink0, sink1] = make_memory_sink_pair(N, {sizeof(int)});
    dealer_cursor dealer(N, {});
    using Prg = dpf::randomness::aes_buffered_prg<std::uint64_t>;
    Prg prg(dpf::randomness::sample_master_seed<>());

    auto s0 = dpf::protocol::make_two_move_session<int, int, empty_pad,
        empty_pad, int>(N, sink0, dealer, prg, Round1Empty{}, Round2Empty{});
    auto s1 = dpf::protocol::make_two_move_session<int, int, empty_pad,
        empty_pad, int>(N, sink1, dealer, prg, Round1Empty{}, Round2Empty{});

    for (std::size_t i = 0; i < N; ++i)
    {
        s0.submit(i, 3);
        s1.submit(i, 4);
    }
    s0.drive();
    s1.drive();
    for (std::size_t i = 0; i < N; ++i)
    {
        // fwd=in, swap=in; out = fwd + peer_swap = in_me + in_peer
        EXPECT_EQ(s0.take(i), 3 + 4);
        EXPECT_EQ(s1.take(i), 4 + 3);
    }
}

TEST(ProtocolBatch, MemoryHoleBlocksPrefix)
{
    auto [sink0, sink1] = make_memory_sink_pair(3, {sizeof(int)});
    int v = 1;
    sink0.submit(0, 0, reinterpret_cast<const std::uint8_t *>(&v), sizeof(v));
    v = 2;
    sink0.submit(0, 2, reinterpret_cast<const std::uint8_t *>(&v), sizeof(v));
    sink0.flush();
    // Index 1 is a hole: only index 0 may be delivered.
    EXPECT_TRUE(sink1.peer_ready(0, 0));
    EXPECT_FALSE(sink1.peer_ready(0, 1));
    EXPECT_FALSE(sink1.peer_ready(0, 2));
    v = 9;
    sink0.submit(0, 1, reinterpret_cast<const std::uint8_t *>(&v), sizeof(v));
    sink0.flush();
    EXPECT_TRUE(sink1.peer_ready(0, 1));
    EXPECT_TRUE(sink1.peer_ready(0, 2));
}

TEST(ProtocolBatch, ScheduleSessionMemory)
{
    constexpr std::size_t N = 8;
    auto [sink0, sink1] = make_memory_sink_pair(N, {sizeof(std::uint64_t)});
    std::vector<std::uint64_t> local0(N);
    std::vector<std::uint64_t> local1(N);
    for (std::size_t i = 0; i < N; ++i)
    {
        local0[i] = 10 + i;
        local1[i] = 100 + i;
    }
    auto make_rounds = [&](std::vector<std::uint64_t> & local) {
        std::vector<dpf::protocol::schedule_round> rounds(1);
        rounds[0].slot_bytes = sizeof(std::uint64_t);
        rounds[0].produce = [&](std::size_t index, const std::uint8_t *,
                                std::size_t, std::uint8_t * out) {
            std::memcpy(out, &local[index], sizeof(std::uint64_t));
        };
        return rounds;
    };
    dpf::protocol::schedule_session a(N, sink0, make_rounds(local0));
    dpf::protocol::schedule_session b(N, sink1, make_rounds(local1));
    for (std::size_t i = 0; i < N; ++i)
    {
        a.submit(i);
        b.submit(N - 1 - i);
    }
    a.drive();
    b.drive();
    for (std::size_t i = 0; i < N; ++i)
    {
        EXPECT_TRUE(a.done(i));
        EXPECT_TRUE(b.done(i));
        std::uint64_t peer = 0;
        sink0.read_peer(0, i, reinterpret_cast<std::uint8_t *>(&peer),
            sizeof(peer));
        EXPECT_EQ(peer, local1[i]);
    }
}

void run_trio_two_move(bool use_stream)
{
    constexpr std::size_t N = 32;
    const std::string dir = "/tmp/libdpf-protocol-batch-XXXXXX";
    std::string path = dir;
    ASSERT_NE(mkdtemp(path.data()), nullptr);

    auto child = [&](role self) {
        using Prg = dpf::randomness::aes_buffered_prg<std::uint64_t, std::uint64_t>;
        const auto seed = simde_mm_setzero_si128();
        Prg prg(seed);
        dealer_cursor dealer(N, {});
        std::unique_ptr<dpf::net::RoundSink> sink;
        dpf::net::trio net;
        if (use_stream)
        {
            auto s = dpf::net::connect_stream_sink(self,
                self == role::p0 ? role::p1 : role::p0, path, N,
                {sizeof(int)});
            struct holder : dpf::net::RoundSink
            {
                dpf::net::stream_sink inner;
                explicit holder(dpf::net::stream_sink s) : inner(std::move(s)) {}
                std::size_t count() const noexcept override
                {
                    return inner.count();
                }
                std::size_t rounds() const noexcept override
                {
                    return inner.rounds();
                }
                std::size_t slot_bytes(std::uint16_t r) const override
                {
                    return inner.slot_bytes(r);
                }
                void submit(std::uint16_t r, std::size_t i,
                    const std::uint8_t * b, std::size_t n) override
                {
                    inner.submit(r, i, b, n);
                }
                bool peer_ready(std::uint16_t r, std::size_t i) const override
                {
                    return inner.peer_ready(r, i);
                }
                void read_peer(std::uint16_t r, std::size_t i, std::uint8_t * o,
                    std::size_t n) const override
                {
                    inner.read_peer(r, i, o, n);
                }
                void flush() override { inner.flush(); }
                void flush_round(std::uint16_t r) override
                {
                    inner.flush_round(r);
                }
                void poll() override { inner.poll(); }
            };
            sink = std::make_unique<holder>(std::move(s));
        }
        else
        {
            net = dpf::net::trio::connect_pair(self, path);
            struct holder : dpf::net::RoundSink
            {
                dpf::net::mux_sink inner;
                explicit holder(dpf::net::mux_sink s) : inner(std::move(s)) {}
                std::size_t count() const noexcept override
                {
                    return inner.count();
                }
                std::size_t rounds() const noexcept override
                {
                    return inner.rounds();
                }
                std::size_t slot_bytes(std::uint16_t r) const override
                {
                    return inner.slot_bytes(r);
                }
                void submit(std::uint16_t r, std::size_t i,
                    const std::uint8_t * b, std::size_t n) override
                {
                    inner.submit(r, i, b, n);
                }
                bool peer_ready(std::uint16_t r, std::size_t i) const override
                {
                    return inner.peer_ready(r, i);
                }
                void read_peer(std::uint16_t r, std::size_t i, std::uint8_t * o,
                    std::size_t n) const override
                {
                    inner.read_peer(r, i, o, n);
                }
                void flush() override { inner.flush(); }
                void flush_round(std::uint16_t r) override
                {
                    inner.flush_round(r);
                }
                void poll() override { inner.poll(); }
            };
            sink = std::make_unique<holder>(
                dpf::net::make_mux_sink(net, N, {sizeof(int)}));
        }
        auto sess = dpf::protocol::make_two_move_session<int, int, std::uint64_t,
            empty_pad, int>(N, *sink, dealer, prg, Round1{}, Round2{});
        const int base = self == role::p0 ? 1 : 7;
        for (std::size_t i = 0; i < N; ++i)
            sess.submit(i, static_cast<int>(i) + base);
        sess.drive();
        for (std::size_t i = 0; i < N; ++i)
        {
            if (sess.take(i) != static_cast<int>(i) + base)
                std::_Exit(2);
        }
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    int st0 = 0;
    int st1 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0);
}

TEST(ProtocolBatch, MuxTwoMove)
{
    run_trio_two_move(false);
}

TEST(ProtocolBatch, StreamTwoMove)
{
    run_trio_two_move(true);
}

/// @brief Counts hooked trio helpers; framing stays on the mesh.
struct counting_hook : dpf::net::mesh_comm_hook
{
    int sends = 0;
    int recvs = 0;
    int exchanges = 0;
    int batches = 0;

    void send_bytes(dpf::net::trio & net, unsigned peer_id, dpf::net::msg tag,
        const void * data, std::size_t n) override
    {
        ++sends;
        mesh_comm_hook::send_bytes(net, peer_id, tag, data, n);
    }

    std::vector<std::uint8_t> recv_bytes(dpf::net::trio & net, unsigned peer_id,
        dpf::net::msg tag) override
    {
        ++recvs;
        return mesh_comm_hook::recv_bytes(net, peer_id, tag);
    }

    std::vector<std::uint8_t> exchange_bytes(dpf::net::trio & net,
        unsigned peer_id, dpf::net::msg tag, const void * data,
        std::size_t n) override
    {
        ++exchanges;
        return mesh_comm_hook::exchange_bytes(net, peer_id, tag, data, n);
    }

    std::vector<std::uint8_t> exchange_vec_bytes(dpf::net::trio & net,
        unsigned peer_id, dpf::net::msg tag, const void * data,
        std::size_t nbytes) override
    {
        ++exchanges;
        return mesh_comm_hook::exchange_vec_bytes(net, peer_id, tag, data,
            nbytes);
    }

    std::unique_ptr<dpf::net::RoundSink> batch(dpf::net::trio & net,
        std::size_t count, std::vector<std::size_t> slot_bytes) override
    {
        ++batches;
        return mesh_comm_hook::batch(net, count, std::move(slot_bytes));
    }
};

TEST(ProtocolBatch, HookRoutesDealExchangeAndBeaver)
{
    char tmpl[] = "/tmp/dpf-hook-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        counting_hook hook;
        auto net = dpf::net::trio::connect_local(self, path);
        net.set_hook(&hook);

        // Dealer pad through hooked send/recv.
        if (self == role::p2)
        {
            struct split
            {
                std::uint64_t p0 = 11;
                std::uint64_t p1 = 31;
            } s;
            net.deal(s);
        }
        else
        {
            const auto pad = net.accept_deal<std::uint64_t>();
            if (self == role::p0 && pad != 11ull)
                std::_Exit(3);
            if (self == role::p1 && pad != 31ull)
                std::_Exit(4);
        }

        // Beaver online through hooked batch (no mux_sink name at the call).
        dpf::beavers::session<std::uint64_t> s;
        auto x = s.input();
        auto y = s.input();
        auto z = s(x * y);
        if (self == role::p2)
        {
            struct Counter
            {
                int draws = 0;
                std::uint64_t operator()()
                {
                    ++draws;
                    return 0x9e3779b97f4a7c15ull * static_cast<std::uint64_t>(draws);
                }
            } rng;
            s.sample(rng);
            dpf::net::deal_session(net, s);
            if (hook.sends < 2)
                std::_Exit(5);
            std::_Exit(0);
        }
        auto tape = dpf::net::accept_session<std::uint64_t>(net);
        s.install_party(self == role::p0 ? 0u : 1u, tape);
        s.bind_party(x, self == role::p0 ? 3ull : 4ull);
        s.bind_party(y, self == role::p0 ? 5ull : 6ull);
        dpf::party::util::evaluate_online(s, net, self);
        const auto open = net.open_with(
            self == role::p0 ? role::p1 : role::p0, s.value_party(z));
        // Shares (3,4) and (5,6) reconstruct to 7 * 11 = 77.
        if (self == role::p0 && open != 77ull)
            std::_Exit(6);
        if (hook.batches < 1)
            std::_Exit(7);
        if (hook.recvs < 1)
            std::_Exit(8);
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    pid_t p2 = fork();
    ASSERT_GE(p2, 0);
    if (p2 == 0)
        child(role::p2);
    int st0 = 0, st1 = 0, st2 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    waitpid(p2, &st2, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
    EXPECT_TRUE(WIFEXITED(st2) && WEXITSTATUS(st2) == 0) << WEXITSTATUS(st2);
}

/// @brief Oversized stream sink + sink_exchange must not hang or pay one
///        exchange per idle round (flush_round path only).
TEST(ProtocolBatch, StreamFlushRoundNoIdleStorm)
{
    constexpr std::size_t R = 64;
    constexpr std::size_t steps = 8;
    char tmpl[] = "/tmp/dpf-stream-flush-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        alarm(15);
        std::vector<std::size_t> slots(R, sizeof(std::uint64_t));
        auto sink = dpf::net::connect_stream_sink(self,
            self == role::p0 ? role::p1 : role::p0, path, 1, std::move(slots));
        dpf::net::sink_exchange ex(sink, 0);
        for (std::size_t i = 0; i < steps; ++i)
        {
            const std::uint64_t mine =
                (self == role::p0 ? 0x1000ull : 0x2000ull) + i;
            const auto peer = ex(mine);
            const std::uint64_t expect =
                (self == role::p0 ? 0x2000ull : 0x1000ull) + i;
            if (peer != expect)
                std::_Exit(2);
        }
        // flush_round path: exactly one duplex per step, not one per idle round.
        if (sink.exchanges() != steps)
            std::_Exit(3);
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    int st0 = 0;
    int st1 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
}

/// @brief Mux flush_round must not invent extra link exchanges beyond the
///        number of interactive steps (no per-idle-round storm).
TEST(ProtocolBatch, MuxFlushRoundExchangeCount)
{
    constexpr std::size_t R = 64;
    constexpr std::size_t steps = 8;
    char tmpl[] = "/tmp/dpf-mux-flush-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        alarm(15);
        auto net = dpf::net::trio::connect_pair(self, path);
        std::vector<std::size_t> slots(R, sizeof(std::uint64_t));
        dpf::net::mux_sink sink = dpf::net::make_mux_sink(net, 1, std::move(slots));
        dpf::net::sink_exchange ex(sink, 0);
        for (std::size_t i = 0; i < steps; ++i)
        {
            const std::uint64_t mine =
                (self == role::p0 ? 0xA000ull : 0xB000ull) + i;
            const auto peer = ex(mine);
            const std::uint64_t expect =
                (self == role::p0 ? 0xB000ull : 0xA000ull) + i;
            if (peer != expect)
                std::_Exit(2);
        }
        if (sink.exchanges() != steps)
            std::_Exit(3);
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    int st0 = 0;
    int st1 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
}

/// @brief Memory sink_exchange across a padded round budget (ds-walk sized).
TEST(ProtocolBatch, MemorySinkExchangePaddedRounds)
{
    const auto slots = dpf::net::ds_walk_slot_bytes(8, false);
    ASSERT_GE(slots.size(), 16u);
    auto [a, b] = make_memory_sink_pair(1, slots);
    // Interleave submits then one flush_round — memory hub moves both directions.
    for (std::size_t i = 0; i < 12; ++i)
    {
        const std::uint64_t va = 0x11ull + i;
        const std::uint64_t vb = 0x22ull + i;
        // memory_sink: one flush_round after both submits (shared hub).
        std::vector<std::uint8_t> ba(slots[i], 0);
        std::vector<std::uint8_t> bb(slots[i], 0);
        std::memcpy(ba.data(), &va, sizeof(va));
        std::memcpy(bb.data(), &vb, sizeof(vb));
        a.submit(static_cast<std::uint16_t>(i), 0, ba.data(), ba.size());
        b.submit(static_cast<std::uint16_t>(i), 0, bb.data(), bb.size());
        a.flush_round(static_cast<std::uint16_t>(i));
        EXPECT_TRUE(a.peer_ready(static_cast<std::uint16_t>(i), 0));
        EXPECT_TRUE(b.peer_ready(static_cast<std::uint16_t>(i), 0));
        std::vector<std::uint8_t> ra(slots[i], 0);
        std::vector<std::uint8_t> rb(slots[i], 0);
        a.read_peer(static_cast<std::uint16_t>(i), 0, ra.data(), slots[i]);
        b.read_peer(static_cast<std::uint16_t>(i), 0, rb.data(), slots[i]);
        std::uint64_t pa = 0;
        std::uint64_t pb = 0;
        std::memcpy(&pa, ra.data(), sizeof(pa));
        std::memcpy(&pb, rb.data(), sizeof(pb));
        EXPECT_EQ(pa, vb);
        EXPECT_EQ(pb, va);
    }
}

/// @brief Hooked beaver online on an oversized batch must not exhaust rounds
///        or deadlock under a wall-clock alarm.
TEST(ProtocolBatch, HookBeaverOnlinePaddedBatch)
{
    char tmpl[] = "/tmp/dpf-hook-pad-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        alarm(20);
        counting_hook hook;
        auto net = dpf::net::trio::connect_local(self, path);
        net.set_hook(&hook);

        dpf::beavers::session<std::uint64_t> s;
        std::vector<dpf::beavers::session<std::uint64_t>::wire> xs;
        xs.reserve(6);
        for (int i = 0; i < 6; ++i)
            xs.push_back(s.input());
        auto z = s(xs[0] * xs[1] + xs[2] * xs[3] + xs[4] * xs[5]);
        if (self == role::p2)
        {
            struct Counter
            {
                int draws = 0;
                std::uint64_t operator()()
                {
                    ++draws;
                    return 0x9e3779b97f4a7c15ull
                        * static_cast<std::uint64_t>(draws);
                }
            } rng;
            s.sample(rng);
            dpf::net::deal_session(net, s);
            std::_Exit(0);
        }
        auto tape = dpf::net::accept_session<std::uint64_t>(net);
        s.install_party(self == role::p0 ? 0u : 1u, tape);
        for (int i = 0; i < 6; ++i)
            s.bind_party(xs[static_cast<std::size_t>(i)],
                self == role::p0 ? static_cast<std::uint64_t>(i + 1)
                                 : static_cast<std::uint64_t>(10 + i));
        // Pad far beyond max_ready_round to mimic ds_walk oversizing.
        const int need = std::max(1, s.max_ready_round());
        std::vector<std::size_t> slots(
            static_cast<std::size_t>(need) + 48,
            sizeof(std::uint32_t) + 64 * sizeof(std::uint64_t));
        auto sink = net.batch(1, std::move(slots));
        dpf::party::util::evaluate_online_on_sink(s, *sink, 0);
        const auto open = net.open_with(
            self == role::p0 ? role::p1 : role::p0, s.value_party(z));
        // (1*11)+(2*12)+(3*13) tied to party shares: p0 binds 1..6, p1 10..15
        // products reconstruct (1+10)*(2+11) wait — each wire is additive share.
        // x_i = p0_i + p1_i = (i+1)+(10+i) = 11+2i
        // z = x0*x1 + x2*x3 + x4*x5
        //   = 11*13 + 15*17 + 19*21 = 143 + 255 + 399 = 797
        if (self == role::p0 && open != 797ull)
            std::_Exit(4);
        if (hook.batches < 1)
            std::_Exit(5);
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    pid_t p2 = fork();
    ASSERT_GE(p2, 0);
    if (p2 == 0)
        child(role::p2);
    int st0 = 0, st1 = 0, st2 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    waitpid(p2, &st2, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
    EXPECT_TRUE(WIFEXITED(st2) && WEXITSTATUS(st2) == 0) << WEXITSTATUS(st2);
}

/// @brief Wide shallow circuits must not allocate O(wires) round windows.
TEST(ProtocolBatch, BeaverOnlineRoundBudgetTight)
{
    char tmpl[] = "/tmp/dpf-beaver-budget-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        alarm(30);
        struct hold_hook : dpf::net::mesh_comm_hook
        {
            std::size_t last_rounds = 0;
            std::uint64_t exchanges = 0;
            struct wrap : dpf::net::RoundSink
            {
                dpf::net::mux_sink inner;
                std::uint64_t * out = nullptr;
                wrap(dpf::net::mux_sink s, std::uint64_t * o)
                  : inner(std::move(s)), out(o)
                { }
                ~wrap() override
                {
                    if (out)
                        *out = inner.exchanges();
                }
                std::size_t count() const noexcept override
                {
                    return inner.count();
                }
                std::size_t rounds() const noexcept override
                {
                    return inner.rounds();
                }
                std::size_t slot_bytes(std::uint16_t r) const override
                {
                    return inner.slot_bytes(r);
                }
                void submit(std::uint16_t r, std::size_t i,
                    const std::uint8_t * b, std::size_t n) override
                {
                    inner.submit(r, i, b, n);
                }
                bool peer_ready(std::uint16_t r, std::size_t i) const override
                {
                    return inner.peer_ready(r, i);
                }
                void read_peer(std::uint16_t r, std::size_t i, std::uint8_t * o,
                    std::size_t n) const override
                {
                    inner.read_peer(r, i, o, n);
                }
                void flush() override { inner.flush(); }
                void flush_round(std::uint16_t r) override
                {
                    inner.flush_round(r);
                }
                void poll() override { inner.poll(); }
            };
            std::unique_ptr<dpf::net::RoundSink> batch(dpf::net::trio & net,
                std::size_t count, std::vector<std::size_t> slot_bytes) override
            {
                last_rounds = slot_bytes.size();
                return std::make_unique<wrap>(
                    dpf::net::make_mux_sink(net, count, std::move(slot_bytes)),
                    &exchanges);
            }
        } h;
        auto net = dpf::net::trio::connect_local(self, path);
        net.set_hook(&h);

        constexpr int N = 256;
        dpf::beavers::session<std::uint64_t> s;
        std::vector<dpf::beavers::session<std::uint64_t>::wire> xs, ys;
        xs.reserve(N);
        ys.reserve(N);
        for (int i = 0; i < N; ++i)
        {
            xs.push_back(s.input());
            ys.push_back(s.input());
        }
        auto acc = xs[0] * ys[0];
        for (int i = 1; i < N; ++i)
            acc = acc + xs[i] * ys[i];
        auto z = s(acc);

        if (self == role::p2)
        {
            if (s.max_ready_round() != 1)
                std::_Exit(2);
            struct Counter
            {
                int draws = 0;
                std::uint64_t operator()()
                {
                    ++draws;
                    return 0x9e3779b97f4a7c15ull
                        * static_cast<std::uint64_t>(draws);
                }
            } rng;
            s.sample(rng);
            dpf::net::deal_session(net, s);
            std::_Exit(0);
        }
        auto tape = dpf::net::accept_session<std::uint64_t>(net);
        s.install_party(self == role::p0 ? 0u : 1u, tape);
        for (int i = 0; i < N; ++i)
        {
            s.bind_party(xs[static_cast<std::size_t>(i)],
                self == role::p0 ? static_cast<std::uint64_t>(i + 1)
                                 : static_cast<std::uint64_t>(1000 + i));
            s.bind_party(ys[static_cast<std::size_t>(i)],
                self == role::p0 ? static_cast<std::uint64_t>(i + 3)
                                 : static_cast<std::uint64_t>(2000 + i));
        }
        dpf::party::util::evaluate_online(s, net, self);
        // Depth-1 product sum: budget is O(1), not O(N).
        const int expect_budget = dpf::party::util::beaver_batch_round_budget(
            s.max_ready_round(), s.wire_count());
        if (static_cast<int>(h.last_rounds) != expect_budget)
            std::_Exit(3);
        if (h.last_rounds > 48)
            std::_Exit(4);
        if (h.exchanges == 0 || h.exchanges > h.last_rounds)
            std::_Exit(5);
        const auto open = net.open_with(
            self == role::p0 ? role::p1 : role::p0, s.value_party(z));
        // Reconstruct expected: sum_i (i+1+1000+i)*(i+3+2000+i)
        // = sum_i (1001+2i)*(2003+2i)
        if (self == role::p0)
        {
            std::uint64_t expect = 0;
            for (int i = 0; i < N; ++i)
            {
                const std::uint64_t xi =
                    static_cast<std::uint64_t>(i + 1) + static_cast<std::uint64_t>(1000 + i);
                const std::uint64_t yi =
                    static_cast<std::uint64_t>(i + 3) + static_cast<std::uint64_t>(2000 + i);
                expect += xi * yi;
            }
            if (open != expect)
                std::_Exit(6);
        }
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    pid_t p2 = fork();
    ASSERT_GE(p2, 0);
    if (p2 == 0)
        child(role::p2);
    int st0 = 0, st1 = 0, st2 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    waitpid(p2, &st2, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
    EXPECT_TRUE(WIFEXITED(st2) && WEXITSTATUS(st2) == 0) << WEXITSTATUS(st2);
}

/// @brief Multi-lane sequential sink_exchange (signum pattern) must not hang.
TEST(ProtocolBatch, MultiLaneSequentialSinkExchange)
{
    constexpr std::size_t lanes = 4;
    constexpr std::size_t steps = 6;
    char tmpl[] = "/tmp/dpf-multilane-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        alarm(20);
        auto net = dpf::net::trio::connect_pair(self, path);
        std::vector<std::size_t> slots(steps, sizeof(std::uint64_t));
        auto sink = dpf::net::make_mux_sink(net, lanes, std::move(slots));
        for (std::size_t lane = 0; lane < lanes; ++lane)
        {
            dpf::net::sink_exchange ex(sink, lane);
            for (std::size_t r = 0; r < steps; ++r)
            {
                const std::uint64_t mine =
                    (self == role::p0 ? 0x10ull : 0x20ull) + lane * 100 + r;
                const auto peer = ex(mine);
                const std::uint64_t expect =
                    (self == role::p0 ? 0x20ull : 0x10ull) + lane * 100 + r;
                if (peer != expect)
                    std::_Exit(2);
            }
        }
        // Sequential lanes: steps exchanges per lane.
        if (sink.exchanges() != lanes * steps)
            std::_Exit(3);
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    int st0 = 0, st1 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
}

/// @brief Mux flush() bundles many pending rounds into one exchange.
TEST(ProtocolBatch, MuxFlushBundlesPendingRounds)
{
    constexpr std::size_t R = 16;
    char tmpl[] = "/tmp/dpf-mux-bundle-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        alarm(15);
        auto net = dpf::net::trio::connect_pair(self, path);
        std::vector<std::size_t> slots(R, sizeof(std::uint32_t));
        auto sink = dpf::net::make_mux_sink(net, 1, std::move(slots));
        for (std::uint16_t r = 0; r < R; ++r)
        {
            const std::uint32_t v =
                (self == role::p0 ? 100u : 200u) + r;
            sink.submit(r, 0, reinterpret_cast<const std::uint8_t *>(&v),
                sizeof(v));
        }
        sink.flush();  // one exchange for all R rounds
        if (sink.exchanges() != 1)
            std::_Exit(2);
        for (std::uint16_t r = 0; r < R; ++r)
        {
            if (!sink.peer_ready(r, 0))
                std::_Exit(3);
            std::uint32_t peer = 0;
            sink.read_peer(r, 0, reinterpret_cast<std::uint8_t *>(&peer),
                sizeof(peer));
            const std::uint32_t expect =
                (self == role::p0 ? 200u : 100u) + r;
            if (peer != expect)
                std::_Exit(4);
        }
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    int st0 = 0, st1 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
}

/// @brief Stream flush() after multi-round submits must rendezvous (no hang).
TEST(ProtocolBatch, StreamFlushAllRoundsRendezvous)
{
    constexpr std::size_t R = 8;
    char tmpl[] = "/tmp/dpf-stream-bundle-XXXXXX";
    ASSERT_NE(mkdtemp(tmpl), nullptr);
    const std::string path = tmpl;

    auto child = [&](role self) {
        alarm(20);
        std::vector<std::size_t> slots(R, sizeof(std::uint32_t));
        auto sink = dpf::net::connect_stream_sink(self,
            self == role::p0 ? role::p1 : role::p0, path, 1, std::move(slots));
        // Stagger: p0 fills even rounds first, p1 fills odd — then both flush().
        for (std::uint16_t r = 0; r < R; ++r)
        {
            const bool mine =
                (self == role::p0) ? (r % 2 == 0) : (r % 2 == 1);
            if (!mine)
                continue;
            const std::uint32_t v = 50u + r;
            sink.submit(r, 0, reinterpret_cast<const std::uint8_t *>(&v),
                sizeof(v));
        }
        sink.flush();
        // After global flush both should have peer data for the rounds the
        // peer filled; own rounds may still be empty on the peer side until
        // the peer also submitted — second pass fills the rest.
        for (std::uint16_t r = 0; r < R; ++r)
        {
            const bool peer_filled =
                (self == role::p0) ? (r % 2 == 1) : (r % 2 == 0);
            if (peer_filled && !sink.peer_ready(r, 0))
                std::_Exit(2);
        }
        for (std::uint16_t r = 0; r < R; ++r)
        {
            const bool mine =
                (self == role::p0) ? (r % 2 == 1) : (r % 2 == 0);
            if (!mine)
                continue;
            const std::uint32_t v = 50u + r;
            sink.submit(r, 0, reinterpret_cast<const std::uint8_t *>(&v),
                sizeof(v));
        }
        sink.flush();
        for (std::uint16_t r = 0; r < R; ++r)
        {
            if (!sink.peer_ready(r, 0))
                std::_Exit(3);
            std::uint32_t peer = 0;
            sink.read_peer(r, 0, reinterpret_cast<std::uint8_t *>(&peer),
                sizeof(peer));
            if (peer != 50u + r)
                std::_Exit(4);
        }
        std::_Exit(0);
    };

    pid_t p0 = fork();
    ASSERT_GE(p0, 0);
    if (p0 == 0)
        child(role::p0);
    pid_t p1 = fork();
    ASSERT_GE(p1, 0);
    if (p1 == 0)
        child(role::p1);
    int st0 = 0, st1 = 0;
    waitpid(p0, &st0, 0);
    waitpid(p1, &st1, 0);
    EXPECT_TRUE(WIFEXITED(st0) && WEXITSTATUS(st0) == 0) << WEXITSTATUS(st0);
    EXPECT_TRUE(WIFEXITED(st1) && WEXITSTATUS(st1) == 0) << WEXITSTATUS(st1);
}

}  // namespace