/// @file party/oblivious_select.hpp
/// @brief Shared-output muxes for the socket Doerner–Shelat walk.
/// @details Included from inside `dpf::party::dist`. Products stay shared.
///          The only opened values are public correction words.

#ifndef LIBDPF_PARTY_OBLIVIOUS_SELECT_HPP__
#define LIBDPF_PARTY_OBLIVIOUS_SELECT_HPP__

template <int Me, typename T>
std::vector<T> exch_vec(net::trio & net, const std::vector<T> & mine)
{
    const role peer = Me == 0 ? role::p1 : role::p0;
    if constexpr (Me == 0)
    {
        net.to(peer).send_vec(mine, net::msg::delta);
        return net.to(peer).template recv_vec<T>(net::msg::delta);
    }
    else
    {
        auto got = net.to(peer).template recv_vec<T>(net::msg::delta);
        net.to(peer).send_vec(mine, net::msg::delta);
        return got;
    }
}

/// @brief This party's XOR share of `bit * block`. The product is not opened.
template <typename Exchange>
simde__m128i beaver_xor_share_bit(Exchange && exch, bool i_hold_block,
    std::uint8_t bit_share, simde__m128i block, const and_share_msg & mine)
{
    and_round1_msg msg{};
    msg.d_bit = static_cast<std::uint8_t>(bit_share ^ mine.a);
    msg.b_bit = mine.b;
    msg.a_m = mine.a;
    msg.e_m = i_hold_block ? detail::ds_xor(block, mine.b) : mine.b;
    const and_round1_msg peer = exch(msg);
    const std::uint8_t d = static_cast<std::uint8_t>(msg.d_bit ^ peer.d_bit);
    const simde__m128i e = i_hold_block
        ? detail::ds_xor(msg.e_m, peer.b_bit)
        : detail::ds_xor(peer.e_m, mine.b);
    simde__m128i z = detail::ds_xor(
        detail::ds_xor(detail::ds_gate(d, mine.b), detail::ds_gate(mine.a, e)),
        mine.c);
    if (i_hold_block)
        z = detail::ds_xor(z, detail::ds_gate(d, e));
    return z;
}

/// @brief Additive leaf share of a XOR-shared block. Limb width matches `leaf_add`.
template <int Me>
simde__m128i b2a_leaf(net::trio & net, simde__m128i xor_share,
    const b2a_pad_msg * pads, std::size_t limb_bytes)
{
    std::uint8_t bytes[16]{};
    std::memcpy(bytes, &xor_share, 16);
    std::vector<std::uint8_t> masked(128);
    for (int i = 0; i < 128; ++i)
    {
        const std::uint8_t bit = static_cast<std::uint8_t>(
            (bytes[i / 8] >> (i % 8)) & 1u);
        masked[static_cast<std::size_t>(i)] = static_cast<std::uint8_t>(
            bit ^ pads[i].r);
    }
    const auto peer = exch_vec<Me>(net, masked);
    const std::size_t limb_bits = limb_bytes * 8;
    const std::size_t nlimbs = 16 / limb_bytes;
    const std::uint64_t limb_mask = limb_bits == 64
        ? ~std::uint64_t{0}
        : ((std::uint64_t{1} << limb_bits) - 1);
    std::uint64_t acc[16]{};
    for (int i = 0; i < 128; ++i)
    {
        const std::uint8_t c = static_cast<std::uint8_t>(
            masked[static_cast<std::size_t>(i)] ^ peer[static_cast<std::size_t>(i)]);
        std::uint64_t s = pads[i].add;
        if (c != 0)
            s = Me == 0 ? (std::uint64_t{1} - pads[i].add)
                        : (std::uint64_t{0} - pads[i].add);
        const std::size_t limb = static_cast<std::size_t>(i) / limb_bits;
        const unsigned shift = static_cast<unsigned>(
            static_cast<std::size_t>(i) % limb_bits);
        acc[limb] += s << shift;
    }
    std::uint8_t outb[16]{};
    for (std::size_t k = 0; k < nlimbs; ++k)
    {
        const std::uint64_t v = acc[k] & limb_mask;
        std::memcpy(outb + k * limb_bytes, &v, limb_bytes);
    }
    simde__m128i out{};
    std::memcpy(&out, outb, 16);
    return out;
}

template <int Me, typename Concrete, typename Leaf>
Leaf open_leaf_sum(net::trio & net, Leaf mine)
{
    const role peer = Me == 0 ? role::p1 : role::p0;
    const Leaf theirs = net.exchange_with(peer, mine, net::msg::ring_vector);
    return leaf_add<Concrete>(mine, theirs);
}

/// @brief Public leaf CW from a shared naked leaf and this party's mask.
template <int Me, typename Concrete, typename Leaf>
Leaf cw_from_selected(net::trio & net, Leaf selected, Leaf mask, bool sign)
{
    constexpr bool char2 = utils::has_characteristic_two_v<Concrete>;
    Leaf share{};
    if constexpr (char2)
    {
        share = leaf_add<Concrete>(selected, mask);
    }
    else if (sign)
    {
        if constexpr (Me == 0)
            share = leaf_add<Concrete>(selected, mask);
        else
            share = dpf::subtract_leaf<Concrete>(selected, mask);
    }
    else if constexpr (Me == 0)
    {
        const Leaf neg = dpf::subtract_leaf<Concrete>(Leaf{}, selected);
        share = dpf::subtract_leaf<Concrete>(neg, mask);
    }
    else
    {
        share = dpf::subtract_leaf<Concrete>(mask, selected);
    }
    return open_leaf_sum<Me, Concrete>(net, share);
}

/// @brief Share of the leaf selected by the low `lg` bits of the target.
/// @details Party 0's `make(i)` is the candidate. Party 1's is zero. Each mux
///          multiplies the selector by both parties' deltas and keeps the sum
///          shared, so neither party sees the lane.
template <int Me, typename Concrete, typename Leaf, typename Exchange,
          typename Make, typename BitAt>
Leaf mux_leaf_share(net::trio & net, Exchange && exch, std::size_t lg,
    Make && make, BitAt && bit_at, const and_share_msg * ands,
    const b2a_pad_msg * b2a)
{
    static_assert(sizeof(Leaf) == sizeof(simde__m128i),
        "shared leaf mux covers a single AES block");
    constexpr bool char2 = utils::has_characteristic_two_v<Concrete>;
    constexpr std::size_t limb = [] {
        constexpr std::size_t s = sizeof(Concrete);
        return (s == 1 || s == 2 || s == 4 || s == 8) ? s : std::size_t{8};
    }();
    std::vector<Leaf> cand(std::size_t{1} << lg);
    for (std::size_t i = 0; i < cand.size(); ++i)
        cand[i] = make(static_cast<unsigned>(i));
    std::size_t pad_i = 0;
    std::size_t b2a_i = 0;
    for (std::size_t b = 0; b < lg; ++b)
    {
        const std::uint8_t bit = bit_at(b);
        std::vector<Leaf> next(cand.size() / 2);
        for (std::size_t k = 0; k < next.size(); ++k)
        {
            const Leaf lo = cand[2 * k];
            const Leaf hi = cand[2 * k + 1];
            const Leaf delta = dpf::subtract_leaf<Concrete>(hi, lo);
            simde__m128i dblock{};
            std::memcpy(&dblock, &delta, sizeof(dblock));
            const simde__m128i zero = simde_mm_setzero_si128();
            simde__m128i zx{};
            simde__m128i zy{};
            if constexpr (Me == 0)
            {
                zx = beaver_xor_share_bit(exch, true, bit, dblock, ands[pad_i++]);
                zy = beaver_xor_share_bit(exch, false, bit, zero, ands[pad_i++]);
            }
            else
            {
                zx = beaver_xor_share_bit(exch, false, bit, zero, ands[pad_i++]);
                zy = beaver_xor_share_bit(exch, true, bit, dblock, ands[pad_i++]);
            }
            Leaf prod{};
            if constexpr (char2)
            {
                const simde__m128i z = detail::ds_xor(zx, zy);
                std::memcpy(&prod, &z, sizeof(prod));
            }
            else
            {
                const simde__m128i ax = b2a_leaf<Me>(net, zx, b2a + b2a_i, limb);
                b2a_i += 128;
                const simde__m128i ay = b2a_leaf<Me>(net, zy, b2a + b2a_i, limb);
                b2a_i += 128;
                Leaf px{};
                Leaf py{};
                std::memcpy(&px, &ax, sizeof(px));
                std::memcpy(&py, &ay, sizeof(py));
                prod = leaf_add<Concrete>(px, py);
            }
            next[k] = leaf_add<Concrete>(lo, prod);
        }
        cand.swap(next);
    }
    return cand[0];
}

template <int Me>
std::uint8_t and_bit_share(net::trio & net, std::uint8_t x, std::uint8_t y,
    const bit_and_pad_msg & pad)
{
    const role peer = Me == 0 ? role::p1 : role::p0;
    const bit_and_mask_msg mine{
        static_cast<std::uint8_t>(x ^ pad.a),
        static_cast<std::uint8_t>(y ^ pad.b)};
    const bit_and_mask_msg got = net.exchange_with(peer, mine, net::msg::delta);
    const std::uint8_t d = static_cast<std::uint8_t>(mine.x ^ got.x);
    const std::uint8_t e = static_cast<std::uint8_t>(mine.y ^ got.y);
    return detail::ds_bit_and_party(d, e, pad.a, pad.b, pad.c, Me == 0);
}

template <int Me>
std::uint64_t b2a_bit(net::trio & net, std::uint8_t bit, const b2a_pad_msg & pad)
{
    const role peer = Me == 0 ? role::p1 : role::p0;
    const std::uint8_t masked = static_cast<std::uint8_t>(bit ^ pad.r);
    const std::uint8_t other = net.exchange_with(peer, masked, net::msg::delta);
    const std::uint8_t c = static_cast<std::uint8_t>(masked ^ other);
    if (c == 0)
        return pad.add;
    if constexpr (Me == 0)
        return std::uint64_t{1} - pad.add;
    else
        return std::uint64_t{0} - pad.add;
}

/// @brief Additive share of `bit * word`. `word` is meaningful only for the holder.
template <int Me>
std::uint64_t mul_known_word(net::trio & net, bool i_hold, std::uint64_t word,
    std::uint8_t bit, const word_bit_pad * pads)
{
    std::vector<bit_and_mask_msg> mine(64);
    for (int i = 0; i < 64; ++i)
    {
        const std::uint8_t wbit = i_hold
            ? static_cast<std::uint8_t>((word >> i) & 1u) : std::uint8_t{0};
        mine[static_cast<std::size_t>(i)] = bit_and_mask_msg{
            static_cast<std::uint8_t>(bit ^ pads[i].a),
            static_cast<std::uint8_t>(wbit ^ pads[i].b)};
    }
    const auto peer = exch_vec<Me>(net, mine);
    std::vector<std::uint8_t> masked(64);
    std::uint8_t prod[64]{};
    for (int i = 0; i < 64; ++i)
    {
        const auto & m = mine[static_cast<std::size_t>(i)];
        const auto & g = peer[static_cast<std::size_t>(i)];
        const std::uint8_t d = static_cast<std::uint8_t>(m.x ^ g.x);
        const std::uint8_t e = static_cast<std::uint8_t>(m.y ^ g.y);
        prod[i] = detail::ds_bit_and_party(
            d, e, pads[i].a, pads[i].b, pads[i].c, Me == 0);
        masked[static_cast<std::size_t>(i)] = static_cast<std::uint8_t>(
            prod[i] ^ pads[i].r);
    }
    const auto opened = exch_vec<Me>(net, masked);
    std::uint64_t acc = 0;
    for (int i = 0; i < 64; ++i)
    {
        const std::uint8_t c = static_cast<std::uint8_t>(
            masked[static_cast<std::size_t>(i)] ^ opened[static_cast<std::size_t>(i)]);
        std::uint64_t s = pads[i].add;
        if (c != 0)
            s = Me == 0 ? (std::uint64_t{1} - pads[i].add)
                        : (std::uint64_t{0} - pads[i].add);
        acc += s << i;
    }
    return acc;
}

/// @brief XOR-share of the AND of every bit in `bits`. `n == 0` returns 1 on p0.
template <int Me>
std::uint8_t and_tree(net::trio & net, const std::uint8_t * bits, std::size_t n,
    const bit_and_pad_msg * pads)
{
    if (n == 0)
        return Me == 0 ? std::uint8_t{1} : std::uint8_t{0};
    std::uint8_t acc = bits[0];
    for (std::size_t i = 1; i < n; ++i)
        acc = and_bit_share<Me>(net, acc, bits[i], pads[i - 1]);
    return acc;
}

template <int Me>
std::uint8_t open_bit(net::trio & net, std::uint8_t mine)
{
    const role peer = Me == 0 ? role::p1 : role::p0;
    const std::uint8_t other = net.exchange_with(peer, mine, net::msg::delta);
    return static_cast<std::uint8_t>(mine ^ other);
}

#endif