/// @file dpf/field64.hpp
/// @brief Prime field GF(2^64 − 2^32 + 1), as a DPF output.
/// @details The modulus is libprio's `Field64` (the Goldilocks prime). Pass
///          `dpf::field64{n}` as a `make_dpf` payload. Leaf addition and
///          scaling are the field operations. A raw PRG block is reduced into
///          the field on the first leaf operation. This is a point-function
///          output, not a comparison payload.
/// @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_FIELD64_HPP__
#define LIBDPF_INCLUDE_DPF_FIELD64_HPP__

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

#include "hedley/hedley.h"
#include "simde/simde/x86/avx2.h"

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

namespace dpf
{

/// @brief Element of GF(2^64 − 2^32 + 1).
class field64
{
  public:
    /// @brief Modulus \f$2^{64}-2^{32}+1\f$.
    static constexpr std::uint64_t mod = 0xffffffff00000001ull;
    using integral_type = std::uint64_t;
    static constexpr bool dpf_point_group = true;

    /// @brief Reduce `v` into the field. A negative value is negated in the field.
    /// @tparam T integral type
    /// @param v the integer to reduce
    template <typename T, typename = std::enable_if_t<std::is_integral_v<T>>>
    HEDLEY_ALWAYS_INLINE
    constexpr field64(T v) noexcept
    {
        val = from_integer(v);
    }

    /// @brief The zero element.
    HEDLEY_ALWAYS_INLINE
    constexpr field64() noexcept = default;

    /// @brief Reduce a PRG block into the field. Uses up to 16 bytes.
    /// @param bytes the PRG output
    /// @param n the number of bytes available
    /// @return the field element
    HEDLEY_ALWAYS_INLINE
    static field64 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(reduce_u128(wide));
    }

    /// @brief Reduced representative in `[0, mod)`.
    /// @return the stored field element
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    constexpr integral_type raw() const noexcept { return reduce(val); }

    /// @brief Same value as `raw()`.
    /// @return the stored field element
    HEDLEY_ALWAYS_INLINE
    HEDLEY_PURE
    explicit constexpr operator integral_type() const noexcept { return raw(); }

    /// @brief Fold a 128-bit integer into the field.
    /// @param x the integer to reduce
    /// @return `x` modulo `mod`
    HEDLEY_ALWAYS_INLINE
    static constexpr integral_type reduce_u128(unsigned __int128 x) noexcept
    {
        // 2^64 ≡ 2^32 − 1, so each step replaces the high half.
        while (x >> 64)
        {
            const auto lo = static_cast<integral_type>(x);
            const auto hi = static_cast<integral_type>(x >> 64);
            x = static_cast<unsigned __int128>(lo)
                + (static_cast<unsigned __int128>(hi) << 32) - hi;
        }
        auto lo = static_cast<integral_type>(x);
        if (lo >= mod)
            lo -= mod;
        return lo;
    }

    /// @brief Reduce a 64-bit word. `mod` itself fits in the word.
    /// @param x the integer to reduce
    /// @return `x` modulo `mod`
    HEDLEY_ALWAYS_INLINE
    static constexpr integral_type reduce(integral_type x) noexcept
    {
        return x >= mod ? static_cast<integral_type>(x - mod) : x;
    }

    /// @brief Field addition.
    /// @param a left addend
    /// @param b right addend
    /// @return `a + b` in the field
    HEDLEY_ALWAYS_INLINE
    friend constexpr field64 operator+(field64 a, field64 b) noexcept
    {
        return from_reduced(reduce_u128(
            static_cast<unsigned __int128>(reduce(a.val)) + reduce(b.val)));
    }

    /// @brief Field subtraction.
    /// @param a minuend
    /// @param b subtrahend
    /// @return `a - b` in the field
    HEDLEY_ALWAYS_INLINE
    friend constexpr field64 operator-(field64 a, field64 b) noexcept
    {
        a.val = reduce(a.val);
        b.val = reduce(b.val);
        if (a.val >= b.val)
            return from_reduced(static_cast<integral_type>(a.val - b.val));
        return from_reduced(static_cast<integral_type>(mod - (b.val - a.val)));
    }

    /// @brief Field negation.
    /// @param a the element to negate
    /// @return `-a`, with `-0 = 0`
    HEDLEY_ALWAYS_INLINE
    friend constexpr field64 operator-(field64 a) noexcept
    {
        a.val = reduce(a.val);
        return from_reduced(a.val == 0 ? 0 : static_cast<integral_type>(mod - a.val));
    }

    /// @brief Field multiplication.
    /// @param a left factor
    /// @param b right factor
    /// @return `a * b` in the field
    HEDLEY_ALWAYS_INLINE
    friend constexpr field64 operator*(field64 a, field64 b) noexcept
    {
        return from_reduced(reduce_u128(
            static_cast<unsigned __int128>(reduce(a.val)) * reduce(b.val)));
    }

    /// @brief Field equality.
    /// @param a left element
    /// @param b right element
    /// @return `true` when the reduced values match
    HEDLEY_ALWAYS_INLINE
    friend constexpr bool operator==(field64 a, field64 b) noexcept
    {
        return reduce(a.val) == reduce(b.val);
    }

    /// @brief Field inequality.
    /// @param a left element
    /// @param b right element
    /// @return `true` when the reduced values differ
    HEDLEY_ALWAYS_INLINE
    friend constexpr bool operator!=(field64 a, field64 b) noexcept
    {
        return reduce(a.val) != reduce(b.val);
    }

    /// @brief Write the reduced representative in decimal.
    /// @param os the output stream
    /// @param a the element to write
    /// @return `os`
    friend std::ostream & operator<<(std::ostream & os, field64 a)
    {
        return os << a.raw();
    }

  private:
    integral_type val{};

    HEDLEY_ALWAYS_INLINE
    static constexpr field64 from_reduced(integral_type v) noexcept
    {
        field64 out;
        out.val = v;
        return out;
    }

    template <typename T>
    HEDLEY_ALWAYS_INLINE
    static constexpr integral_type from_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);
                return reduce_u128(static_cast<unsigned __int128>(mag)) == 0
                    ? 0
                    : static_cast<integral_type>(
                        mod - reduce_u128(static_cast<unsigned __int128>(mag)));
            }
        }
        return reduce_u128(static_cast<unsigned __int128>(v));
    }
};

namespace utils
{

template <>
struct bitlength_of<field64>
    : std::integral_constant<std::size_t, 64>
{ };

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

}  // namespace utils

namespace leaf_arithmetic
{

namespace detail
{

template <std::size_t Lanes>
HEDLEY_ALWAYS_INLINE
void field64_lanes(const void *a, const void *b, void *out,
    field64 (*op)(field64, field64)) noexcept
{
    std::uint64_t aa[Lanes], bb[Lanes], cc[Lanes];
    std::memcpy(aa, a, sizeof(aa));
    std::memcpy(bb, b, sizeof(bb));
    for (std::size_t i = 0; i < Lanes; ++i)
        cc[i] = op(field64{aa[i]}, field64{bb[i]}).raw();
    std::memcpy(out, cc, sizeof(cc));
}

template <std::size_t Lanes>
HEDLEY_ALWAYS_INLINE
void field64_scale(const void *a, field64 b, void *out) noexcept
{
    std::uint64_t aa[Lanes], cc[Lanes];
    std::memcpy(aa, a, sizeof(aa));
    for (std::size_t i = 0; i < Lanes; ++i)
        cc[i] = (field64{aa[i]} * b).raw();
    std::memcpy(out, cc, sizeof(cc));
}

}  // namespace detail

HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")

template <>
struct add_t<field64, simde__m128i>
{
    auto operator()(const simde__m128i &a, const simde__m128i &b) const
    {
        simde__m128i out;
        detail::field64_lanes<2>(&a, &b, &out, [](field64 x, field64 y) {
            return x + y;
        });
        return out;
    }
};

template <>
struct subtract_t<field64, simde__m128i>
{
    auto operator()(const simde__m128i &a, const simde__m128i &b) const
    {
        simde__m128i out;
        detail::field64_lanes<2>(&a, &b, &out, [](field64 x, field64 y) {
            return x - y;
        });
        return out;
    }
};

template <>
struct multiply_t<field64, simde__m128i>
{
    auto operator()(const simde__m128i &a, field64 b) const
    {
        simde__m128i out;
        detail::field64_scale<2>(&a, b, &out);
        return out;
    }
};

template <>
struct add_t<field64, simde__m256i>
{
    auto operator()(const simde__m256i &a, const simde__m256i &b) const
    {
        simde__m256i out;
        detail::field64_lanes<4>(&a, &b, &out, [](field64 x, field64 y) {
            return x + y;
        });
        return out;
    }
};

template <>
struct subtract_t<field64, simde__m256i>
{
    auto operator()(const simde__m256i &a, const simde__m256i &b) const
    {
        simde__m256i out;
        detail::field64_lanes<4>(&a, &b, &out, [](field64 x, field64 y) {
            return x - y;
        });
        return out;
    }
};

template <>
struct multiply_t<field64, simde__m256i>
{
    auto operator()(const simde__m256i &a, field64 b) const
    {
        simde__m256i out;
        detail::field64_scale<4>(&a, b, &out);
        return out;
    }
};

HEDLEY_PRAGMA(GCC diagnostic pop)

}  // namespace leaf_arithmetic

/// @brief Sample a uniform field element by rejection.
/// @return an element of the field
template <>
HEDLEY_NO_THROW
inline auto uniform_sample<field64>() noexcept
{
    for (;;)
    {
        const auto x = uniform_sample<std::uint64_t>();
        if (x < field64::mod)
            return field64{x};
    }
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_FIELD64_HPP__