/// @file grotto/ring_switch.hpp
/// @brief Exact share conversion from \f$\mathbb{Z}/2^n\f$ into a residue group.
/// @details For an unsigned \f$n\f$-bit limb (\f$n\le 64\f$) with representatives
///          in \f$[0,2^n)\f$,
///          \f$\eta + r = x + w\cdot 2^n,\qquad w=\mathbf{1}[r+\eta\ge 2^n].\f$
///          In any modulus \f$M\f$,
///          \f$x \equiv \eta + (r\bmod M) - w\cdot(2^n\bmod M)\pmod M.\f$
///          The wrap bit times \f$2^n\bmod M\f$ must be shared inside the
///          destination group: a `uint64` comparison share is not a share
///          mod \f$M\f$. The dealer keys `lt(2^n \bmod M)` at the secret \f$r\f$
///          and, after \f$\eta\f$ opens, each party evaluates at the public
///          query \f$2^n-1-\eta\f$. That indicator is hot exactly on wrap,
///          including the \f$\eta=0\f$ case (no wrap). Party 0 adds public
///          \f$\eta\f$; both parties hold an additive split of \f$r\f$ in the
///          residue group.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license.

#ifndef LIBDPF_INCLUDE_GROTTO_RING_SWITCH_HPP__
#define LIBDPF_INCLUDE_GROTTO_RING_SWITCH_HPP__

#include "dpf.hpp"
#include "dpf/field128.hpp"
#include "dpf/p256_scalar.hpp"
#include "grotto/residue.hpp"

#include <cstddef>
#include <cstdint>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <utility>

namespace grotto
{

namespace ring_switch_detail
{

template <typename InputT>
HEDLEY_NO_THROW
constexpr unsigned limb_bits() noexcept
{
    static_assert(std::is_unsigned_v<InputT>,
        "ring switch limb domain must be an unsigned integer");
    constexpr unsigned bits = static_cast<unsigned>(dpf::utils::bitlength_of_v<InputT>);
    static_assert(bits >= 1 && bits <= 64, "ring switch supports 1..64 bit limbs");
    return bits;
}

template <typename InputT>
HEDLEY_NO_THROW
constexpr InputT limb_mask() noexcept
{
    constexpr unsigned bits = limb_bits<InputT>();
    if constexpr (bits == 64)
        return ~InputT{0};
    else
        return static_cast<InputT>((InputT{1} << bits) - InputT{1});
}

/// @brief Public query \f$2^n - 1 - \eta\f$ for the wrap `lt` key.
template <typename InputT>
HEDLEY_NO_THROW
constexpr InputT wrap_query(InputT eta) noexcept
{
    return static_cast<InputT>(limb_mask<InputT>() - (eta & limb_mask<InputT>()));
}

/// @brief \f$2^n\f$ as an integer for constructing the residue payload.
template <typename InputT>
HEDLEY_NO_THROW
constexpr unsigned __int128 pow2_n() noexcept
{
    constexpr unsigned bits = limb_bits<InputT>();
    if constexpr (bits == 64)
        return static_cast<unsigned __int128>(1) << 64;
    else
        return static_cast<unsigned __int128>(1) << bits;
}

template <typename T, typename = void>
struct has_modulus64 : std::false_type {};
template <typename T>
struct has_modulus64<T, std::void_t<decltype(T::modulus)>> : std::true_type {};

template <typename T, typename = void>
struct has_modulus128 : std::false_type {};
template <typename T>
struct has_modulus128<T, std::void_t<decltype(T::modulus_lo)>> : std::true_type {};

template <typename Residue, typename InputT>
Residue pow2_payload()
{
    if constexpr (std::is_same_v<Residue, dpf::field128>)
    {
        constexpr unsigned bits = limb_bits<InputT>();
        if constexpr (bits == 64)
            return dpf::field128::from_lane(0, 1);
        else
            return Residue{static_cast<std::uint64_t>(InputT{1} << bits)};
    }
    else if constexpr (std::is_same_v<Residue, dpf::p256_scalar>)
    {
        constexpr unsigned bits = limb_bits<InputT>();
        if constexpr (bits == 64)
            return Residue{static_cast<unsigned __int128>(1) << 64};
        else
            return Residue{static_cast<std::uint64_t>(InputT{1} << bits)};
    }
    else if constexpr (has_modulus64<Residue>::value)
    {
        return Residue{static_cast<std::uint64_t>(
            pow2_n<InputT>() % Residue::modulus)};
    }
    else if constexpr (has_modulus128<Residue>::value)
    {
        return Residue::from_u128(pow2_n<InputT>());
    }
    else
    {
        static_assert(sizeof(Residue) == 0, "unsupported ring switch residue");
        return Residue{};
    }
}

template <typename Residue, typename InputT>
Residue reduce_limb(InputT v)
{
    return Residue{static_cast<std::uint64_t>(v & limb_mask<InputT>())};
}

}  // namespace ring_switch_detail

/// @brief Dealer material for one exact ring switch into `Residue`.
template <typename InputT, typename Residue>
struct ring_switch_keys
{
    static_assert(std::is_unsigned_v<InputT>,
        "ring switch limb domain must be an unsigned integer");
    static_assert(dpf::utils::bitlength_of_v<InputT> <= 64,
        "ring switch supports at most 64-bit limbs");

    using input_type = InputT;
    using residue_type = Residue;
    using key_pair = decltype(dpf::make_dpf(std::declval<InputT>(),
        dpf::lt(std::declval<Residue>())));
    using key_pair_v = decltype(dpf::make_dpf(std::declval<InputT>(),
        dpf::lt(std::declval<Residue>()), dpf::verifiable{}));

    InputT r{};
    bool verifiable = false;
    std::optional<key_pair> keys{};
    std::optional<key_pair_v> keys_v{};
    /// @brief Additive split of `r` in `Residue`, indexed `[party]`.
    Residue r_share[2]{};
};
/// \complexity One `dpf::make_dpf` of `lt(2^n mod M)` at the masked limb `r`, plus one `from_seed` to split `r` in `Residue`.
/// `n` is `bitlength_of_v<InputT>` and the header rejects `n` outside `1 .. 64`.
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One comparison key and two `Residue` words (`r_share[2]`).
/// @note The comparison payload is `2^n mod M` in the destination group. A `uint64_t` comparison share is not a share modulo `M`.
/// @see grotto::zn64
/// @see grotto::nmod
/// @see [Exact ring switch](@ref ring_switch)

template <typename Residue, typename InputT>
ring_switch_keys<InputT, Residue> make_ring_switch_keys(InputT r)
{
    using namespace ring_switch_detail;
    ring_switch_keys<InputT, Residue> mat;
    mat.r = static_cast<InputT>(r & limb_mask<InputT>());
    mat.verifiable = false;
    const Residue payload = pow2_payload<Residue, InputT>();
    mat.keys = dpf::make_dpf(mat.r, dpf::lt(payload));
    const auto seed = dpf::uniform_sample<simde__m128i>();
    mat.r_share[0] = Residue::from_seed(&seed, sizeof(seed));
    mat.r_share[1] = reduce_limb<Residue>(mat.r) - mat.r_share[0];
    return mat;
}
/// \complexity One `dpf::make_dpf` of `lt(2^n mod M)` at the masked limb `r`, plus one `from_seed` to split `r` in `Residue`.
/// `n` is `bitlength_of_v<InputT>` and the header rejects `n` outside `1 .. 64`.
/// \rounds No party interaction.
/// \communication None inside this function.
/// \preprocessing One comparison key and two `Residue` words (`r_share[2]`).
/// @note The comparison payload is `2^n mod M` in the destination group. A `uint64_t` comparison share is not a share modulo `M`.
/// @see grotto::zn64
/// @see grotto::nmod
/// @see [Exact ring switch](@ref ring_switch)

template <typename Residue, typename InputT>
ring_switch_keys<InputT, Residue> make_ring_switch_keys(InputT r, dpf::verifiable)
{
    using namespace ring_switch_detail;
    ring_switch_keys<InputT, Residue> mat;
    mat.r = static_cast<InputT>(r & limb_mask<InputT>());
    mat.verifiable = true;
    const Residue payload = pow2_payload<Residue, InputT>();
    mat.keys_v = dpf::make_dpf(mat.r, dpf::lt(payload), dpf::verifiable{});
    const auto seed = dpf::uniform_sample<simde__m128i>();
    mat.r_share[0] = Residue::from_seed(&seed, sizeof(seed));
    mat.r_share[1] = reduce_limb<Residue>(mat.r) - mat.r_share[0];
    return mat;
}

/// @brief Cleartext \f$x \bmod M\f$ from the limb identity.
/// \complexity A constant number of limb adds and one comparison of `r + eta` against `2^n`. `Θ(1)`. No keys.
/// @see grotto::ring_switch_eval
template <typename Residue, typename InputT>
Residue ring_switch_clear(InputT x, InputT r, InputT eta)
{
    using namespace ring_switch_detail;
    const InputT mask = limb_mask<InputT>();
    x = static_cast<InputT>(x & mask);
    r = static_cast<InputT>(r & mask);
    eta = static_cast<InputT>(eta & mask);
    const unsigned __int128 sum =
        static_cast<unsigned __int128>(r) + eta;
    const unsigned __int128 modn = pow2_n<InputT>();
    const unsigned w = (sum >= modn) ? 1u : 0u;
    Residue out = reduce_limb<Residue>(eta) + reduce_limb<Residue>(r);
    if (w)
        out = out - pow2_payload<Residue, InputT>();
    return out;
}

/// @brief One party's share of \f$x\f$ in `Residue` after `eta` is public.
/// \complexity One `dpf::eval_point` on the wrap key at the public query `2^n - 1 - eta`, then a constant number of `Residue` additions.
/// Party 0 also adds `eta` reduced into `Residue`. `n ≤ 64`.
/// Extra space `Θ(1)`.
/// \rounds None. `eta` is an argument; the wrap bit is the comparison output, not a separate opening.
/// \communication None.
/// \preprocessing None created here. Uses the one key and the two `r_share` words from `make_ring_switch_keys`.
/// @note The value subtracted on wrap is the keyed payload `2^n mod M`, not a `uint64_t` share.
/// @see grotto::zn64
/// @see grotto::nmod
/// @see [Exact ring switch](@ref ring_switch)
template <std::size_t Party, typename InputT, typename Residue>
Residue ring_switch_eval(
    const ring_switch_keys<InputT, Residue> & mat,
    InputT eta,
    dpf::proof_token * pi = nullptr)
{
    static_assert(Party < 2, "ring switch party is 0 or 1");
    using namespace ring_switch_detail;
    eta = static_cast<InputT>(eta & limb_mask<InputT>());
    const InputT query = wrap_query(eta);

    Residue wrap{};
    if (mat.verifiable)
    {
        if (!mat.keys_v)
            throw std::invalid_argument("ring switch: missing verifiable keys");
        auto & key = std::get<Party>(*mat.keys_v);
        if (pi != nullptr)
        {
            wrap = dpf::eval_point<Residue>(dpf::cmp, key, query,
                dpf::prove(*pi)).raw();
        }
        else
        {
            wrap = dpf::eval_point<Residue>(dpf::cmp, key, query).raw();
        }
    }
    else
    {
        if (!mat.keys)
            throw std::invalid_argument("ring switch: missing keys");
        auto & key = std::get<Party>(*mat.keys);
        wrap = dpf::eval_point<Residue>(dpf::cmp, key, query).raw();
    }

    Residue out = mat.r_share[Party] - wrap;
    if constexpr (Party == 0)
        out = out + reduce_limb<Residue>(eta);
    return out;
}

}  // namespace grotto

#endif  // LIBDPF_INCLUDE_GROTTO_RING_SWITCH_HPP__