/// @file grotto/residue.hpp
/// @brief Compile-time modular residues as comparison payload groups.
/// @details `zn64<Mod>` and `zn128<Lo,Hi>` supply `from_seed`, `+`, and unary
///          `-`, which selects the payload-group comparison path. The modulus
///          is a template argument so keygen needs no runtime side channel.
/// @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_RESIDUE_HPP__
#define LIBDPF_INCLUDE_GROTTO_RESIDUE_HPP__

#include <cstddef>
#include <cstdint>
#include <cstring>
#include <type_traits>

#include "hedley/hedley.h"

namespace grotto
{

/// @brief Residue class modulo a 64-bit compile-time modulus.
/// @tparam Mod compile-time modulus, greater than 1
/// @see grotto::make_ring_switch_keys
/// @see grotto::ring_switch_factor
/// \complexity Construction, `from_seed`, addition, and negation are one remainder each. `Θ(1)` time and extra space.
template <std::uint64_t Mod>
class zn64
{
    static_assert(Mod > 1, "zn64 modulus must be greater than 1");

  public:
    /// \complexity One remainder modulo `Mod`, with a sign branch for negative `T`. `Θ(1)`.
    /// @see grotto::ring_switch_eval
    static constexpr std::uint64_t modulus = Mod;
    static constexpr bool dpf_point_group = true;

    HEDLEY_ALWAYS_INLINE
    constexpr zn64() noexcept = default;

    template <typename T, typename = std::enable_if_t<std::is_integral_v<T>>>
    HEDLEY_ALWAYS_INLINE
    constexpr zn64(T v) noexcept
    {
        assign_integer(v);
    }

    HEDLEY_ALWAYS_INLINE
    constexpr zn64(unsigned __int128 v) noexcept
    {
        val_ = static_cast<std::uint64_t>(v % Mod);
    }
    /// \complexity Copies at most 16 seed bytes and reduces the 128-bit word modulo `Mod`. `Θ(1)`.
    /// @see grotto::make_ring_switch_keys

    HEDLEY_ALWAYS_INLINE
    static zn64 from_seed(const void * bytes, std::size_t n) noexcept
    {
        unsigned char buf[16]{};
        if (n > sizeof(buf))
            n = sizeof(buf);
        std::memcpy(buf, bytes, n);
        std::uint64_t w[2]{};
        std::memcpy(w, buf, sizeof(w));
        const unsigned __int128 wide = static_cast<unsigned __int128>(w[0])
            | (static_cast<unsigned __int128>(w[1]) << 64);
        return from_reduced(static_cast<std::uint64_t>(wide % Mod));
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr std::uint64_t raw() const noexcept { return val_; }

    HEDLEY_ALWAYS_INLINE
    friend constexpr zn64 operator+(zn64 a, zn64 b) noexcept
    {
        const unsigned __int128 s =
            static_cast<unsigned __int128>(a.val_) + b.val_;
        return from_reduced(static_cast<std::uint64_t>(s % Mod));
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr zn64 operator-(zn64 a) noexcept
    {
        if (a.val_ == 0)
            return a;
        return from_reduced(Mod - a.val_);
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr zn64 operator-(zn64 a, zn64 b) noexcept
    {
        return a + (-b);
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr bool operator==(zn64 a, zn64 b) noexcept
    {
        return a.val_ == b.val_;
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr bool operator!=(zn64 a, zn64 b) noexcept
    {
        return !(a == b);
    }

  private:
    std::uint64_t val_{};

    HEDLEY_ALWAYS_INLINE
    static constexpr zn64 from_reduced(std::uint64_t v) noexcept
    {
        zn64 out;
        out.val_ = v;
        return out;
    }

    template <typename T>
    HEDLEY_ALWAYS_INLINE
    constexpr void assign_integer(T v) noexcept
    {
        if constexpr (std::is_signed_v<T>)
        {
            if (v < 0)
            {
                using U = std::make_unsigned_t<T>;
                const auto mag = static_cast<U>(0) - static_cast<U>(v);
                val_ = static_cast<std::uint64_t>(
                    Mod - (static_cast<std::uint64_t>(mag) % Mod));
                if (val_ == Mod)
                    val_ = 0;
                return;
            }
            val_ = static_cast<std::uint64_t>(
                static_cast<std::make_unsigned_t<T>>(v) % Mod);
        }
        else
        {
            val_ = static_cast<std::uint64_t>(
                static_cast<unsigned __int128>(v) % Mod);
        }
    }
};

/// @brief Residue class modulo a 128-bit compile-time modulus `(Hi << 64) | Lo`.
/// @tparam Lo low 64 bits of the modulus
/// @tparam Hi high 64 bits of the modulus
/// @see grotto::make_ring_switch_keys
/// @see grotto::zn64
/// \complexity Constructors reduce one integer modulo that modulus. `from_seed` folds at most 16 bytes. Addition is one 128-bit sum and a remainder. `Θ(1)` time and extra space.
template <std::uint64_t Lo, std::uint64_t Hi>
class zn128
{
    static_assert(Lo != 0 || Hi != 0, "zn128 modulus must be nonzero");
    static_assert(Hi != 0 || Lo > 1, "zn128 modulus must be greater than 1");

  public:
    static constexpr std::uint64_t modulus_lo = Lo;
    static constexpr std::uint64_t modulus_hi = Hi;
    static constexpr bool dpf_point_group = true;

    HEDLEY_ALWAYS_INLINE
    constexpr zn128() noexcept = default;

    template <typename T, typename = std::enable_if_t<std::is_integral_v<T>>>
    HEDLEY_ALWAYS_INLINE
    constexpr zn128(T v) noexcept
    {
        assign_integer(v);
    }

    HEDLEY_ALWAYS_INLINE
    static zn128 from_u128(unsigned __int128 v) noexcept
    {
        return reduce(static_cast<std::uint64_t>(v),
            static_cast<std::uint64_t>(v >> 64));
    }

    HEDLEY_ALWAYS_INLINE
    static zn128 from_seed(const void * bytes, std::size_t n) noexcept
    {
        unsigned char buf[16]{};
        if (n > sizeof(buf))
            n = sizeof(buf);
        std::memcpy(buf, bytes, n);
        std::uint64_t w[2]{};
        std::memcpy(w, buf, sizeof(w));
        return reduce(w[0], w[1]);
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr std::uint64_t lo() const noexcept { return lo_; }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr std::uint64_t hi() const noexcept { return hi_; }

    HEDLEY_ALWAYS_INLINE
    friend constexpr zn128 operator+(zn128 a, zn128 b) noexcept
    {
        const unsigned __int128 low =
            static_cast<unsigned __int128>(a.lo_) + b.lo_;
        const unsigned __int128 high =
            static_cast<unsigned __int128>(a.hi_) + b.hi_ + (low >> 64);
        std::uint64_t lo = static_cast<std::uint64_t>(low);
        std::uint64_t hi = static_cast<std::uint64_t>(high);
        const std::uint64_t top = static_cast<std::uint64_t>(high >> 64);
        // Two reduced residues sum to less than 2M < 2^129.
        if (top != 0)
            sub_words(lo, hi, Lo, Hi, lo, hi);
        if (hi > Hi || (hi == Hi && lo >= Lo))
            sub_words(lo, hi, Lo, Hi, lo, hi);
        return from_reduced(lo, hi);
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr zn128 operator-(zn128 a) noexcept
    {
        if (a.lo_ == 0 && a.hi_ == 0)
            return a;
        std::uint64_t ol = 0, oh = 0;
        sub_words(Lo, Hi, a.lo_, a.hi_, ol, oh);
        return from_reduced(ol, oh);
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr zn128 operator-(zn128 a, zn128 b) noexcept
    {
        return a + (-b);
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr bool operator==(zn128 a, zn128 b) noexcept
    {
        return a.lo_ == b.lo_ && a.hi_ == b.hi_;
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr bool operator!=(zn128 a, zn128 b) noexcept
    {
        return !(a == b);
    }

  private:
    std::uint64_t lo_{};
    std::uint64_t hi_{};

    HEDLEY_ALWAYS_INLINE
    static constexpr zn128 from_reduced(std::uint64_t lo, std::uint64_t hi) noexcept
    {
        zn128 out;
        out.lo_ = lo;
        out.hi_ = hi;
        return out;
    }

    HEDLEY_ALWAYS_INLINE
    static constexpr bool less(std::uint64_t ah, std::uint64_t al,
        std::uint64_t bh, std::uint64_t bl) noexcept
    {
        return ah < bh || (ah == bh && al < bl);
    }

    HEDLEY_ALWAYS_INLINE
    static constexpr void sub_words(std::uint64_t al, std::uint64_t ah,
        std::uint64_t bl, std::uint64_t bh,
        std::uint64_t & ol, std::uint64_t & oh) noexcept
    {
        const unsigned borrow = al < bl ? 1u : 0u;
        ol = static_cast<std::uint64_t>(al - bl);
        oh = static_cast<std::uint64_t>(ah - bh - borrow);
    }

    HEDLEY_ALWAYS_INLINE
    static constexpr zn128 reduce(std::uint64_t lo, std::uint64_t hi) noexcept
    {
        const unsigned __int128 mod =
            static_cast<unsigned __int128>(Lo)
            | (static_cast<unsigned __int128>(Hi) << 64);
        const unsigned __int128 v =
            static_cast<unsigned __int128>(lo)
            | (static_cast<unsigned __int128>(hi) << 64);
        const unsigned __int128 r = v % mod;
        return from_reduced(static_cast<std::uint64_t>(r),
            static_cast<std::uint64_t>(r >> 64));
    }

    template <typename T>
    HEDLEY_ALWAYS_INLINE
    constexpr void assign_integer(T v) noexcept
    {
        bool neg = false;
        unsigned __int128 mag = 0;
        if constexpr (std::is_signed_v<T>)
        {
            if (v < 0)
            {
                neg = true;
                using U = std::make_unsigned_t<T>;
                mag = static_cast<U>(0) - static_cast<U>(v);
            }
            else
            {
                mag = static_cast<std::make_unsigned_t<T>>(v);
            }
        }
        else
        {
            mag = static_cast<unsigned __int128>(v);
        }
        *this = from_u128(mag);
        if (neg)
            *this = -*this;
    }
};

/// @brief Reduce a `zn64<Mod>` share into `zn64<Factor>` when `Factor | Mod`.
/// \complexity One `raw()` read and a `zn64<Factor>` construction (one remainder). `Θ(1)`.
/// The `static_assert` requires `Factor` divides `Mod` (and, for `zn128`, that the modulus is a 64-bit multiple of `Factor`).
/// @see grotto::ring_switch_eval
/// @see grotto::zn64
/// @param share a residue already reduced modulo `Mod`
/// @return the same integer modulo `Factor`
template <std::uint64_t Factor, std::uint64_t Mod>
HEDLEY_ALWAYS_INLINE
constexpr zn64<Factor> ring_switch_factor(zn64<Mod> share) noexcept
{
    static_assert(Mod % Factor == 0, "ring_switch_factor: Factor must divide Mod");
    return zn64<Factor>{share.raw()};
}

/// @brief Reduce a `zn128` share into a 64-bit factor when the factor divides.
/// \complexity One `raw()` read and a `zn64<Factor>` construction (one remainder). `Θ(1)`.
/// The `static_assert` requires `Factor` divides `Mod` (and, for `zn128`, that the modulus is a 64-bit multiple of `Factor`).
/// @see grotto::ring_switch_eval
/// @see grotto::zn64
/// @param share a residue already reduced modulo `Mod`
/// @return the same integer modulo `Factor`
template <std::uint64_t Factor, std::uint64_t Lo, std::uint64_t Hi>
HEDLEY_ALWAYS_INLINE
constexpr zn64<Factor> ring_switch_factor(zn128<Lo, Hi> share) noexcept
{
    static_assert(Hi == 0 && Lo % Factor == 0,
        "ring_switch_factor: Factor must divide the 128-bit modulus");
    return zn64<Factor>{share.lo()};
}

}  // namespace grotto

#endif  // LIBDPF_INCLUDE_GROTTO_RESIDUE_HPP__