/// @file dpf/p256_scalar.hpp
/// @brief NIST P-256 scalar field as a comparison payload group.
/// @details Integers modulo the curve order
///          `0xffffffff00000000ffffffffffffffffbce6faada7179e84f3b9cac2fc632551`.
///          This is the scalar field, not the point group in `p256.hpp`.
///          `from_seed`, `+`, and unary `-` select the payload-group path.
/// @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_DPF_P256_SCALAR_HPP__
#define LIBDPF_INCLUDE_DPF_P256_SCALAR_HPP__

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

#include "hedley/hedley.h"

#include "dpf/random.hpp"
#include "dpf/utils.hpp"

namespace dpf
{

/// @brief Element of the NIST P-256 scalar field.
class p256_scalar
{
  public:
    /// @brief Little-endian limbs of the group order \f$n\f$.
    static constexpr std::uint64_t order[4] = {
        0xf3b9cac2fc632551ull,
        0xbce6faada7179e84ull,
        0xffffffffffffffffull,
        0xffffffff00000000ull,
    };
    static constexpr bool dpf_point_group = true;

    HEDLEY_ALWAYS_INLINE
    constexpr p256_scalar() noexcept = default;

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

    HEDLEY_ALWAYS_INLINE
    constexpr p256_scalar(unsigned __int128 v) noexcept
    {
        std::uint64_t z[4] = {
            static_cast<std::uint64_t>(v),
            static_cast<std::uint64_t>(v >> 64),
            0, 0};
        if (ge_order(z))
            sub_order(z);
        d_[0] = z[0];
        d_[1] = z[1];
        d_[2] = z[2];
        d_[3] = z[3];
    }

    /// @brief Canonical representative in `[0, n)`.
    HEDLEY_ALWAYS_INLINE
    static constexpr p256_scalar canonicalize(p256_scalar a) noexcept
    {
        std::uint64_t z[4] = {a.d_[0], a.d_[1], a.d_[2], a.d_[3]};
        if (ge_order(z))
            sub_order(z);
        if (ge_order(z))
            sub_order(z);
        return from_limbs(z);
    }

    /// @brief Stretch a PRG block to ≥256 bits and rejection-sample into `[0, n)`.
    HEDLEY_ALWAYS_INLINE
    static p256_scalar from_seed(const void * bytes, std::size_t n) noexcept
    {
        for (std::uint32_t counter = 0; ; ++counter)
        {
            class SHA256 h;
            h.add(bytes, n);
            const unsigned char ctr[4] = {
                static_cast<unsigned char>(counter),
                static_cast<unsigned char>(counter >> 8),
                static_cast<unsigned char>(counter >> 16),
                static_cast<unsigned char>(counter >> 24)};
            h.add(ctr, sizeof(ctr));
            unsigned char block[SHA256::HashBytes];
            h.getHash(block);
            std::uint64_t w[4]{};
            std::memcpy(w, block, sizeof(w));
            if (!ge_order(w))
                return from_limbs(w);
        }
    }

    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr std::uint64_t limb(std::size_t i) const noexcept { return d_[i]; }

    HEDLEY_ALWAYS_INLINE
    friend constexpr p256_scalar operator+(p256_scalar a, p256_scalar b) noexcept
    {
        a = canonicalize(a);
        b = canonicalize(b);
        std::uint64_t z[4]{};
        unsigned __int128 carry = 0;
        for (std::size_t i = 0; i < 4; ++i)
        {
            carry += static_cast<unsigned __int128>(a.d_[i]) + b.d_[i];
            z[i] = static_cast<std::uint64_t>(carry);
            carry >>= 64;
        }
        if (carry != 0 || ge_order(z))
            sub_order(z);
        if (ge_order(z))
            sub_order(z);
        return from_limbs(z);
    }

    HEDLEY_ALWAYS_INLINE
    friend constexpr p256_scalar operator-(p256_scalar a) noexcept
    {
        a = canonicalize(a);
        if (is_zero(a))
            return a;
        std::uint64_t z[4]{};
        unsigned borrow = 0;
        for (std::size_t i = 0; i < 4; ++i)
        {
            const unsigned __int128 diff =
                static_cast<unsigned __int128>(order[i]) - a.d_[i] - borrow;
            z[i] = static_cast<std::uint64_t>(diff);
            borrow = (diff >> 64) ? 1u : 0u;
        }
        return from_limbs(z);
    }

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

    HEDLEY_ALWAYS_INLINE
    friend constexpr bool operator==(p256_scalar a, p256_scalar b) noexcept
    {
        a = canonicalize(a);
        b = canonicalize(b);
        return a.d_[0] == b.d_[0] && a.d_[1] == b.d_[1]
            && a.d_[2] == b.d_[2] && a.d_[3] == b.d_[3];
    }

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

  private:
    std::uint64_t d_[4]{};

    HEDLEY_ALWAYS_INLINE
    static constexpr bool is_zero(p256_scalar a) noexcept
    {
        return (a.d_[0] | a.d_[1] | a.d_[2] | a.d_[3]) == 0;
    }

    HEDLEY_ALWAYS_INLINE
    static constexpr bool ge_order(const std::uint64_t z[4]) noexcept
    {
        for (std::size_t i = 4; i-- > 0; )
        {
            if (z[i] > order[i])
                return true;
            if (z[i] < order[i])
                return false;
        }
        return true;
    }

    HEDLEY_ALWAYS_INLINE
    static constexpr void sub_order(std::uint64_t z[4]) noexcept
    {
        unsigned borrow = 0;
        for (std::size_t i = 0; i < 4; ++i)
        {
            const unsigned __int128 diff =
                static_cast<unsigned __int128>(z[i]) - order[i] - borrow;
            z[i] = static_cast<std::uint64_t>(diff);
            borrow = (diff >> 64) ? 1u : 0u;
        }
    }

    HEDLEY_ALWAYS_INLINE
    static constexpr p256_scalar from_limbs(const std::uint64_t z[4]) noexcept
    {
        p256_scalar out;
        out.d_[0] = z[0];
        out.d_[1] = z[1];
        out.d_[2] = z[2];
        out.d_[3] = z[3];
        return out;
    }

    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);
        }
        std::uint64_t z[4] = {
            static_cast<std::uint64_t>(mag),
            static_cast<std::uint64_t>(mag >> 64),
            0, 0};
        if (ge_order(z))
            sub_order(z);
        *this = from_limbs(z);
        if (neg)
            *this = -*this;
    }
};

namespace utils
{

template <>
struct bitlength_of<p256_scalar>
    : std::integral_constant<std::size_t, 256>
{ };

template <>
struct has_characteristic_two<p256_scalar> : std::false_type
{ };

}  // namespace utils

/// @brief Sample a uniform scalar by rejection into `[0, n)`.
template <>
HEDLEY_NO_THROW
inline auto uniform_sample<p256_scalar>() noexcept
{
    for (;;)
    {
        std::uint64_t z[4] = {
            uniform_sample<std::uint64_t>(),
            uniform_sample<std::uint64_t>(),
            uniform_sample<std::uint64_t>(),
            uniform_sample<std::uint64_t>(),
        };
        bool ge = true;
        for (std::size_t i = 4; i-- > 0; )
        {
            if (z[i] > p256_scalar::order[i])
            {
                ge = true;
                break;
            }
            if (z[i] < p256_scalar::order[i])
            {
                ge = false;
                break;
            }
        }
        if (!ge)
        {
            p256_scalar out;
            std::memcpy(&out, z, sizeof(z));
            return out;
        }
    }
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_P256_SCALAR_HPP__