/// @file dpf/yao_share.hpp
/// @brief Move a DPF leaf share into the XOR bits a garbled netlist consumes.
/// @details A point leaf is subtractive (`b2y` / `y2b`). A comparison leaf is
///          additive (`a2y` / `y2a`). An `fss_share` opens like a point leaf
///          (`fss2y` / `y2fss`). A replicated leaf that parties 0 and 1 will
///          garble uses `rss2y`; `y2rss` deals a fresh triple of the bits'
///          integer value.
///
///          The Yao share of that integer is one XOR bit per party, least-
///          significant bit at index 0. Those bytes are what
///          `netlist::shared_in` takes and what `session::eval` returns.
///          The forward calls run edaBit A2B with a GMW carry. The reverse
///          calls run daBit B2A.
///
///          Both parties' shares are arguments, as in `edabit::a2b_gmw_pair`:
///          the opens are resolved inside the call. Top-level `dpf::rss2y` and
///          `dpf::y2rss` are the local (3,3) casts and are not these functions.
///
///          `rss2y` keeps parties 0 and 1, the garbler and the evaluator.
///          Party 0's additive share is `x0 + x1` (both components it holds).
///          Party 1's is `x2` (its `next`). Party 2 sends nothing.
///          `y2rss` is called where both Yao shares already sit, and deals a
///          fresh replicated triple of their sum.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_YAO_SHARE_HPP__
#define LIBDPF_INCLUDE_DPF_YAO_SHARE_HPP__

#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <tuple>
#include <utility>
#include <vector>

#include "hedley/hedley.h"

#include "dpf/edabit.hpp"
#include "dpf/ot_pack.hpp"
#include "dpf/random.hpp"
#include "dpf/secret_share.hpp"

namespace dpf
{
namespace yao
{

namespace detail
{

HEDLEY_ALWAYS_INLINE
unsigned width_or_full(unsigned width, unsigned full)
{
    if (width == 0)
        return full;
    if (width > full)
        throw std::invalid_argument("yao: bit width exceeds the ring");
    return width;
}

HEDLEY_ALWAYS_INLINE
HEDLEY_WARN_UNUSED_RESULT
std::vector<std::uint8_t> unpack_lsb(const std::vector<std::uint8_t> & packed,
    unsigned width)
{
    std::vector<std::uint8_t> out(width);
    for (unsigned i = 0; i < width; ++i)
        out[i] = edabit::detail::get_bit(packed, i);
    return out;
}

template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<std::vector<std::uint8_t>, std::vector<std::uint8_t>>
a2y_raw(Ring x0, Ring x1, unsigned width)
{
    width = width_or_full(width, static_cast<unsigned>(8u * sizeof(Ring)));
    auto eda = edabit::sample_edabit_pair<Ring>(width);
    auto packed = edabit::a2b_gmw_pair(eda, x0, x1);
    return {unpack_lsb(packed.first, width), unpack_lsb(packed.second, width)};
}

template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<Ring, Ring> y2a_raw(const std::vector<std::uint8_t> & y0,
    const std::vector<std::uint8_t> & y1, unsigned width)
{
    width = width_or_full(width, static_cast<unsigned>(8u * sizeof(Ring)));
    if (y0.size() != width || y1.size() != width)
        throw std::invalid_argument("yao: bit vector length");
    for (unsigned i = 0; i < width; ++i)
    {
        if (y0[i] > 1u || y1[i] > 1u)
            throw std::invalid_argument("yao: bit value");
    }
    Ring s0{};
    Ring s1{};
    for (unsigned i = 0; i < width; ++i)
    {
        auto dab = ot::sample_dabit_pair<Ring>();
        const std::uint8_t mask = static_cast<std::uint8_t>(
            (y0[i] ^ dab.p0.bit) ^ (y1[i] ^ dab.p1.bit));
        s0 = static_cast<Ring>(s0 + edabit::b2a_party_bit<Ring>(y0[i], dab.p0, mask, 0, i));
        s1 = static_cast<Ring>(s1 + edabit::b2a_party_bit<Ring>(y1[i], dab.p1, mask, 1, i));
    }
    return {s0, s1};
}

}  // namespace detail

/// @brief Additive shares to Yao bit shares. LSB at index 0.
/// @param x0 party 0's additive share
/// @param x1 party 1's additive share
/// @param width bit length. 0 means the whole ring
/// @return party 0's bits, then party 1's bits. XOR is the low `width` bits of `x0 + x1`
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<std::vector<std::uint8_t>, std::vector<std::uint8_t>>
a2y(const additive_share<Ring, 0> & x0, const additive_share<Ring, 1> & x1,
    unsigned width = 0)
{
    return detail::a2y_raw<Ring>(x0.raw(), x1.raw(), width);
}

/// @brief Point-leaf (subtractive) shares to Yao bit shares.
/// @param x0 party 0
/// @param x1 party 1
/// @param width bit length. 0 means the whole ring
/// @return XOR bits of `x0 - x1`
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<std::vector<std::uint8_t>, std::vector<std::uint8_t>>
b2y(const subtractive_share<Ring, 0> & x0, const subtractive_share<Ring, 1> & x1,
    unsigned width = 0)
{
    return a2y(x0.as_additive(), x1.as_additive(), width);
}

/// @brief `fss_share` to Yao bit shares. Same opening as a point leaf.
/// @param x0 party 0
/// @param x1 party 1
/// @param width bit length. 0 means the whole ring
/// @return XOR bits of the subtractive opening
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<std::vector<std::uint8_t>, std::vector<std::uint8_t>>
fss2y(const fss_share<Ring, 0> & x0, const fss_share<Ring, 1> & x1,
    unsigned width = 0)
{
    return a2y(x0.as_additive(), x1.as_additive(), width);
}

/// @brief Replicated shares to Yao bit shares for the garbler and the evaluator.
/// @details Party 0 contributes `own + next` (`x0 + x1`). Party 1 contributes
///          `next` (`x2`). `r0.next` must equal `r1.own`. Party 2 is idle.
/// @param r0 party 0
/// @param r1 party 1
/// @param width bit length. 0 means the whole ring
/// @return XOR bits of `x0 + x1 + x2` for parties 0 and 1
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<std::vector<std::uint8_t>, std::vector<std::uint8_t>>
rss2y(const replicated_share<Ring, 0> & r0, const replicated_share<Ring, 1> & r1,
    unsigned width = 0)
{
    if (r0.next != r1.own)
        throw std::invalid_argument("yao: rss components disagree on x1");
    const Ring a0 = static_cast<Ring>(r0.own + r0.next);
    const Ring a1 = r1.next;
    return detail::a2y_raw<Ring>(a0, a1, width);
}

/// @brief Yao bit shares back to additive shares.
/// @param y0 party 0's bits, LSB at index 0, one byte per bit
/// @param y1 party 1's bits
/// @param width bit length. 0 means `y0.size()`
/// @return additive shares of the integer those bits encode
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<additive_share<Ring, 0>, additive_share<Ring, 1>>
y2a(const std::vector<std::uint8_t> & y0, const std::vector<std::uint8_t> & y1,
    unsigned width = 0)
{
    if (width == 0)
        width = static_cast<unsigned>(y0.size());
    auto [s0, s1] = detail::y2a_raw<Ring>(y0, y1, width);
    return {additive_share<Ring, 0>::from_raw(s0),
        additive_share<Ring, 1>::from_raw(s1)};
}

/// @brief Yao bit shares back to a point-leaf (subtractive) sharing.
/// @param y0 party 0's bits
/// @param y1 party 1's bits
/// @param width bit length. 0 means `y0.size()`
/// @return subtractive shares of the integer those bits encode
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<subtractive_share<Ring, 0>, subtractive_share<Ring, 1>>
y2b(const std::vector<std::uint8_t> & y0, const std::vector<std::uint8_t> & y1,
    unsigned width = 0)
{
    auto [a0, a1] = y2a<Ring>(y0, y1, width);
    return {a0.as_subtractive(), a1.as_subtractive()};
}

/// @brief Yao bit shares back to `fss_share`.
/// @param y0 party 0's bits
/// @param y1 party 1's bits
/// @param width bit length. 0 means `y0.size()`
/// @return FSS shares of the integer those bits encode
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::pair<fss_share<Ring, 0>, fss_share<Ring, 1>>
y2fss(const std::vector<std::uint8_t> & y0, const std::vector<std::uint8_t> & y1,
    unsigned width = 0)
{
    auto [a0, a1] = y2a<Ring>(y0, y1, width);
    return {a0.as_fss(), a1.as_fss()};
}

/// @brief Yao bit shares to a fresh replicated triple of the same integer.
/// @details Both bit vectors are arguments, so their sum is already available
///          to the caller. The three components are a new sharing of that sum.
/// @param y0 party 0's bits
/// @param y1 party 1's bits
/// @param width bit length. 0 means `y0.size()`
/// @return replicated shares for parties 0, 1, and 2
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::tuple<replicated_share<Ring, 0>, replicated_share<Ring, 1>, replicated_share<Ring, 2>>
y2rss(const std::vector<std::uint8_t> & y0, const std::vector<std::uint8_t> & y1,
    unsigned width = 0)
{
    auto [a0, a1] = y2a<Ring>(y0, y1, width);
    const Ring secret = static_cast<Ring>(a0.raw() + a1.raw());
    const Ring x0 = dpf::uniform_sample<Ring>();
    const Ring x1 = dpf::uniform_sample<Ring>();
    const Ring x2 = static_cast<Ring>(secret - x0 - x1);
    return dpf::make_replicated_shares(x0, x1, x2);
}

}  // namespace yao
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_YAO_SHARE_HPP__