/// @file dpf/net/sink_exchange.hpp
/// @brief One-at-a-time typed exchange on a RoundSink (count lanes).
#ifndef LIBDPF_INCLUDE_DPF_NET_SINK_EXCHANGE_HPP__
#define LIBDPF_INCLUDE_DPF_NET_SINK_EXCHANGE_HPP__

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

#include "dpf/net/round_sink.hpp"

namespace dpf
{
namespace net
{

/// @brief Advance through RoundSink rounds with trivially copyable slots.
class sink_exchange
{
public:
    sink_exchange(RoundSink & sink, std::size_t index = 0)
      : sink_(&sink), index_(index)
    { }

    std::uint16_t round() const noexcept { return round_; }
    std::uint16_t & round_ref() noexcept { return round_; }

    template <typename T>
    T operator()(const T & mine)
    {
        static_assert(std::is_trivially_copyable_v<T>,
            "sink_exchange requires a trivially copyable type");
        if (!sink_)
            throw std::logic_error("sink_exchange: no sink");
        if (round_ >= sink_->rounds())
            throw std::runtime_error("sink_exchange: out of rounds");
        const std::size_t slot = sink_->slot_bytes(round_);
        if (sizeof(T) > slot)
            throw std::runtime_error("sink_exchange: slot too small");
        std::vector<std::uint8_t> buf(slot, 0);
        if constexpr (sizeof(T) != 0)
            std::memcpy(buf.data(), &mine, sizeof(T));
        sink_->submit(round_, index_, buf.data(), buf.size());
        sink_->flush_round(round_);
        sink_->poll();
        if (!sink_->peer_ready(round_, index_))
            throw std::runtime_error("sink_exchange: peer missing");
        std::vector<std::uint8_t> peer_buf(slot);
        sink_->read_peer(round_, index_, peer_buf.data(), slot);
        T peer{};
        if constexpr (sizeof(T) != 0)
            std::memcpy(&peer, peer_buf.data(), sizeof(T));
        ++round_;
        return peer;
    }

private:
    RoundSink * sink_ = nullptr;
    std::size_t index_ = 0;
    std::uint16_t round_ = 0;
};

/// @brief Slot widths for a classic non-verifiable point walk of `depth` levels.
/// @details Per level: blind_msg, share_msg, advice_msg, and one 16-byte
///          bit-block product. Callers pad every slot to `max` so a single
///          sink sizing works for the whole walk.
inline std::vector<std::size_t> point_walk_slot_bytes(std::size_t depth,
    std::size_t max_slot)
{
    std::vector<std::size_t> out;
    out.reserve(depth * 4);
    for (std::size_t i = 0; i < depth * 4; ++i)
        out.push_back(max_slot);
    return out;
}

/// @brief Conservative RoundSink sizing for a Doerner–Shelat point or
///        comparison walk, including optional oblivious-hash AND layers.
/// @param depth tree depth (`dpf_type::depth`)
/// @param with_oblivious_hash verifiable non-reveal path (shared prefix hash)
/// @param lg_outputs `dpf_type::lg_outputs_per_leaf` — packed leaf mux cost

/// Boyar–Peralta SubBytes AND-depth layers recorded by compose / party OH walks.
inline constexpr std::size_t ds_oh_subbytes_per_level = 10;
inline constexpr std::size_t ds_oh_and_layers_per_subbytes = 8;
inline constexpr std::size_t ds_oh_exchanges_per_level =
    ds_oh_subbytes_per_level * ds_oh_and_layers_per_subbytes;
inline constexpr std::size_t ds_oh_slot_bytes = 4096;

/// @brief Slot widths matching a composed DS walk (4 opens/level + OH layers).
/// @details Prefer this when sizing a sink for `composer::level_walk_ds` /
///          `level_walk_ds_sized`. `ds_walk_slot_bytes` stays the conservative
///          party oversize used by hand `dist_ds` paths.
inline std::vector<std::size_t> compose_ds_slot_bytes(std::size_t depth,
    std::size_t fss_slot_bytes, bool with_oblivious_hash,
    std::size_t lg_outputs = 0)
{
    if (depth == 0)
        throw std::invalid_argument("compose_ds_slot_bytes depth must be > 0");
    if (fss_slot_bytes == 0)
        throw std::invalid_argument("compose_ds_slot_bytes slot must be > 0");
    std::vector<std::size_t> out;
    out.reserve(depth * (5 + (with_oblivious_hash ? ds_oh_exchanges_per_level : 0))
        + 64);
    for (std::size_t level = 0; level < depth; ++level)
    {
        // blind, share, advice, AND round 1, AND round 2 — same as point_party.
        for (int i = 0; i < 5; ++i)
            out.push_back(fss_slot_bytes);
        if (with_oblivious_hash)
        {
            for (std::size_t i = 0; i < ds_oh_exchanges_per_level; ++i)
                out.push_back(ds_oh_slot_bytes);
        }
    }
    if (lg_outputs > 0 && lg_outputs < 32)
    {
        const std::size_t mux_nodes =
            2 * ((std::size_t{1} << lg_outputs) - 1);
        for (std::size_t i = 0; i < mux_nodes; ++i)
            out.push_back(fss_slot_bytes);
    }
    return out;
}

inline std::vector<std::size_t> ds_walk_slot_bytes(std::size_t depth,
    bool with_oblivious_hash, std::size_t lg_outputs = 0)
{
    // Envelope of the hand walk (typed frames, leaf extras) and a composed
    // walk at 16-byte FSS slots. Every compose round fits in this sink;
    // extra rounds keep the hand walk from running out of slots.
    const auto composed = compose_ds_slot_bytes(depth, 16, with_oblivious_hash,
        lg_outputs);
    // Per level: blind, share, advice, two beaver_bit_block exchanges, plus
    // reveal / leaf extras. Oblivious hash: 10 SubBytes × ~8 BP AND-depth
    // layers (one schedule round each), not one round per AND gate.
    std::size_t n = depth * 10 + 16;
    // Packed leaf mux: two sink exchanges per internal node of a complete
    // binary tree with 2^lg candidates. Verifiable non-reveal wildcards run
    // a second mux for the payload vector — budget both.
    if (lg_outputs > 0 && lg_outputs < 32)
    {
        const std::size_t mux =
            2 * ((std::size_t{1} << lg_outputs) - 1) + 8;
        n += mux;
        if (with_oblivious_hash)
            n += mux;
    }
    std::size_t slot = 64;
    if (with_oblivious_hash)
    {
        n += depth * ds_oh_exchanges_per_level;
        // Packed mask bits for a full SubBytes AND layer across 8×16 bytes.
        slot = ds_oh_slot_bytes;
    }
    std::vector<std::size_t> out(std::max(n, composed.size()), slot);
    for (std::size_t i = 0; i < composed.size(); ++i)
        out[i] = std::max(out[i], composed[i]);
    return out;
}

}  // namespace net
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_NET_SINK_EXCHANGE_HPP__