/// @file dpf/rss_seed.hpp
/// @brief Pairwise PRG seeds for honest-majority 3PC RSS preprocessing.
/// @details Parties 0,1,2 hold keys `k01`, `k12`, `k20`. Party `i` derives
///          both components of a replicated random value, a zero-sharing, and
///          the local cross term of an RSS multiply with no message.
#ifndef LIBDPF_INCLUDE_DPF_RSS_SEED_HPP__
#define LIBDPF_INCLUDE_DPF_RSS_SEED_HPP__

#include <array>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <utility>

#include "hedley/hedley.h"

#include "dpf/buffered_prg.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/random.hpp"
#include "dpf/secret_share.hpp"

namespace dpf
{
namespace rss
{

using seed_block = typename prg::aes128::block_type;

/// @brief Three pairwise master seeds. Role `i` holds seeds with `i±1 mod 3`.
struct seed_bundle
{
    seed_block k01{};
    seed_block k12{};
    seed_block k20{};
};

/// @brief Sample three independent pairwise seeds (dealer / setup).
HEDLEY_WARN_UNUSED_RESULT
inline seed_bundle sample_seed_bundle()
{
    seed_bundle b;
    b.k01 = randomness::sample_master_seed<prg::aes128>();
    b.k12 = randomness::sample_master_seed<prg::aes128>();
    b.k20 = randomness::sample_master_seed<prg::aes128>();
    return b;
}

/// @brief Party `me`'s view: the two seeds it shares with its neighbors.
struct party_seeds
{
    unsigned me = 0;
    seed_block with_prev{};  ///< shared with (me+2)%3
    seed_block with_next{};  ///< shared with (me+1)%3

    static party_seeds from_bundle(const seed_bundle & b, unsigned me)
    {
        if (me > 2)
            throw std::invalid_argument("rss party_seeds: me must be 0..2");
        party_seeds s;
        s.me = me;
        if (me == 0)
        {
            s.with_prev = b.k20;
            s.with_next = b.k01;
        }
        else if (me == 1)
        {
            s.with_prev = b.k01;
            s.with_next = b.k12;
        }
        else
        {
            s.with_prev = b.k12;
            s.with_next = b.k20;
        }
        return s;
    }
};

namespace detail
{

template <typename T>
T derive(seed_block master, std::uint32_t role, std::uint64_t index)
{
    randomness::lane_table<T, prg::aes128> table(master);
    return table.value_at(role, index);
}

inline seed_block pair_seed(unsigned a, unsigned b, const seed_bundle & bundle)
{
    const unsigned lo = a < b ? a : b;
    const unsigned hi = a < b ? b : a;
    if (lo == 0 && hi == 1)
        return bundle.k01;
    if (lo == 1 && hi == 2)
        return bundle.k12;
    if (lo == 0 && hi == 2)
        return bundle.k20;
    throw std::invalid_argument("rss pair_seed: bad pair");
}

}  // namespace detail

/// @brief Replicated random `r` at `index`. Party `me` holds `(r_me, r_{me+1})`.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
replicated_share<T, 0> random_replicated0(const party_seeds & s, std::uint64_t index)
{
    if (s.me != 0)
        throw std::invalid_argument("random_replicated0: wrong party");
    // r0 from k20 (parties 2,0), r1 from k01 (parties 0,1)
    const T r0 = detail::derive<T>(s.with_prev, 0, index);
    const T r1 = detail::derive<T>(s.with_next, 1, index);
    return replicated_share<T, 0>::from_raw(r0, r1);
}

template <typename T>
HEDLEY_WARN_UNUSED_RESULT
replicated_share<T, 1> random_replicated1(const party_seeds & s, std::uint64_t index)
{
    if (s.me != 1)
        throw std::invalid_argument("random_replicated1: wrong party");
    const T r1 = detail::derive<T>(s.with_prev, 1, index);
    const T r2 = detail::derive<T>(s.with_next, 2, index);
    return replicated_share<T, 1>::from_raw(r1, r2);
}

template <typename T>
HEDLEY_WARN_UNUSED_RESULT
replicated_share<T, 2> random_replicated2(const party_seeds & s, std::uint64_t index)
{
    if (s.me != 2)
        throw std::invalid_argument("random_replicated2: wrong party");
    const T r2 = detail::derive<T>(s.with_prev, 2, index);
    const T r0 = detail::derive<T>(s.with_next, 0, index);
    return replicated_share<T, 2>::from_raw(r2, r0);
}

/// @brief Party-erased random replicated share.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
std::pair<T, T> random_replicated_components(const party_seeds & s,
    std::uint64_t index)
{
    if (s.me == 0)
    {
        auto sh = random_replicated0<T>(s, index);
        return {sh.own, sh.next};
    }
    if (s.me == 1)
    {
        auto sh = random_replicated1<T>(s, index);
        return {sh.own, sh.next};
    }
    auto sh = random_replicated2<T>(s, index);
    return {sh.own, sh.next};
}

/// @brief Zero-sharing for reshare: party `i` holds `z_i` with sum zero.
/// @details `z_i = r_i - r_{i-1}` where `r` is a common pairwise stream on
///          each edge; the three differences cancel.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
T zero_share(const party_seeds & s, std::uint64_t index)
{
    // On edge with_next: parties me and next share stream role 0 → value α
    // On edge with_prev: parties me and prev share stream role 0 → value β
    // z_me = α - β  so sum over the ring is zero.
    const T alpha = detail::derive<T>(s.with_next, 10u + s.me, index);
    const T beta = detail::derive<T>(s.with_prev, 10u + ((s.me + 2) % 3), index);
    return static_cast<T>(alpha - beta);
}

/// @brief Local y-component of an RSS multiply before neighbor refresh.
/// @details Party `i` holds `(x_i, x_{i+1})` and `(y_i, y_{i+1})`. The local
///          product factor is `x_i*y_i + x_i*y_{i+1} + x_{i+1}*y_i` plus a
///          zero-mask so the three factors sum to `xy` after refresh.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
T rss_mul_local(const party_seeds & s, T x_own, T x_next, T y_own, T y_next,
    std::uint64_t index)
{
    const T cross = static_cast<T>(
        x_own * y_own + x_own * y_next + x_next * y_own);
    const T mask = zero_share<T>(s, index);
    return static_cast<T>(cross + mask);
}

/// @brief Full dealer view: derive every party's components for tests.
template <typename T>
struct replicated_triple
{
    replicated_share<T, 0> p0;
    replicated_share<T, 1> p1;
    replicated_share<T, 2> p2;
};

template <typename T>
HEDLEY_WARN_UNUSED_RESULT
replicated_triple<T> random_replicated_all(const seed_bundle & b,
    std::uint64_t index)
{
    auto s0 = party_seeds::from_bundle(b, 0);
    auto s1 = party_seeds::from_bundle(b, 1);
    auto s2 = party_seeds::from_bundle(b, 2);
    return replicated_triple<T>{
        random_replicated0<T>(s0, index),
        random_replicated1<T>(s1, index),
        random_replicated2<T>(s2, index)};
}

}  // namespace rss
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_RSS_SEED_HPP__