#include <cstddef>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <thread>
#include <vector>

#include "dpf.hpp"
#include "dpf/net/stream_array.hpp"
#include "dpf/protocol_factory.hpp"

// Protocol composition schedules (compose.hpp). Records FSS walks, ABY
// products, and RSS refreshes on one RoundSink plan — the shapes Express,
// Sabre, Pika early-stop, Poplar prefixes, and Duoram scale use by hand.
//
// The final block drives a composed open on the stream framework
// (drive_both_on_streams). The identical schedule also runs over the truly
// asynchronous backends (dpf::net::async_stream_array + async_round_sink, or
// dpf::async::overlapped_byte_protocol) and, on Linux, over real SCTP
// (async_sctp_stream_array). The experiment_bench harness picks the transport
// via DPF_TRANSPORT=memory|stream|async|mux|parallel|sctp.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty
//       examples/protocol/compose_schedule.cpp

namespace
{

using dpf::protocol::composer;
using dpf::protocol::domain;
using dpf::protocol::effect;
namespace opcodes = dpf::protocol::opcodes;

}  // namespace

int main()
{
    // Express / Sabre: audit rides in the last CW flush (depth exchanges, not
    // depth+1).
    {
        composer naive(0);
        auto seed = naive.input(domain::fss, 16);
        auto sketch = naive.input(domain::a, 8);
        auto leaf = naive.fss_point(seed, 4, 16);
        auto dep = naive.compute(opcodes::user_base + 1, {leaf, sketch},
            domain::a, 8);
        (void)naive.exchange(dep);
        const auto naive_ex = naive.schedule().exchange_waves();

        composer fused(0);
        auto seed_f = fused.input(domain::fss, 16);
        auto sketch_f = fused.input(domain::a, 8);
        auto wr = fused.fss_point_fused(seed_f, 4, 16, sketch_f);
        const auto fused_ex = fused.schedule().exchange_waves();
        if (fused_ex >= naive_ex || fused.domain_of(wr.trailer_open) != domain::a)
        {
            std::cerr << "compose fuse\n";
            return 1;
        }
        std::cout << "express_fuse " << naive_ex << "->" << fused_ex << "\n";
    }

    // Pika / small-output PIR: BGI Remark 3.4 early-stop drops ν CW rounds.
    {
        composer full(0);
        auto seed = full.input(domain::fss, 16);
        (void)full.fss_point(seed, 8, 16);
        composer early(0);
        auto seed_e = early.input(domain::fss, 16);
        (void)early.fss_point_early_stop(seed_e, 8, /*early_stop=*/3, 16);
        const auto a = full.schedule().exchange_waves();
        const auto b = early.schedule().exchange_waves();
        if (b != a - 3)
        {
            std::cerr << "compose early_stop\n";
            return 1;
        }
        std::cout << "early_stop " << a << "->" << b << "\n";
    }

    // Poplar: prefix share after each CW, same exchange-wave depth.
    {
        composer c(0);
        auto seed = c.input(domain::fss, 16);
        auto wr = c.level_walk_prefixes(seed, 5, 16, /*prefix_bytes=*/8);
        auto p = c.schedule();
        if (wr.at_level.size() != 5 || p.exchange_waves() != 5)
        {
            std::cerr << "compose prefixes\n";
            return 1;
        }
        std::cout << "prefixes " << p.exchange_waves() << "\n";
    }

    // DCF block_width: variable CW slot sizes.
    {
        composer c(0);
        auto seed = c.input(domain::fss, 16);
        const std::vector<std::size_t> slots = {16, 4, 4, 16};
        (void)c.level_walk_sized(seed, slots);
        auto p = c.schedule();
        if (p.value_bytes_of(p.wave(1).exchanges[0].id) != 4)
        {
            std::cerr << "compose sized\n";
            return 1;
        }
        std::cout << "sized_slots ok\n";
    }

    // RSS product → neighbor y-exchange → RSS (one round).
    {
        composer c(0);
        auto x = c.input(domain::rss, 16);
        auto y = c.input(domain::rss, 16);
        auto z = c.rss_product_replicated(x, y);
        if (c.domain_of(z) != domain::rss || c.schedule().exchange_waves() != 1)
        {
            std::cerr << "compose rss\n";
            return 1;
        }
        std::cout << "rss_refresh 1\n";
    }

    // Duoram write-scale: FSS leaf feeds ABY; beaver waits for the leaf wave.
    {
        composer c(0);
        auto seed = c.input(domain::fss, 16);
        auto leaf = c.fss_point(seed, 4, 16);
        auto scale = c.input(domain::a, 8);
        auto scaled = c.aby_product<std::uint64_t>(leaf, scale);
        auto p = c.schedule();
        if (p.wave_of(scaled) < p.wave_of(leaf))
        {
            std::cerr << "compose duoram_scale\n";
            return 1;
        }
        std::cout << "duoram_scale wave " << p.wave_of(scaled) << "\n";
    }

    // Round-aware ABY: sign×linear stays one online round.
    {
        composer c(0);
        auto & s = c.aby<std::uint64_t>();
        auto sgn = s.input();
        auto x = s.input();
        auto a0 = s.input();
        auto a1 = s.input();
        auto lin = s(sgn * (a1 * x + a0));
        if (s.round_of(lin) != 1)
        {
            std::cerr << "compose aby_rounds\n";
            return 1;
        }
        std::cout << "aby_rounds 1\n";
    }

    // Doerner–Shelat: 5 opens per level; OH adds 80 AND-layers / level.
    {
        composer c(0);
        auto seed = c.input(domain::fss, 16);
        auto tip = c.level_walk_ds(seed, 2, 16);
        if (c.schedule().exchange_waves() != 10 || tip.id == seed.id)
        {
            std::cerr << "compose ds_walk\n";
            return 1;
        }
        std::cout << "ds_walk 10\n";
        composer c_oh(0);
        auto tip_oh = c_oh.level_walk_ds(c_oh.input(domain::fss, 16), 1, 16,
            /*oh=*/true);
        (void)tip_oh;
        const auto want = dpf::net::compose_ds_slot_bytes(1, 16, true).size();
        if (c_oh.schedule().exchange_waves() != want)
        {
            std::cerr << "compose ds_oh\n";
            return 1;
        }
        std::cout << "ds_oh " << want << "\n";
    }

    // Adaptive idpf: one packed L‖R open per step.
    {
        composer c(0);
        auto seed = c.input(domain::fss, 16);
        auto f = c.begin_adaptive_prefix(seed);
        f = c.step_adaptive_prefix(f, 16, 8);
        f = c.retain_adaptive_prefix(f, 0);
        if (c.schedule().exchange_waves() != 1)
        {
            std::cerr << "compose adaptive\n";
            return 1;
        }
        std::cout << "adaptive 1\n";
    }

    // default_plan: rounds == exchange_waves (sink-aligned).
    {
        composer c(0);
        auto leaf = c.fss_point(c.input(domain::fss, 16), 4, 16);
        (void)leaf;
        auto p = c.default_plan();
        if (p.rounds() != p.exchange_waves()
            || p.rounds() != p.slot_bytes_all().size() || p.rounds() != 4)
        {
            std::cerr << "compose default_plan\n";
            return 1;
        }
        std::cout << "default_plan 4\n";
    }

    // Multipoint buckets: CW waves pack; answers one open.
    {
        composer c(0);
        auto mr = c.multipoint_fan(2,
            [&](std::size_t) { return c.input(domain::fss, 16); }, 3, 16, 8);
        (void)mr;
        if (c.schedule().exchange_waves() != 4)
        {
            std::cerr << "compose multipoint\n";
            return 1;
        }
        std::cout << "multipoint 4\n";
    }

    // Multi-lane ABY: independent barriers, one wave.
    {
        composer c(0);
        auto z0 = c.aby_product<std::uint64_t>(c.input(domain::a, 8),
            c.input(domain::a, 8), 0);
        auto z1 = c.aby_product<std::uint64_t>(c.input(domain::a, 8),
            c.input(domain::a, 8), 1);
        if (c.schedule().exchange_waves() != 1
            || c.schedule().effect_count(effect::exchange) != 2
            || z0.id == z1.id)
        {
            std::cerr << "compose multilane\n";
            return 1;
        }
        std::cout << "multilane 1\n";
    }

    // Prepaid defer + rotate: zero online FSS rounds.
    {
        composer c(0);
        auto buf = c.defer_expand(c.input(domain::fss, 16), 4, 16);
        auto rot = c.rotate_share(buf, 7);
        if (c.schedule().exchange_waves() != 0 || rot.id == buf.id)
        {
            std::cerr << "compose defer\n";
            return 1;
        }
        std::cout << "defer 0\n";
    }

    // stream_array: drive a compose open through drive_both_on_streams.
    {
        composer c0(0);
        composer c1(1);
        auto x0 = c0.input(domain::a, 8);
        auto x1 = c1.input(domain::a, 8);
        auto e0 = c0.exchange(x0);
        auto e1 = c1.exchange(x1);
        auto p0 = c0.schedule();
        auto p1 = c1.schedule();
        std::vector<std::vector<std::uint8_t>> v0(p0.nodes().size()),
            v1(p1.nodes().size());
        const std::uint64_t a = 3, b = 5;
        v0[x0.id].assign(8, 0);
        v1[x1.id].assign(8, 0);
        std::memcpy(v0[x0.id].data(), &a, 8);
        std::memcpy(v1[x1.id].data(), &b, 8);
        dpf::protocol::drive_both_on_streams(p0, p1, v0, v1);
        std::uint64_t open0 = 0, open1 = 0;
        std::memcpy(&open0, v0[e0.id].data(), 8);
        std::memcpy(&open1, v1[e1.id].data(), 8);
        if (open0 != 8 || open1 != 8)
        {
            std::cerr << "drive_plan_on_streams open\n";
            return 1;
        }
        std::cout << "drive_plan_on_streams_ok " << open0 << "\n";
    }

    return 0;
}