#include <gtest/gtest.h>

#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <map>
#include <mutex>
#include <string>
#include <thread>
#include <vector>

#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include "dpf/app_flow.hpp"
#include "dpf/compose.hpp"
#include "dpf/compose_async.hpp"
#include "dpf/launch.hpp"
#include "dpf/net/async_round_sink.hpp"
#include "dpf/net/async_sctp_stream_array.hpp"
#include "dpf/net/async_stream_array.hpp"
#include "dpf/net/connect.hpp"
#include "dpf/net/party_session.hpp"
#include "dpf/net/policy.hpp"
#include "dpf/net/round_lane.hpp"
#include "dpf/net/stream_array.hpp"
#include "dpf/net/sync_stream_array.hpp"
#include "dpf/online_session.hpp"
#include "dpf/party_run.hpp"
#include "dpf/party_runner.hpp"
#include "dpf/run_config.hpp"

namespace
{

using dpf::protocol::domain;
using clock_type = std::chrono::steady_clock;

void put(std::vector<std::vector<std::uint8_t>> & values, dpf::protocol::node n,
    std::uint64_t v)
{
    if (values.size() <= n.id)
        values.resize(n.id + 1);
    values[n.id].assign(8, 0);
    std::memcpy(values[n.id].data(), &v, 8);
}

std::uint64_t get(const std::vector<std::vector<std::uint8_t>> & values,
    dpf::protocol::node n, std::size_t lane = 0)
{
    std::uint64_t v = 0;
    std::memcpy(&v, values[n.id].data() + 8 * lane, 8);
    return v;
}

/// Pump `io` until `n` of the counted completions arrive or 5 s pass.
void pump(asio::io_context & io, std::atomic<int> & left)
{
    const auto deadline = clock_type::now() + std::chrono::seconds(5);
    while (left.load() > 0 && clock_type::now() < deadline)
    {
        if (io.stopped())
            io.restart();
        io.run_one_for(std::chrono::milliseconds(10));
    }
}

void step(asio::io_context & io)
{
    if (io.stopped())
        io.restart();
    io.run_one_for(std::chrono::milliseconds(5));
}

unsigned short free_port()
{
    asio::io_context io;
    asio::ip::tcp::acceptor a(io);
    dpf::net::open_listener(a, 0);
    return a.local_endpoint().port();
}

/// Two parties: x0 + x1 opened, then `k` dependent steps (+1 and reopen).
struct chain
{
    dpf::protocol::plan plan;
    dpf::protocol::node x;
    dpf::protocol::node out;
};

constexpr std::uint32_t k_inc = 94001;

chain make_chain(std::size_t party, int steps)
{
    dpf::protocol::composer c(party);
    chain ch;
    ch.x = c.input(domain::a, 8);
    auto cur = ch.x;
    for (int i = 0; i < steps; ++i)
    {
        auto e = c.exchange(cur);
        cur = c.compute(k_inc, {e}, domain::a, 8);
    }
    ch.out = cur;
    ch.plan = c.schedule();
    return ch;
}

dpf::protocol::kernel_fn inc_kernel(std::chrono::milliseconds sleep = {})
{
    return [sleep](std::uint32_t, const std::vector<dpf::protocol::node> &,
               const std::vector<dpf::protocol::block_span> & inputs,
               dpf::protocol::block_span output, std::size_t lanes) {
        if (sleep.count() > 0)
            std::this_thread::sleep_for(sleep);
        for (std::size_t l = 0; l < lanes; ++l)
        {
            std::uint64_t v = 0;
            std::memcpy(&v, inputs[0].at(l), 8);
            v += 1;
            std::memcpy(output.at(l), &v, 8);
        }
    };
}

/// Expected result of `make_chain` from inputs a, b.
std::uint64_t chain_value(std::uint64_t a, std::uint64_t b, int steps)
{
    // Each step opens (both parties hold the prior value) and adds 1.
    std::uint64_t v = a + b + 1;
    for (int i = 1; i < steps; ++i)
        v = 2 * v + 1;
    return v;
}

}  // namespace

// ---------------------------------------------------------------------------
// Explicit configuration
// ---------------------------------------------------------------------------

TEST(RunConfig, EnvAndArgsShareKeys)
{
    ::setenv("DPF_LANES", "4", 1);
    ::setenv("DPF_WIRE_WINDOW", "65536", 1);
    auto cfg = dpf::app::run_config::from_env();
    ::unsetenv("DPF_LANES");
    ::unsetenv("DPF_WIRE_WINDOW");
    EXPECT_EQ(cfg.n_lanes, 4u);
    EXPECT_EQ(cfg.policy.window_bytes, 65536u);
    const char * argv[] = {"x", "--transport=parallel", "--framing=always",
        "--instances=3", "--sndbuf=262144", "--trials=5", "positional"};
    const auto rest = cfg.apply_args(7, const_cast<char **>(argv));
    ASSERT_EQ(rest.size(), 1u);
    EXPECT_EQ(rest[0], "positional");
    EXPECT_EQ(cfg.kind, dpf::net::transport::parallel);
    EXPECT_EQ(cfg.framing, dpf::net::framing_mode::always);
    EXPECT_EQ(cfg.instances, 3u);
    EXPECT_EQ(cfg.policy.socket.send_buffer, 262144);
    EXPECT_EQ(cfg.trials, 5u);
    EXPECT_NE(cfg.summary().find("transport=parallel"), std::string::npos);
    EXPECT_THROW(cfg.set("lanse", "3"), std::invalid_argument);
    EXPECT_THROW(cfg.set("lanes", "three"), std::invalid_argument);
    EXPECT_THROW(cfg.set("transport", "udp"), std::invalid_argument);
    EXPECT_THROW(cfg.set("chunk", std::to_string(std::size_t{64} << 20)),
        std::invalid_argument);
}

TEST(RoundLane, FramingModesAndLimits)
{
    dpf::net::round_lane_map a(4, 3, dpf::net::framing_mode::always);
    EXPECT_TRUE(a.framed);
    EXPECT_EQ(a.lane(2), 2u);
    dpf::net::round_lane_map b(2, 6);
    EXPECT_TRUE(b.framed);
    EXPECT_EQ(b.lane(5), 1u);
    EXPECT_THROW(dpf::net::round_lane_map(2, 6, dpf::net::framing_mode::never),
        std::invalid_argument);
    EXPECT_THROW(dpf::net::round_lane_map(1, 70000), std::invalid_argument);
    EXPECT_EQ(dpf::net::lane_count_for_rounds(20), 8u);
    EXPECT_EQ(dpf::net::lane_count_for_rounds(20, dpf::net::lanes_one_per_round),
        20u);
    EXPECT_EQ(dpf::net::lane_count_for_rounds(3, 8), 3u);
}

// ---------------------------------------------------------------------------
// Sink hello and framing
// ---------------------------------------------------------------------------

namespace
{

std::string sink_mismatch(std::vector<std::size_t> s0, std::vector<std::size_t> s1,
    std::size_t c0, std::size_t c1, dpf::net::framing_mode f0,
    dpf::net::framing_mode f1)
{
    asio::io_context io;
    auto ends = dpf::net::make_async_memory_stream_pair(io, 4);
    dpf::net::sink_options o0, o1;
    o0.framing = f0;
    o1.framing = f1;
    dpf::net::async_round_sink a(ends.first, s0, c0, o0);
    dpf::net::async_round_sink b(ends.second, s1, c1, o1);
    const auto deadline = clock_type::now() + std::chrono::seconds(2);
    while (clock_type::now() < deadline)
    {
        try
        {
            (void)a.peer_ready(0, 0);
            (void)b.peer_ready(0, 0);
        }
        catch (const std::exception & e)
        {
            return e.what();
        }
        step(io);
    }
    return {};
}

}  // namespace

TEST(Sink, HelloNamesTheDisagreement)
{
    using fm = dpf::net::framing_mode;
    auto rounds = sink_mismatch({8, 8}, {8}, 1, 1, fm::automatic, fm::automatic);
    EXPECT_NE(rounds.find("rounds"), std::string::npos) << rounds;
    auto widths = sink_mismatch({8, 8}, {8, 16}, 1, 1, fm::automatic, fm::automatic);
    EXPECT_NE(widths.find("slot widths"), std::string::npos) << widths;
    auto inst = sink_mismatch({8}, {8}, 2, 1, fm::automatic, fm::automatic);
    EXPECT_NE(inst.find("instances 1 vs 2"), std::string::npos) << inst;
    auto frame = sink_mismatch({8}, {8}, 1, 1, fm::always, fm::automatic);
    EXPECT_NE(frame.find("framing"), std::string::npos) << frame;
}

TEST(Sink, FramedPartialPrefixReadyPerInstance)
{
    asio::io_context io;
    auto ends = dpf::net::make_async_memory_stream_pair(io, 1);
    dpf::net::sink_options so;
    so.framing = dpf::net::framing_mode::always;
    dpf::net::async_round_sink a(ends.first, {8}, 4, so);
    dpf::net::async_round_sink b(ends.second, {8}, 4, so);
    std::uint8_t slot[8] = {1};
    a.submit(0, 0, slot, 8);
    a.submit(0, 1, slot, 8);
    a.flush_round(0);
    const auto deadline = clock_type::now() + std::chrono::seconds(2);
    while (!b.peer_ready(0, 1) && clock_type::now() < deadline)
        step(io);
    EXPECT_TRUE(b.peer_ready(0, 0));
    EXPECT_TRUE(b.peer_ready(0, 1));
    EXPECT_FALSE(b.peer_ready(0, 2));
    a.submit(0, 2, slot, 8);
    a.submit(0, 3, slot, 8);
    a.flush_round(0);
    while (!b.peer_ready(0, 3) && clock_type::now() < deadline)
        step(io);
    EXPECT_TRUE(b.peer_ready(0, 3));
    EXPECT_TRUE(a.framed());
}

TEST(Sink, PeerFailureSurfacesImmediately)
{
    asio::io_context io;
    auto ends = dpf::net::make_async_memory_stream_pair(io, 2);
    dpf::net::async_round_sink a(ends.first, {8, 8}, 1);
    {
        dpf::net::async_round_sink b(ends.second, {8, 8}, 1);
        for (int i = 0; i < 50; ++i)
            io.poll();
    }
    ends.second.close();
    const auto t0 = clock_type::now();
    bool threw = false;
    try
    {
        for (int i = 0; i < 1000 && !threw; ++i)
        {
            (void)a.peer_ready(0, 0);
            a.wait_io_for(std::chrono::milliseconds(5));
        }
    }
    catch (const std::system_error &)
    {
        threw = true;
    }
    EXPECT_TRUE(threw);
    EXPECT_LT(clock_type::now() - t0, std::chrono::seconds(2));
}

// ---------------------------------------------------------------------------
// Drive-loop budgets
// ---------------------------------------------------------------------------

TEST(Drive, HealthyPlanLongerThanWaitTimeout)
{
    const int steps = 6;
    auto c0 = make_chain(0, steps);
    auto c1 = make_chain(1, steps);
    std::map<std::uint32_t, dpf::protocol::kernel_fn> k{
        {k_inc, inc_kernel(std::chrono::milliseconds(25))}};
    std::vector<std::vector<std::uint8_t>> v0(c0.plan.nodes().size()),
        v1(c1.plan.nodes().size());
    put(v0, c0.x, 1);
    put(v1, c1.x, 2);
    dpf::protocol::drive_options opt;
    opt.wait_timeout = std::chrono::milliseconds(100);
    const auto t0 = clock_type::now();
    dpf::protocol::drive_both_on_async_streams(c0.plan, c1.plan, v0, v1, k, 1, opt);
    EXPECT_GT(clock_type::now() - t0, std::chrono::milliseconds(120));
    EXPECT_EQ(get(v0, c0.out), chain_value(1, 2, steps));
}

TEST(Drive, EdgeAndRoundBudgetsNameTheWait)
{
    auto c0 = make_chain(0, 2);
    asio::io_context io0, io1;
    auto ends = dpf::net::make_async_dual_memory_stream_pair(io0, io1, 2);
    std::vector<std::vector<std::uint8_t>> v0(c0.plan.nodes().size());
    put(v0, c0.x, 1);
    std::map<std::uint32_t, dpf::protocol::kernel_fn> k{{k_inc, inc_kernel()}};

    dpf::protocol::drive_options edge;
    edge.wait_timeout = std::chrono::seconds(20);
    edge.edge_timeout[dpf::net::edge_peer] = std::chrono::milliseconds(40);
    auto t0 = clock_type::now();
    try
    {
        dpf::protocol::drive_plan_on_async_streams(c0.plan, ends.first, v0, k, 0, 1,
            edge);
        FAIL() << "one-sided drive should time out";
    }
    catch (const std::runtime_error & e)
    {
        const std::string what = e.what();
        EXPECT_NE(what.find("peer"), std::string::npos) << what;
        EXPECT_NE(what.find("budget 40"), std::string::npos) << what;
    }
    EXPECT_LT(clock_type::now() - t0, std::chrono::seconds(5));

    auto ends2 = dpf::net::make_async_dual_memory_stream_pair(io0, io1, 2);
    std::vector<std::vector<std::uint8_t>> v1(c0.plan.nodes().size());
    put(v1, c0.x, 1);
    dpf::protocol::drive_options round;
    round.wait_timeout = std::chrono::seconds(20);
    round.round_timeout[0] = std::chrono::milliseconds(30);
    t0 = clock_type::now();
    EXPECT_THROW(dpf::protocol::drive_plan_on_async_streams(c0.plan, ends2.first,
                     v1, k, 0, 1, round),
        std::runtime_error);
    EXPECT_LT(clock_type::now() - t0, std::chrono::seconds(5));
}

// ---------------------------------------------------------------------------
// Backends: windows, stats, scheduling, graceful close
// ---------------------------------------------------------------------------

TEST(Memory, WindowSetterAndGracefulClose)
{
    asio::io_context io;
    auto ends = dpf::net::make_async_memory_stream_pair(io, 1);
    ends.first.set_window_bytes(100);
    EXPECT_EQ(ends.first.window_bytes(), 100u);
    std::vector<std::uint8_t> a(300, 1), b(300, 2);
    std::atomic<int> left{2};
    ends.first.async_write(0, a.data(), a.size(), [&](auto) { --left; });
    ends.first.async_write(0, b.data(), b.size(), [&](auto) { --left; });
    EXPECT_EQ(ends.first.buffered_bytes(), 600u);
    ends.first.close();
    std::vector<std::uint8_t> got(600);
    std::atomic<int> rd{1};
    std::error_code rec;
    ends.second.async_read(0, got.data(), got.size(), [&](const std::error_code & ec) {
        rec = ec;
        --rd;
    });
    pump(io, rd);
    EXPECT_FALSE(rec);
    EXPECT_EQ(got[0], 1);
    EXPECT_EQ(got[599], 2);
    const auto st = ends.first.stats();
    EXPECT_EQ(st.payload_out, 600u);
    EXPECT_TRUE(st.closed);
}

namespace
{

/// Localhost mux pair, both ends on `io`.
struct mux_pair
{
    std::unique_ptr<dpf::net::async_mux_stream_array> a;
    std::unique_ptr<dpf::net::async_mux_stream_array> b;
};

mux_pair make_mux_pair(asio::io_context & io, std::size_t lanes,
    const dpf::net::wire_policy & pol = {})
{
    asio::ip::tcp::acceptor acc(io);
    dpf::net::open_listener(acc, 0);
    asio::ip::tcp::socket s0(io), s1(io);
    std::thread t([&] {
        dpf::net::connect_until(s1, "127.0.0.1", acc.local_endpoint().port(),
            std::chrono::seconds(5));
    });
    dpf::net::accept_until(acc, s0, std::chrono::seconds(5));
    t.join();
    mux_pair out;
    out.a = std::make_unique<dpf::net::async_mux_stream_array>(io, std::move(s0),
        0, 1, lanes, pol);
    out.b = std::make_unique<dpf::net::async_mux_stream_array>(io, std::move(s1),
        1, 0, lanes, pol);
    return out;
}

}  // namespace

TEST(Mux, StatsCountHeaders)
{
    asio::io_context io;
    auto m = make_mux_pair(io, 2);
    std::vector<std::uint8_t> out(100, 7), in(100);
    std::atomic<int> left{2};
    m.a->async_write(0, out.data(), out.size(), [&](auto) { --left; });
    m.b->async_read(0, in.data(), in.size(), [&](auto) { --left; });
    pump(io, left);
    const auto s = m.a->stats();
    EXPECT_EQ(s.payload_out, 100u);
    EXPECT_EQ(s.bytes_out, 110u);
    EXPECT_EQ(s.frames_out, 1u);
    EXPECT_EQ(in, out);
}

TEST(Mux, SmallWriteIsNotQueuedBehindLargeWrite)
{
    asio::io_context io;
    dpf::net::wire_policy pol;
    pol.chunk_bytes = 16 << 10;
    pol.window_bytes = 0;
    auto m = make_mux_pair(io, 2, pol);
    std::vector<std::uint8_t> big(8u << 20, 1), small(8, 2);
    std::vector<std::uint8_t> big_in(big.size()), small_in(8);
    std::atomic<int> left{2};
    std::atomic<int> order{0};
    int big_at = 0, small_at = 0;
    m.a->async_write(0, big.data(), big.size(), [](auto) {});
    m.a->async_write(1, small.data(), small.size(), [](auto) {});
    m.b->async_read(0, big_in.data(), big_in.size(), [&](auto) {
        big_at = ++order;
        --left;
    });
    m.b->async_read(1, small_in.data(), small_in.size(), [&](auto) {
        small_at = ++order;
        --left;
    });
    pump(io, left);
    EXPECT_EQ(small_at, 1);
    EXPECT_EQ(big_at, 2);
    EXPECT_EQ(small_in, small);
    EXPECT_EQ(big_in, big);
    EXPECT_GT(m.a->stats().frames_out, 100u);
}

TEST(Mux, DestroyAfterWriteStillDelivers)
{
    asio::io_context io;
    auto m = make_mux_pair(io, 1);
    std::vector<std::uint8_t> out(4u << 20, 9), in(out.size());
    m.a->async_write(0, out.data(), out.size(), [](auto) {});
    m.a.reset();
    std::atomic<int> left{1};
    std::error_code rec;
    m.b->async_read(0, in.data(), in.size(), [&](const std::error_code & ec) {
        rec = ec;
        --left;
    });
    pump(io, left);
    EXPECT_FALSE(rec);
    EXPECT_EQ(in, out);
}

TEST(Parallel, PerLaneWindows)
{
    asio::io_context io;
    std::atomic<unsigned short> port{0};
    std::vector<asio::ip::tcp::socket> a, b;
    std::thread t([&] { a = dpf::net::accept_parallel_tcp(io, port, 2); });
    while (port.load() == 0)
        std::this_thread::yield();
    b = dpf::net::connect_parallel_tcp(io, "127.0.0.1", port, 2);
    t.join();
    dpf::net::wire_policy pol;
    pol.window_bytes = 1000;
    dpf::net::async_parallel_stream_array pa(io, std::move(a), pol);
    dpf::net::async_parallel_stream_array pb(io, std::move(b), pol);
    EXPECT_EQ(pa.lane_window_bytes(0), 1000u);
    EXPECT_EQ(pa.window_bytes(), 2000u);
    pa.set_window_bytes(500);
    EXPECT_EQ(pa.lane_window_bytes(1), 500u);
    std::vector<std::uint8_t> out(64, 3), in(64);
    std::atomic<int> left{2};
    pa.async_write(1, out.data(), out.size(), [&](auto) { --left; });
    pb.async_read(1, in.data(), in.size(), [&](auto) { --left; });
    pump(io, left);
    EXPECT_EQ(in, out);
    EXPECT_EQ(pa.stats().bytes_out, 68u);
}

// ---------------------------------------------------------------------------
// Synchronous mux: same wire as async mux, no flush deadlock
// ---------------------------------------------------------------------------

TEST(SyncMux, InteroperatesWithAsyncMux)
{
    // Blocking sync faces need one side to read first so the peer's write can
    // complete (the old fd mux pumped both directions while blocked).
    std::atomic<unsigned short> port{0};
    std::vector<std::uint8_t> from_sync(3000, 5), from_async(5000, 6);
    std::vector<std::uint8_t> got_sync(from_async.size()), got_async(from_sync.size());
    std::exception_ptr err;
    std::mutex err_mu;
    auto note = [&](std::exception_ptr e) {
        std::lock_guard<std::mutex> lock(err_mu);
        if (!err)
            err = std::move(e);
    };
    std::thread sync_side([&] {
        try
        {
            dpf::run::tcp_pair_mux(1, "127.0.0.1", port, 2,
                [&](unsigned, dpf::net::mux_stream_array & mux) {
                    mux.read(0, got_sync.data(), got_sync.size());
                    mux.write(1, from_sync.data(), from_sync.size());
                    mux.flush(1);
                });
        }
        catch (...)
        {
            note(std::current_exception());
        }
    });
    try
    {
        dpf::run::tcp_pair_mux(0, "127.0.0.1", port, 2,
            [&](unsigned, dpf::net::mux_stream_array & mux) {
                mux.write(0, from_async.data(), from_async.size());
                mux.flush(0);
                mux.read(1, got_async.data(), got_async.size());
            });
    }
    catch (...)
    {
        note(std::current_exception());
    }
    sync_side.join();
    if (err)
        std::rethrow_exception(err);
    EXPECT_EQ(got_async, from_sync);
    EXPECT_EQ(got_sync, from_async);
}

TEST(SyncMux, LargeCrossFlushDoesNotDeadlock)
{
    std::atomic<unsigned short> port{0};
    std::vector<std::uint8_t> big(4u << 20);
    for (std::size_t i = 0; i < big.size(); ++i)
        big[i] = static_cast<std::uint8_t>(i);
    std::atomic<int> ok{0};
    auto side = [&](unsigned party) {
        dpf::run::tcp_pair_mux(party, "127.0.0.1", port, 1,
            [&](unsigned, dpf::net::mux_stream_array & mux) {
                mux.write(0, big.data(), big.size());
                mux.flush(0);
                std::vector<std::uint8_t> in(big.size());
                mux.read(0, in.data(), in.size());
                if (in == big)
                    ++ok;
            });
    };
    std::thread t0([&] { side(0); });
    std::thread t1([&] { side(1); });
    t0.join();
    t1.join();
    EXPECT_EQ(ok.load(), 2);
}

// ---------------------------------------------------------------------------
// Sessions: static tables, deadlines, per-edge transport, reconnect
// ---------------------------------------------------------------------------

TEST(Session, ConnectHasADeadline)
{
    asio::io_context io;
    asio::ip::tcp::socket s(io);
    const auto port = free_port();
    const auto t0 = clock_type::now();
    EXPECT_THROW(dpf::net::connect_until(s, "127.0.0.1", port,
                     std::chrono::milliseconds(200)),
        std::system_error);
    EXPECT_LT(clock_type::now() - t0, std::chrono::seconds(2));
}

TEST(Session, StaticTableAnyStartOrder)
{
    const std::vector<dpf::net::peer_address> table = {
        {"127.0.0.1", free_port()}, {"127.0.0.1", free_port()}};
    std::atomic<int> ok{0};
    auto run = [&](unsigned me, std::chrono::milliseconds delay) {
        std::this_thread::sleep_for(delay);
        asio::io_context io;
        dpf::net::party_session s(io, me, 2, std::size_t{1});
        s.join(table);
        std::uint64_t mine = me + 10, got = 0;
        std::atomic<int> left{2};
        s.peer(1 - me).async_write(0, &mine, 8, [&](auto) { --left; });
        s.peer(1 - me).async_read(0, &got, 8, [&](auto) { --left; });
        pump(io, left);
        if (got == 11u - me)
            ++ok;
    };
    std::thread late([&] { run(0, std::chrono::milliseconds(300)); });
    std::thread early([&] { run(1, std::chrono::milliseconds(0)); });
    late.join();
    early.join();
    EXPECT_EQ(ok.load(), 2);
}

TEST(Session, TransportMismatchIsNamed)
{
    auto ports = dpf::net::make_mesh_ports(2);
    std::string e0, e1;
    auto run = [&](unsigned me, dpf::net::transport t, std::string & err) {
        try
        {
            asio::io_context io;
            dpf::net::session_options so;
            so.n_lanes = 2;
            so.limits.accept = std::chrono::milliseconds(2000);
            dpf::net::party_session s(io, me, 2, so);
            s.set_edge_transport(1 - me, t);
            s.join("127.0.0.1", ports);
        }
        catch (const std::exception & e)
        {
            err = e.what();
        }
    };
    std::thread t0([&] { run(0, dpf::net::transport::mux, e0); });
    std::thread t1([&] { run(1, dpf::net::transport::parallel, e1); });
    t0.join();
    t1.join();
    const std::string both = e0 + " | " + e1;
    EXPECT_NE(both.find("transport"), std::string::npos) << both;
}

TEST(Session, ParallelAndSctpEdges)
{
    std::vector<dpf::net::transport> kinds = {dpf::net::transport::parallel};
    if (dpf::net::sctp_available())
        kinds.push_back(dpf::net::transport::sctp);
    for (auto kind : kinds)
    {
        auto ports = dpf::net::make_mesh_ports(2);
        std::atomic<int> ok{0};
        std::string err;
        std::mutex mu;
        auto run = [&](unsigned me) {
            try
            {
                asio::io_context io;
                dpf::net::session_options so;
                so.n_lanes = 3;
                so.kind = kind;
                // SCTP links cannot be encrypted.
                so.security.encrypt = kind != dpf::net::transport::sctp;
                dpf::net::party_session s(io, me, 2, so);
                s.join("127.0.0.1", ports);
                std::vector<std::uint8_t> out(1000, static_cast<std::uint8_t>(me + 1));
                std::vector<std::uint8_t> in(1000);
                std::atomic<int> left{2};
                s.peer(1 - me).async_write(2, out.data(), out.size(), [&](auto) { --left; });
                s.peer(1 - me).async_read(2, in.data(), in.size(), [&](auto) { --left; });
                pump(io, left);
                if (in[0] == static_cast<std::uint8_t>(2 - me)
                    && s.edge_stats(1 - me).payload_out == 1000u)
                    ++ok;
            }
            catch (const std::exception & e)
            {
                std::lock_guard<std::mutex> lock(mu);
                err = e.what();
            }
        };
        std::thread t0([&] { run(0); });
        std::thread t1([&] { run(1); });
        t0.join();
        t1.join();
        EXPECT_EQ(ok.load(), 2) << dpf::net::transport_name(kind) << ": " << err;
    }
}

TEST(Session, ReconnectResumesMidPlan)
{
    const int steps = 5;
    auto c0 = make_chain(0, steps);
    auto c1 = make_chain(1, steps);
    auto ports = dpf::net::make_mesh_ports(2);
    std::uint64_t result[2] = {0, 0};
    std::uint64_t resumes[2] = {0, 0};
    std::string err;
    std::mutex mu;
    auto run = [&](unsigned me, const chain & ch) {
        try
        {
            asio::io_context io;
            dpf::net::party_session s(io, me, 2, std::size_t{steps});
            s.join("127.0.0.1", ports);
            int calls = 0;
            std::map<std::uint32_t, dpf::protocol::kernel_fn> k{
                {k_inc,
                    [&, base = inc_kernel()](std::uint32_t op,
                        const std::vector<dpf::protocol::node> & ns,
                        const std::vector<dpf::protocol::block_span> & in,
                        dpf::protocol::block_span out, std::size_t lanes) {
                        if (me == 0 && ++calls == 2)
                            s.peer(1).close();
                        base(op, ns, in, out, lanes);
                    }}};
            dpf::net::sink_options so;
            so.reconnect = s.reconnector(1 - me);
            dpf::net::async_round_sink sink(s.peer(1 - me), ch.plan.slot_bytes_all(),
                1, so);
            std::vector<std::vector<std::uint8_t>> v(ch.plan.nodes().size());
            put(v, ch.x, me + 1);
            dpf::protocol::drive_via_schedule(ch.plan, sink, v, k, me);
            result[me] = get(v, ch.out);
            resumes[me] = sink.stats().resumes;
        }
        catch (const std::exception & e)
        {
            std::lock_guard<std::mutex> lock(mu);
            err += std::string(e.what()) + "; ";
        }
    };
    std::thread t0([&] { run(0, c0); });
    std::thread t1([&] { run(1, c1); });
    t0.join();
    t1.join();
    ASSERT_TRUE(err.empty()) << err;
    EXPECT_EQ(result[0], chain_value(1, 2, steps));
    EXPECT_EQ(result[1], chain_value(1, 2, steps));
    EXPECT_EQ(resumes[0], 1u);
    EXPECT_EQ(resumes[1], 1u);
}

TEST(Session, DealerLinkHasItsOwnEpoch)
{
    asio::io_context dio, pio;
    std::atomic<unsigned short> port{0};
    std::string err;
    std::mutex mu;
    std::thread dealer([&] {
        try
        {
            dpf::net::dealer_session d(dio, 1);
            port.store(d.listen());
            d.accept_parties();
            std::uint64_t v = 5;
            std::atomic<int> left{1};
            d.party(0).async_write(0, &v, 8, [&](auto) { --left; });
            pump(dio, left);
            d.reconnect(0);
            v = 6;
            left = 1;
            d.party(0).async_write(0, &v, 8, [&](auto) { --left; });
            pump(dio, left);
        }
        catch (const std::exception & e)
        {
            std::lock_guard<std::mutex> lock(mu);
            err += e.what();
        }
    });
    while (port.load() == 0)
        std::this_thread::yield();
    dpf::net::party_session p(pio, 0, 2, std::size_t{1});
    p.connect_dealer("127.0.0.1", port.load());
    std::uint64_t got = 0;
    std::atomic<int> left{1};
    p.dealer().async_read(0, &got, 8, [&](auto) { --left; });
    pump(pio, left);
    EXPECT_EQ(got, 5u);
    p.reconnect_dealer();
    left = 1;
    p.dealer().async_read(0, &got, 8, [&](auto) { --left; });
    pump(pio, left);
    dealer.join();
    EXPECT_TRUE(err.empty()) << err;
    EXPECT_EQ(got, 6u);
}

// ---------------------------------------------------------------------------
// N-party runner and harness
// ---------------------------------------------------------------------------

TEST(Runner, TwoPartyOnEveryTransport)
{
    std::vector<dpf::net::transport> kinds = {dpf::net::transport::async_memory,
        dpf::net::transport::local, dpf::net::transport::mux,
        dpf::net::transport::parallel};
    if (dpf::net::sctp_available())
        kinds.push_back(dpf::net::transport::sctp);
    const int steps = 4;
    for (auto kind : kinds)
    {
        auto c0 = make_chain(0, steps);
        auto c1 = make_chain(1, steps);
        dpf::app::party_values v0(c0.plan.nodes().size()), v1(c1.plan.nodes().size());
        put(v0, c0.x, 3);
        put(v1, c1.x, 4);
        dpf::app::run_config cfg;
        cfg.kind = kind;
        cfg.n_lanes = 2;
        cfg.security.encrypt = kind != dpf::net::transport::sctp;
        std::map<std::uint32_t, dpf::protocol::kernel_fn> k{{k_inc, inc_kernel()}};
        const auto r = dpf::run_two_party(c0.plan, c1.plan, v0, v1, k, cfg);
        EXPECT_EQ(get(v0, c0.out), chain_value(3, 4, steps))
            << dpf::net::transport_name(kind);
        EXPECT_EQ(get(v1, c1.out), chain_value(3, 4, steps))
            << dpf::net::transport_name(kind);
        EXPECT_GT(r.wire[0].payload_out, 0u) << dpf::net::transport_name(kind);
        if (kind == dpf::net::transport::async_memory || kind == dpf::net::transport::sctp)
            EXPECT_EQ(r.wire[0].bytes_out, r.wire[0].payload_out);
        else
            EXPECT_GT(r.wire[0].bytes_out, r.wire[0].payload_out)
                << dpf::net::transport_name(kind);
    }
}

TEST(Runner, InstancesAndComputePool)
{
    const int steps = 3;
    auto c0 = make_chain(0, steps);
    auto c1 = make_chain(1, steps);
    dpf::app::party_values v0(c0.plan.nodes().size()), v1(c1.plan.nodes().size());
    v0[c0.x.id].assign(8 * 4, 0);
    v1[c1.x.id].assign(8 * 4, 0);
    for (std::size_t l = 0; l < 4; ++l)
    {
        std::uint64_t a = l, b = 10;
        std::memcpy(v0[c0.x.id].data() + 8 * l, &a, 8);
        std::memcpy(v1[c1.x.id].data() + 8 * l, &b, 8);
    }
    dpf::app::run_config cfg;
    cfg.instances = 4;
    cfg.compute_threads = 2;
    cfg.framing = dpf::net::framing_mode::always;
    std::map<std::uint32_t, dpf::protocol::kernel_fn> k{{k_inc, inc_kernel()}};
    (void)dpf::run_two_party(c0.plan, c1.plan, v0, v1, k, cfg);
    for (std::size_t l = 0; l < 4; ++l)
        EXPECT_EQ(get(v0, c0.out, l), chain_value(l, 10, steps));
}

TEST(Runner, StarOnEveryTransport)
{
    std::vector<dpf::net::transport> kinds = {dpf::net::transport::async_memory,
        dpf::net::transport::local, dpf::net::transport::mux,
        dpf::net::transport::parallel};
    if (dpf::net::sctp_available())
        kinds.push_back(dpf::net::transport::sctp);
    // pirsona_bitmore_fetch(L) is a star over 2^L servers.
    constexpr std::size_t L = 2;
    constexpr std::size_t n = std::size_t{1} << L;
    constexpr std::size_t query = 16 * L;
    for (auto kind : kinds)
    {
        auto seeds = std::make_shared<std::vector<std::vector<std::uint8_t>>>(n);
        auto answers = std::make_shared<std::vector<std::vector<std::uint8_t>>>(n);
        for (std::size_t i = 0; i < n; ++i)
        {
            (*seeds)[i].assign(query, static_cast<std::uint8_t>(i + 1));
            (*answers)[i].assign(8, static_cast<std::uint8_t>(0x40 + i));
        }
        auto client = dpf::protocol::pirsona_bitmore_fetch(L, 16, 8, seeds, answers);
        dpf::app::run_config cfg;
        cfg.kind = kind;
        cfg.n_lanes = 1;
        cfg.wait_timeout = std::chrono::seconds(10);
        cfg.security.encrypt = kind != dpf::net::transport::sctp;
        EXPECT_NO_THROW(dpf::session::drive_async_star(n, {query, 8u},
            std::move(client),
            [&](std::size_t i) {
                return dpf::protocol::star_server_reply_rounds(query, 8, (*answers)[i]);
            },
            cfg))
            << dpf::net::transport_name(kind);
    }
}

TEST(Runner, ShortInputIsRejectedNotZeroed)
{
    auto c0 = make_chain(0, 1);
    auto c1 = make_chain(1, 1);
    dpf::app::party_values v0(c0.plan.nodes().size()), v1(c1.plan.nodes().size());
    put(v0, c0.x, 1);
    put(v1, c1.x, 2);
    dpf::app::run_config cfg;
    cfg.instances = 4;
    std::map<std::uint32_t, dpf::protocol::kernel_fn> k{{k_inc, inc_kernel()}};
    try
    {
        (void)dpf::run_two_party(c0.plan, c1.plan, v0, v1, k, cfg);
        FAIL() << "8-byte input for 4 instances should be rejected";
    }
    catch (const std::exception & e)
    {
        EXPECT_NE(std::string(e.what()).find("input node"), std::string::npos)
            << e.what();
    }
}

TEST(Runner, ThreePartyDealerRingAndPeer)
{
    struct built
    {
        dpf::protocol::plan plan;
        dpf::protocol::node y, ring, pad, mask, x, open;
    };
    auto make = [](std::size_t party) {
        dpf::protocol::composer c(party);
        built b;
        b.y = c.input(domain::y, 8);
        b.ring = c.rss_from_y(b.y);
        b.pad = c.input(domain::a, 8);
        b.mask = c.dealer_deliver(b.pad);
        b.x = c.input(domain::a, 8);
        b.open = c.exchange(b.x);
        b.plan = c.schedule();
        return b;
    };
    for (auto kind : {dpf::net::transport::async_memory, dpf::net::transport::mux})
    {
        built b[3] = {make(0), make(1), make(2)};
        std::vector<dpf::app::party_values> values(3);
        for (int i = 0; i < 3; ++i)
        {
            values[i].resize(b[i].plan.nodes().size());
            put(values[i], b[i].y, 100 + i);
            put(values[i], b[i].pad, i == 2 ? 77 : 0);
            put(values[i], b[i].x, i == 0 ? 20 : 22);
        }
        dpf::app::run_config cfg;
        cfg.kind = kind;
        (void)dpf::app::run_parties({b[0].plan, b[1].plan, b[2].plan}, values, {},
            cfg);
        for (int i = 0; i < 3; ++i)
        {
            std::uint64_t own = 0, next = 0;
            std::memcpy(&own, values[i][b[i].ring.id].data(), 8);
            std::memcpy(&next, values[i][b[i].ring.id].data() + 8, 8);
            EXPECT_EQ(own, 100u + i) << dpf::net::transport_name(kind);
            EXPECT_EQ(next, 100u + (i + 1) % 3) << dpf::net::transport_name(kind);
        }
        for (int i = 0; i < 2; ++i)
        {
            EXPECT_EQ(get(values[i], b[i].mask), 77u);
            EXPECT_EQ(get(values[i], b[i].open), 42u);
        }
    }
}

TEST(Runner, SeparateProcessesFromAStaticTable)
{
    const std::vector<std::string> peers = {
        "127.0.0.1:" + std::to_string(free_port()),
        "127.0.0.1:" + std::to_string(free_port())};
    const std::string list = peers[0] + "," + peers[1];
    const int steps = 3;
    pid_t kids[2];
    for (int me = 0; me < 2; ++me)
    {
        const pid_t pid = ::fork();
        ASSERT_GE(pid, 0);
        if (pid == 0)
        {
            int code = 1;
            try
            {
                const std::string party = "--party=" + std::to_string(me);
                const std::string pl = "--peers=" + list;
                const char * argv[] = {"node", party.c_str(), pl.c_str(),
                    "--transport=mux", "--lanes=2"};
                const auto args = dpf::app::parse_node_args(5,
                    const_cast<char **>(argv), dpf::app::run_config{});
                auto ch = make_chain(static_cast<std::size_t>(me), steps);
                dpf::app::party_values v(ch.plan.nodes().size());
                put(v, ch.x, me == 0 ? 5 : 6);
                std::map<std::uint32_t, dpf::protocol::kernel_fn> k{
                    {k_inc, inc_kernel()}};
                (void)dpf::app::run_node(args, ch.plan, v, k);
                code = get(v, ch.out) == chain_value(5, 6, steps) ? 0 : 2;
            }
            catch (const std::exception & e)
            {
                std::fprintf(stderr, "party %d: %s\n", me, e.what());
                code = 3;
            }
            ::_exit(code);
        }
        kids[me] = pid;
    }
    for (int me = 0; me < 2; ++me)
    {
        int status = 0;
        ASSERT_EQ(::waitpid(kids[me], &status, 0), kids[me]);
        ASSERT_TRUE(WIFEXITED(status));
        EXPECT_EQ(WEXITSTATUS(status), 0) << "party " << me;
    }
}

TEST(Harness, TrialsConfigAndWireRecorded)
{
    auto c0 = make_chain(0, 2);
    auto c1 = make_chain(1, 2);
    std::vector<dpf::app::party_values> inputs(2);
    inputs[0].resize(c0.plan.nodes().size());
    inputs[1].resize(c1.plan.nodes().size());
    put(inputs[0], c0.x, 1);
    put(inputs[1], c1.x, 2);
    std::map<std::uint32_t, dpf::protocol::kernel_fn> k{{k_inc, inc_kernel()}};
    dpf::app::run_config cfg;
    cfg.kind = dpf::net::transport::mux;
    cfg.warmup = 1;
    cfg.trials = 3;
    dpf::experiment ex("chain", "p0");
    const auto cost =
        dpf::app::exercise_parties({c0.plan, c1.plan}, inputs, k, &ex, cfg);
    EXPECT_EQ(ex.trials().size(), 3u);
    EXPECT_GT(cost.wall_ns, 0u);
    EXPECT_GT(cost.wire_out, cost.bytes);
    EXPECT_EQ(ex.wire().bytes_out, cost.wire_out);
    bool has_transport = false;
    for (const auto & kv : ex.config())
        has_transport = has_transport || (kv.first == "transport" && kv.second == "mux");
    EXPECT_TRUE(has_transport);
    const std::string dir = "/tmp/libdpf_net_control_" + std::to_string(::getpid());
    ex.write_csv(dir);
    for (const char * f : {"/config.csv", "/trials.csv", "/wire.csv"})
    {
        std::ifstream in(dir + f);
        EXPECT_TRUE(static_cast<bool>(in)) << f;
    }
    (void)::system(("rm -rf '" + dir + "'").c_str());
}

TEST(Prep, ShippedOverADealerSession)
{
    dpf::prep::demand d;
    d.ring_triples = 2;
    dpf::app::run_config cfg;
    cfg.kind = dpf::net::transport::mux;
    auto shipped = dpf::session::ship_prep(d, cfg);
    std::uint8_t a0[8], b0[8], c0[8], a1[8], b1[8], c1[8];
    shipped.party0.take_ring(a0, b0, c0);
    shipped.party1.take_ring(a1, b1, c1);
    std::uint64_t ta[2], tb[2], tc[2];
    std::memcpy(&ta[0], a0, 8);
    std::memcpy(&tb[0], b0, 8);
    std::memcpy(&tc[0], c0, 8);
    std::memcpy(&ta[1], a1, 8);
    std::memcpy(&tb[1], b1, 8);
    std::memcpy(&tc[1], c1, 8);
    EXPECT_EQ((ta[0] + ta[1]) * (tb[0] + tb[1]), tc[0] + tc[1]);
}

TEST(Sink, WindowGatesPipelining)
{
    asio::io_context io0, io1;
    auto ends = dpf::net::make_async_dual_memory_stream_pair(io0, io1, 1, 64);
    dpf::net::sink_options so;
    so.drain_timeout = std::chrono::milliseconds(300);
    dpf::net::async_round_sink a(ends.first, {256}, 1, so);
    std::vector<std::uint8_t> slot(256, 1);
    a.submit(0, 0, slot.data(), slot.size());
    EXPECT_THROW(a.flush_round(0), std::runtime_error);
    EXPECT_FALSE(a.can_send_ahead());
}