/// @file dpf/pad_graphs.hpp
/// @brief Setup / pad protocols as real `schedule_round` lists.
#ifndef LIBDPF_INCLUDE_DPF_PAD_GRAPHS_HPP__
#define LIBDPF_INCLUDE_DPF_PAD_GRAPHS_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <memory>
#include <stdexcept>
#include <vector>

#include "dpf/protocol.hpp"

namespace dpf
{
namespace protocol
{

/// @brief Dealer Beaver / Du-Atallah blinding tape: one delivery round per triple.
inline std::vector<schedule_round> dealer_tape_graph(std::size_t n_triples,
    std::size_t slot_bytes,
    const std::shared_ptr<std::vector<std::uint8_t>> & tape)
{
    return make_pad_rounds(n_triples, slot_bytes, tape, edge_channel::dealer);
}

/// @brief Du-Atallah 3PC multiply: P2 blinds (offline) + one P0↔P1 exchange.
/// @details Round 0 on dealer edge carries (X0,Y0,…) tokens; round 1 on peer
///          exchanges blinded shares. Bodies are schedule-shaped stubs that
///          XOR into `tape` so PIRsona / hushmap can splice them before online.
inline std::vector<schedule_round> du_atallah_mul_graph(
    const std::shared_ptr<std::vector<std::uint8_t>> & tape,
    std::size_t slot_bytes = 24)
{
    auto dealer = make_pad_rounds(1, slot_bytes, tape, edge_channel::dealer);
    auto peer = make_pad_rounds(1, slot_bytes, tape, edge_channel::peer);
    peer[0].sink_round = 0;
    return splice_rounds(std::move(dealer), std::move(peer));
}

/// @brief Key-ship: one round delivering `key_bytes` on `edge` (PIR upload seed).
inline std::vector<schedule_round> key_ship_graph(std::size_t key_bytes,
    edge_id edge = edge_peer)
{
    std::vector<schedule_round> rounds(1);
    rounds[0].slot_bytes = key_bytes;
    rounds[0].edge = edge;
    rounds[0].channel = edge <= edge_dealer
        ? static_cast<edge_channel>(edge)
        : edge_channel::peer;
    rounds[0].recv = receive_rule::copy_peer;
    rounds[0].sink_round = 0;
    rounds[0].produce = [key_bytes](std::size_t, const std::uint8_t *,
                            std::size_t, std::uint8_t * out) {
        if (out != nullptr && key_bytes != 0)
            std::memset(out, 0x5a, key_bytes);
    };
    return rounds;
}

/// @brief Server end of one star spoke: absorb upload, produce answer.
inline std::vector<schedule_round> star_server_reply_rounds(
    std::size_t query_bytes, std::size_t answer_bytes,
    std::vector<std::uint8_t> answer)
{
    std::vector<schedule_round> sr(2);
    sr[0].slot_bytes = query_bytes;
    sr[0].edge = edge_peer;
    sr[0].sink_round = 0;
    sr[0].produce = [query_bytes](std::size_t, const std::uint8_t *, std::size_t,
                         std::uint8_t * out) {
        if (out != nullptr && query_bytes != 0)
            std::memset(out, 0, query_bytes);
    };
    sr[1].slot_bytes = answer_bytes;
    sr[1].edge = edge_peer;
    sr[1].sink_round = 1;
    sr[1].produce = [answer = std::move(answer), answer_bytes](std::size_t,
                         const std::uint8_t *, std::size_t,
                         std::uint8_t * out) {
        if (out == nullptr)
            return;
        if (answer.size() >= answer_bytes)
            std::memcpy(out, answer.data(), answer_bytes);
        else
            std::memset(out, 0, answer_bytes);
    };
    return sr;
}

/// @brief Client star: upload then answer on each of `n_servers` edges.
/// @details Rounds `[0, n)` upload on edge i; rounds `[n, 2n)` answer on edge i.
inline std::vector<schedule_round> star_upload_answer_graph(
    std::size_t n_servers, std::size_t query_bytes, std::size_t answer_bytes,
    const std::shared_ptr<std::vector<std::vector<std::uint8_t>>> & queries,
    const std::shared_ptr<std::vector<std::vector<std::uint8_t>>> & answers)
{
    if (n_servers < 2)
        throw std::invalid_argument("star_upload_answer_graph servers");
    std::vector<schedule_round> rounds(2 * n_servers);
    for (std::size_t i = 0; i < n_servers; ++i)
    {
        auto & up = rounds[i];
        up.slot_bytes = query_bytes;
        up.edge = static_cast<edge_id>(i);
        up.recv = receive_rule::copy_peer;
        up.sink_round = 0;
        up.produce = [i, query_bytes, queries](std::size_t, const std::uint8_t *,
                         std::size_t, std::uint8_t * out) {
            if (out == nullptr)
                return;
            if (queries && i < queries->size()
                && (*queries)[i].size() >= query_bytes)
                std::memcpy(out, (*queries)[i].data(), query_bytes);
            else
                std::memset(out, static_cast<int>(i + 1), query_bytes);
        };
        auto & ans = rounds[n_servers + i];
        ans.slot_bytes = answer_bytes;
        ans.edge = static_cast<edge_id>(i);
        ans.recv = receive_rule::copy_peer;
        ans.sink_round = 1;
        ans.produce = [i, answer_bytes, answers](std::size_t,
                            const std::uint8_t * peer, std::size_t peer_n,
                            std::uint8_t * out) {
            if (out == nullptr)
                return;
            if (answers && i < answers->size()
                && (*answers)[i].size() >= answer_bytes)
                std::memcpy(out, (*answers)[i].data(), answer_bytes);
            else if (peer != nullptr && peer_n >= answer_bytes)
                std::memcpy(out, peer, answer_bytes);
            else
                std::memset(out, 0, answer_bytes);
        };
    }
    return rounds;
}

}  // namespace protocol
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PAD_GRAPHS_HPP__