/// @file dpf/secret_share.hpp
/// @brief Thin (2,2) additive and subtractive secret-share wrappers.
/// @details Layout-identical to `T`. Party is a compile-time `0` or `1`.
///          Reconstruction: additive opens by sum, subtractive by
///          `share0 - share1`. Linear combinations of same-party shares are
///          supported; mixing additive with subtractive applies the correct
///          party coefficient. A plaintext absorbs on party 0 only.
/// @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_SECRET_SHARE_HPP__
#define LIBDPF_INCLUDE_DPF_SECRET_SHARE_HPP__

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <ostream>
#include <type_traits>
#include <utility>

#include "hedley/hedley.h"

#include "dpf/twiddle.hpp"

namespace dpf
{

/// @brief Sharing scheme tag.
enum class sharing : unsigned char
{
    additive = 0,
    subtractive = 1
};

template <typename T, std::size_t Party, sharing Scheme>
struct secret_share;

template <typename T, std::size_t Party>
using additive_share = secret_share<T, Party, sharing::additive>;

template <typename T, std::size_t Party>
using subtractive_share = secret_share<T, Party, sharing::subtractive>;

template <typename T>
struct is_secret_share : std::false_type
{
};

template <typename T, std::size_t Party, sharing Scheme>
struct is_secret_share<secret_share<T, Party, Scheme>> : std::true_type
{
};

template <typename T>
inline constexpr bool is_secret_share_v = is_secret_share<std::decay_t<T>>::value;

template <typename T>
struct share_party;

template <typename T, std::size_t Party, sharing Scheme>
struct share_party<secret_share<T, Party, Scheme>>
  : std::integral_constant<std::size_t, Party>
{
};

template <typename T>
inline constexpr std::size_t share_party_v = share_party<std::decay_t<T>>::value;

template <typename T>
struct share_scheme;

template <typename T, std::size_t Party, sharing Scheme>
struct share_scheme<secret_share<T, Party, Scheme>>
  : std::integral_constant<sharing, Scheme>
{
};

template <typename T>
inline constexpr sharing share_scheme_v = share_scheme<std::decay_t<T>>::value;

template <typename T>
struct share_value_type;

template <typename T, std::size_t Party, sharing Scheme>
struct share_value_type<secret_share<T, Party, Scheme>>
{
    using type = T;
};

template <typename T>
using share_value_type_t = typename share_value_type<std::decay_t<T>>::type;

namespace detail
{

/// @brief Party coefficient of the secret for this scheme: additive always +1;
/// subtractive is +1 for party 0 and −1 for party 1.
/// @tparam Scheme scheme
/// @tparam Party party index, `0` or `1`
/// @tparam T value type
/// @param v the `v`
/// @return Party coefficient of the secret for this scheme: additive always +1; subtractive is +1
/// for party 0 and −1 for party 1
template <sharing Scheme, std::size_t Party, typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr T party_coeff_times(const T & v) noexcept
{
    if constexpr (Scheme == sharing::additive || Party == 0)
        return v;
    else if constexpr (std::is_integral_v<T> && std::is_signed_v<T>)
    {
        // Signed negation of the minimum is undefined. The two's-complement
        // negation is well-defined on the unsigned width.
        using unsigned_type = std::make_unsigned_t<T>;
        return static_cast<T>(static_cast<unsigned_type>(0)
            - static_cast<unsigned_type>(v));
    }
    else
        return static_cast<T>(-v);
}

}  // namespace detail

template <typename T, std::size_t Party, sharing Scheme>
struct secret_share
{
    static_assert(Party == 0 || Party == 1,
        "secret_share party must be 0 or 1");

    using value_type = T;
    static constexpr std::size_t party = Party;
    static constexpr sharing scheme = Scheme;

    T value{};

    secret_share() = default;
    HEDLEY_NO_THROW
    secret_share(const secret_share &) noexcept = default;
    HEDLEY_NO_THROW
    secret_share(secret_share &&) noexcept = default;
    HEDLEY_NO_THROW
    secret_share & operator=(const secret_share &) noexcept = default;
    HEDLEY_NO_THROW
    secret_share & operator=(secret_share &&) noexcept = default;
    ~secret_share() = default;

    /// @brief Bit-preserving construction. Does not apply a party coefficient.
    /// @param v the `v`
    /// @return Bit-preserving construction
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    static constexpr secret_share from_raw(T v) noexcept
    {
        secret_share s;
        s.value = v;
        return s;
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_CONST
    constexpr const T & raw() const noexcept { return value; }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr T & raw() noexcept { return value; }

    /// @brief Secret-preserving conversion to an additive share of the same party.
    /// @return Secret-preserving conversion to an additive share of the same party
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr additive_share<T, Party> as_additive() const noexcept
    {
        if constexpr (Scheme == sharing::additive)
            return additive_share<T, Party>::from_raw(value);
        // subtractive → additive: party 0 keeps bits; party 1 negates.
        return additive_share<T, Party>::from_raw(
            detail::party_coeff_times<sharing::subtractive, Party>(value));
    }

    /// @brief Secret-preserving conversion to a subtractive share of the same party.
    /// @return Secret-preserving conversion to a subtractive share of the same party
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr subtractive_share<T, Party> as_subtractive() const noexcept
    {
        if constexpr (Scheme == sharing::subtractive)
            return subtractive_share<T, Party>::from_raw(value);
        // additive → subtractive: party 0 keeps bits; party 1 negates.
        return subtractive_share<T, Party>::from_raw(
            detail::party_coeff_times<sharing::additive, Party>(value));
    }

    /// @brief Bit-preserving retag (no secret-preserving sign fix).
    /// @tparam NewScheme new scheme
    /// @tparam NewParty new party
    /// @return Bit-preserving retag (no secret-preserving sign fix)
    template <sharing NewScheme, std::size_t NewParty = Party>
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    HEDLEY_NO_THROW
    constexpr secret_share<T, NewParty, NewScheme> retag() const noexcept
    {
        return secret_share<T, NewParty, NewScheme>::from_raw(value);
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr secret_share operator-() const noexcept
    {
        if constexpr (std::is_integral_v<T> && std::is_signed_v<T>)
        {
            using unsigned_type = std::make_unsigned_t<T>;
            return from_raw(static_cast<T>(static_cast<unsigned_type>(0)
                - static_cast<unsigned_type>(value)));
        }
        else
            return from_raw(static_cast<T>(-value));
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr secret_share & operator+=(const secret_share & rhs) noexcept
    {
        value = static_cast<T>(value + rhs.value);
        return *this;
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    constexpr secret_share & operator-=(const secret_share & rhs) noexcept
    {
        value = static_cast<T>(value - rhs.value);
        return *this;
    }

    template <typename Scalar,
              std::enable_if_t<!is_secret_share_v<Scalar>, int> = 0>
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    constexpr secret_share & operator*=(const Scalar & c) noexcept
    {
        value = static_cast<T>(value * static_cast<T>(c));
        return *this;
    }

    /// @brief Absorb a public plaintext on party 0 only.
    /// @tparam Plain plain
    /// @tparam T value type
    /// @param c the `c`
    /// @return `*this`
    template <typename Plain,
              std::enable_if_t<!is_secret_share_v<Plain>
                  && std::is_convertible_v<Plain, T>, int> = 0>
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    constexpr secret_share & operator+=(const Plain & c) noexcept
    {
        if constexpr (Party == 0)
            value = static_cast<T>(value + static_cast<T>(c));
        return *this;
    }

    template <typename Plain,
              std::enable_if_t<!is_secret_share_v<Plain>
                  && std::is_convertible_v<Plain, T>, int> = 0>
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NO_THROW
    constexpr secret_share & operator-=(const Plain & c) noexcept
    {
        if constexpr (Party == 0)
            value = static_cast<T>(value - static_cast<T>(c));
        return *this;
    }
};

// ---------------------------------------------------------------------------
// Same-scheme, same-party arithmetic
// ---------------------------------------------------------------------------

template <typename T, std::size_t Party, sharing Scheme>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, Scheme> operator+(
    secret_share<T, Party, Scheme> lhs,
    const secret_share<T, Party, Scheme> & rhs) noexcept
{
    lhs += rhs;
    return lhs;
}

template <typename T, std::size_t Party, sharing Scheme>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, Scheme> operator-(
    secret_share<T, Party, Scheme> lhs,
    const secret_share<T, Party, Scheme> & rhs) noexcept
{
    lhs -= rhs;
    return lhs;
}

template <typename T, std::size_t Party, sharing Scheme, typename Scalar,
          std::enable_if_t<!is_secret_share_v<Scalar>, int> = 0>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, Scheme> operator*(
    secret_share<T, Party, Scheme> lhs, const Scalar & c) noexcept
{
    lhs *= c;
    return lhs;
}

template <typename T, std::size_t Party, sharing Scheme, typename Scalar,
          std::enable_if_t<!is_secret_share_v<Scalar>, int> = 0>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, Scheme> operator*(
    const Scalar & c, secret_share<T, Party, Scheme> rhs) noexcept
{
    rhs *= c;
    return rhs;
}

// ---------------------------------------------------------------------------
// Cross-scheme, same-party: keep the left-hand scheme; party 1 flips the
// operand whose scheme differs from the result.
// ---------------------------------------------------------------------------

template <typename T, std::size_t Party, sharing LhsScheme, sharing RhsScheme,
          std::enable_if_t<LhsScheme != RhsScheme, int> = 0>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, LhsScheme> operator+(
    const secret_share<T, Party, LhsScheme> & lhs,
    const secret_share<T, Party, RhsScheme> & rhs) noexcept
{
    if constexpr (Party == 0)
        return secret_share<T, Party, LhsScheme>::from_raw(
            static_cast<T>(lhs.raw() + rhs.raw()));
    else
        return secret_share<T, Party, LhsScheme>::from_raw(
            static_cast<T>(lhs.raw() - rhs.raw()));
}

template <typename T, std::size_t Party, sharing LhsScheme, sharing RhsScheme,
          std::enable_if_t<LhsScheme != RhsScheme, int> = 0>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, LhsScheme> operator-(
    const secret_share<T, Party, LhsScheme> & lhs,
    const secret_share<T, Party, RhsScheme> & rhs) noexcept
{
    if constexpr (Party == 0)
        return secret_share<T, Party, LhsScheme>::from_raw(
            static_cast<T>(lhs.raw() - rhs.raw()));
    else
        return secret_share<T, Party, LhsScheme>::from_raw(
            static_cast<T>(lhs.raw() + rhs.raw()));
}

// ---------------------------------------------------------------------------
// Plaintext absorb (party 0 only)
// ---------------------------------------------------------------------------

template <typename T, std::size_t Party, sharing Scheme, typename Plain,
          std::enable_if_t<!is_secret_share_v<Plain>
              && std::is_convertible_v<Plain, T>, int> = 0>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, Scheme> operator+(
    secret_share<T, Party, Scheme> lhs, const Plain & c) noexcept
{
    lhs += c;
    return lhs;
}

template <typename T, std::size_t Party, sharing Scheme, typename Plain,
          std::enable_if_t<!is_secret_share_v<Plain>
              && std::is_convertible_v<Plain, T>, int> = 0>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, Scheme> operator+(
    const Plain & c, secret_share<T, Party, Scheme> rhs) noexcept
{
    rhs += c;
    return rhs;
}

template <typename T, std::size_t Party, sharing Scheme, typename Plain,
          std::enable_if_t<!is_secret_share_v<Plain>
              && std::is_convertible_v<Plain, T>, int> = 0>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr secret_share<T, Party, Scheme> operator-(
    secret_share<T, Party, Scheme> lhs, const Plain & c) noexcept
{
    lhs -= c;
    return lhs;
}

// ---------------------------------------------------------------------------
// Equality (same party, same scheme) — compare raw bits
// ---------------------------------------------------------------------------

template <typename T, std::size_t Party, sharing Scheme>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr bool operator==(const secret_share<T, Party, Scheme> & lhs,
    const secret_share<T, Party, Scheme> & rhs) noexcept
{
    return lhs.raw() == rhs.raw();
}

template <typename T, std::size_t Party, sharing Scheme>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr bool operator!=(const secret_share<T, Party, Scheme> & lhs,
    const secret_share<T, Party, Scheme> & rhs) noexcept
{
    return !(lhs == rhs);
}

// ---------------------------------------------------------------------------
// Reconstruction
// ---------------------------------------------------------------------------

template <typename T, sharing Scheme>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr T reconstruct(const secret_share<T, 0, Scheme> & s0,
    const secret_share<T, 1, Scheme> & s1) noexcept
{
    if constexpr (std::is_integral_v<T> && std::is_signed_v<T>)
    {
        using unsigned_type = std::make_unsigned_t<T>;
        if constexpr (Scheme == sharing::additive)
            return static_cast<T>(static_cast<unsigned_type>(s0.raw())
                + static_cast<unsigned_type>(s1.raw()));
        else
            return static_cast<T>(static_cast<unsigned_type>(s0.raw())
                - static_cast<unsigned_type>(s1.raw()));
    }
    else if constexpr (Scheme == sharing::additive)
        return static_cast<T>(s0.raw() + s1.raw());
    else
        return static_cast<T>(s0.raw() - s1.raw());
}

template <typename T, sharing Scheme>
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
HEDLEY_NO_THROW
constexpr T reconstruct(const secret_share<T, 1, Scheme> & s1,
    const secret_share<T, 0, Scheme> & s0) noexcept
{
    return reconstruct(s0, s1);
}

// ---------------------------------------------------------------------------
// Plaintext splits (share1 = 0)
// ---------------------------------------------------------------------------

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_CONST
HEDLEY_NO_THROW
constexpr auto make_additive_shares(T secret) noexcept
{
    using T_ = std::remove_cv_t<std::remove_reference_t<T>>;
    return std::make_pair(
        additive_share<T_, 0>::from_raw(static_cast<T_>(secret)),
        additive_share<T_, 1>::from_raw(T_{}));
}

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_CONST
HEDLEY_NO_THROW
constexpr auto make_subtractive_shares(T secret) noexcept
{
    using T_ = std::remove_cv_t<std::remove_reference_t<T>>;
    return std::make_pair(
        subtractive_share<T_, 0>::from_raw(static_cast<T_>(secret)),
        subtractive_share<T_, 1>::from_raw(T_{}));
}

// ---------------------------------------------------------------------------
// Party-tagged DPF key wrapper
// ---------------------------------------------------------------------------

template <typename T>
struct is_party_key : std::false_type
{
};

template <std::size_t Party, typename Key>
struct party_key : Key
{
    static_assert(Party == 0 || Party == 1, "party_key party must be 0 or 1");

    static constexpr std::size_t party = Party;
    using key_type = Key;

    party_key() = default;

    HEDLEY_ALWAYS_INLINE
    explicit party_key(Key k)
      : Key(std::move(k))
    {
#ifndef NDEBUG
        assert(static_cast<std::size_t>(
            static_cast<bool>(dpf::get_lo_bit(this->root()))) == Party);
#endif
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    Key & key() noexcept { return static_cast<Key &>(*this); }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    const Key & key() const noexcept { return static_cast<const Key &>(*this); }

    /// @brief Party-tagged additive share of the comparison absorb addend.
    /// @return Party-tagged additive share of the comparison absorb addend
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    auto cmp_addend() const noexcept
    {
        return additive_share<std::uint64_t, Party>::from_raw(
            Key::cmp_addend());
    }
};

template <std::size_t Party, typename Key>
struct is_party_key<party_key<Party, Key>> : std::true_type
{
};

template <typename T>
inline constexpr bool is_party_key_v = is_party_key<std::decay_t<T>>::value;

template <typename T>
struct party_of;  // incomplete for non-`party_key` (fail loudly on misuse)

template <std::size_t Party, typename Key>
struct party_of<party_key<Party, Key>>
  : std::integral_constant<std::size_t, Party>
{
};

template <typename T>
inline constexpr std::size_t party_of_v = party_of<std::decay_t<T>>::value;

/// @brief Strip a `party_key` wrapper; bare keys are unchanged. Memoizers and other
/// tree-layout helpers key on the underlying DPF key type so a memoizer built
/// for party 0 also accepts party 1.
/// @tparam T value type
template <typename T>
struct unwrap_party_key
{
    using type = std::decay_t<T>;
};
template <std::size_t Party, typename Key>
struct unwrap_party_key<party_key<Party, Key>>
{
    using type = Key;
};
template <typename T>
using unwrap_party_key_t = typename unwrap_party_key<std::decay_t<T>>::type;

template <std::size_t Party, typename Key>
HEDLEY_ALWAYS_INLINE
auto make_party_key(Key && k)
{
    return party_key<Party, std::decay_t<Key>>(std::forward<Key>(k));
}

template <typename Key0, typename Key1>
HEDLEY_ALWAYS_INLINE
auto make_party_key_pair(Key0 && k0, Key1 && k1)
{
    using K = std::decay_t<Key0>;
    static_assert(std::is_same_v<K, std::decay_t<Key1>>,
        "make_party_key_pair: both keys must have the same type");
    return std::make_pair(
        party_key<0, K>(std::forward<Key0>(k0)),
        party_key<1, K>(std::forward<Key1>(k1)));
}

template <typename CharT, typename Traits, typename T, std::size_t Party,
          sharing Scheme>
std::basic_ostream<CharT, Traits> & operator<<(
    std::basic_ostream<CharT, Traits> & os,
    const secret_share<T, Party, Scheme> & s)
{
    return os << s.raw();
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_SECRET_SHARE_HPP__