/// @file dpf/net/party_tape_io.hpp
/// @brief Deal / accept a `dpf::beavers::party_tape` over a framed channel.
#ifndef LIBDPF_INCLUDE_DPF_NET_PARTY_TAPE_IO_HPP__
#define LIBDPF_INCLUDE_DPF_NET_PARTY_TAPE_IO_HPP__

#include <cstdint>
#include <stdexcept>
#include <type_traits>
#include <vector>

#include "dpf/beaver.hpp"
#include "dpf/net/channel.hpp"
#include "dpf/net/trio.hpp"

namespace dpf
{
namespace net
{
namespace detail
{

template <typename Ring>
void send_ring_vec(channel & c, const std::vector<Ring> & v)
{
    static_assert(std::is_trivially_copyable_v<Ring>,
        "party_tape ring must be trivially copyable");
    std::uint64_t n = v.size();
    c.send(msg::ring_vector, n);
    if (n != 0)
        c.send_vec(v.data(), v.size(), msg::ring_vector);
}

template <typename Ring>
std::vector<Ring> recv_ring_vec(channel & c)
{
    auto n = c.recv<std::uint64_t>(msg::ring_vector);
    if (n == 0)
        return {};
    return c.recv_vec<Ring>(msg::ring_vector);
}

inline void send_flags(channel & c, const std::vector<std::uint8_t> & v)
{
    std::uint64_t n = v.size();
    c.send(msg::bytes, n);
    if (n != 0)
        c.send_bytes(msg::bytes, v);
}

inline std::vector<std::uint8_t> recv_flags(channel & c)
{
    auto n = c.recv<std::uint64_t>(msg::bytes);
    if (n == 0)
        return {};
    auto body = c.recv_bytes(msg::bytes);
    if (body.size() != n)
        throw std::runtime_error("party_tape flag size mismatch");
    return body;
}

}  // namespace detail

template <typename Ring>
void send_party_tape(channel & c, const beavers::party_tape<Ring> & tape)
{
    c.send(msg::beaver_tape, std::uint8_t{2});
    c.send(msg::beaver_tape, static_cast<std::uint8_t>(tape.has_mac ? 1 : 0));
    detail::send_ring_vec(c, tape.lambda);
    detail::send_flags(c, tape.lambda_ready);
    detail::send_ring_vec(c, tape.monomial);
    detail::send_flags(c, tape.monomial_ready);
    detail::send_ring_vec(c, tape.bundles);
    detail::send_flags(c, tape.bundles_ready);
    detail::send_ring_vec(c, tape.dot_cross);
    detail::send_flags(c, tape.dot_ready);
    if (tape.has_mac)
    {
        detail::send_ring_vec(c, tape.lambda_tag);
        detail::send_ring_vec(c, tape.monomial_tag);
        detail::send_ring_vec(c, tape.bundles_tag);
        detail::send_ring_vec(c, tape.dot_cross_tag);
    }
}

template <typename Ring>
beavers::party_tape<Ring> recv_party_tape(channel & c)
{
    const auto ver = c.recv<std::uint8_t>(msg::beaver_tape);
    beavers::party_tape<Ring> tape;
    if (ver >= 2)
        tape.has_mac = c.recv<std::uint8_t>(msg::beaver_tape) != 0;
    tape.lambda = detail::recv_ring_vec<Ring>(c);
    tape.lambda_ready = detail::recv_flags(c);
    tape.monomial = detail::recv_ring_vec<Ring>(c);
    tape.monomial_ready = detail::recv_flags(c);
    tape.bundles = detail::recv_ring_vec<Ring>(c);
    tape.bundles_ready = detail::recv_flags(c);
    tape.dot_cross = detail::recv_ring_vec<Ring>(c);
    tape.dot_ready = detail::recv_flags(c);
    if (tape.has_mac)
    {
        tape.lambda_tag = detail::recv_ring_vec<Ring>(c);
        tape.monomial_tag = detail::recv_ring_vec<Ring>(c);
        tape.bundles_tag = detail::recv_ring_vec<Ring>(c);
        tape.dot_cross_tag = detail::recv_ring_vec<Ring>(c);
    }
    return tape;
}

/// @brief Dealer exports and sends each party's tape.
/// @tparam Ring Beaver ring
/// @param net the connected trio
/// @param s the sampled session
/// @throws std::logic_error if this process is not p2
template <typename Ring>
void send_party_tape(trio & net, role peer, const beavers::party_tape<Ring> & tape)
{
    net.send_to(peer, msg::beaver_tape, std::uint8_t{2});
    net.send_to(peer, msg::beaver_tape,
        static_cast<std::uint8_t>(tape.has_mac ? 1 : 0));
    auto send_ring_vec = [&](const std::vector<Ring> & v) {
        const std::uint64_t n = v.size();
        net.send_to(peer, msg::ring_vector, n);
        if (n != 0)
            net.send_vec_to(peer, msg::ring_vector, v);
    };
    auto send_flags = [&](const std::vector<std::uint8_t> & v) {
        const std::uint64_t n = v.size();
        net.send_to(peer, msg::bytes, n);
        if (n != 0)
            net.send_bytes_to(peer, msg::bytes, v.data(), v.size());
    };
    send_ring_vec(tape.lambda);
    send_flags(tape.lambda_ready);
    send_ring_vec(tape.monomial);
    send_flags(tape.monomial_ready);
    send_ring_vec(tape.bundles);
    send_flags(tape.bundles_ready);
    send_ring_vec(tape.dot_cross);
    send_flags(tape.dot_ready);
    if (tape.has_mac)
    {
        send_ring_vec(tape.lambda_tag);
        send_ring_vec(tape.monomial_tag);
        send_ring_vec(tape.bundles_tag);
        send_ring_vec(tape.dot_cross_tag);
    }
}

template <typename Ring>
beavers::party_tape<Ring> recv_party_tape(trio & net, role peer)
{
    const auto ver = net.recv_from<std::uint8_t>(peer, msg::beaver_tape);
    beavers::party_tape<Ring> tape;
    if (ver >= 2)
        tape.has_mac = net.recv_from<std::uint8_t>(peer, msg::beaver_tape) != 0;
    auto recv_ring_vec = [&]() {
        const auto n = net.recv_from<std::uint64_t>(peer, msg::ring_vector);
        if (n == 0)
            return std::vector<Ring>{};
        return net.recv_vec_from<Ring>(peer, msg::ring_vector);
    };
    auto recv_flags = [&]() {
        const auto n = net.recv_from<std::uint64_t>(peer, msg::bytes);
        if (n == 0)
            return std::vector<std::uint8_t>{};
        auto body = net.recv_bytes_from(peer, msg::bytes);
        if (body.size() != n)
            throw std::runtime_error("party_tape flag size mismatch");
        return body;
    };
    tape.lambda = recv_ring_vec();
    tape.lambda_ready = recv_flags();
    tape.monomial = recv_ring_vec();
    tape.monomial_ready = recv_flags();
    tape.bundles = recv_ring_vec();
    tape.bundles_ready = recv_flags();
    tape.dot_cross = recv_ring_vec();
    tape.dot_ready = recv_flags();
    if (tape.has_mac)
    {
        tape.lambda_tag = recv_ring_vec();
        tape.monomial_tag = recv_ring_vec();
        tape.bundles_tag = recv_ring_vec();
        tape.dot_cross_tag = recv_ring_vec();
    }
    return tape;
}

template <typename Ring>
void deal_session(trio & net, const beavers::session<Ring> & s)
{
    if (net.self() != role::p2)
        throw std::logic_error("deal_session is for the dealer");
    send_party_tape(net, role::p0, s.export_party(0));
    send_party_tape(net, role::p1, s.export_party(1));
}

/// @brief Computing party receives its tape from the dealer.
/// @tparam Ring Beaver ring
/// @param net the connected trio
/// @return this party's tape
/// @throws std::logic_error if this process is p2
/// @throws std::runtime_error if a flag vector's length disagrees with its header
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
beavers::party_tape<Ring> accept_session(trio & net)
{
    if (net.self() == role::p2)
        throw std::logic_error("dealer does not accept_session");
    return recv_party_tape<Ring>(net, role::p2);
}

}  // namespace net
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_NET_PARTY_TAPE_IO_HPP__