/// @file dpf/p256.hpp
/// @brief NIST P-256 points as a DPF output group.
/// @details A point is 33 bytes in SEC1 compressed form. The identity is 33
///          zero bytes. `p256{k}` is `k` times the generator, and
///          `p256::from_compressed` reads a 33-byte encoding. Leaf addition is
///          point addition on valid encodings. A PRG block is stretched to at
///          least 256 bits and hashed onto the curve via `from_seed`. There is
///          no point×point product. This is a point-function output; it also
///          supplies `from_seed` for comparison payloads.
/// @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_P256_HPP__
#define LIBDPF_INCLUDE_DPF_P256_HPP__

#include <array>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <ostream>
#include <stdexcept>
#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
{

namespace p256_detail
{

struct fe
{
    std::uint64_t d[4]{};
};

inline constexpr std::uint64_t P[4] = {
    0xffffffffffffffffull,
    0x00000000ffffffffull,
    0x0000000000000000ull,
    0xffffffff00000001ull,
};

inline constexpr std::uint64_t B[4] = {
    0x3bce3c3e27d2604bull,
    0x651d06b0cc53b0f6ull,
    0xb3ebbd55769886bcull,
    0x5ac635d8aa3a93e7ull,
};

inline constexpr std::uint64_t GX[4] = {
    0xf4a13945d898c296ull,
    0x77037d812deb33a0ull,
    0xf8bce6e563a440f2ull,
    0x6b17d1f2e12c4247ull,
};

inline constexpr std::uint64_t GY[4] = {
    0xcbb6406837bf51f5ull,
    0x2bce33576b315eceull,
    0x8ee7eb4a7c0f9e16ull,
    0x4fe342e2fe1a7f9bull,
};

inline constexpr std::uint64_t N[4] = {
    0xf3b9cac2fc632551ull,
    0xbce6faada7179e84ull,
    0xffffffffffffffffull,
    0xffffffff00000000ull,
};

inline constexpr std::uint64_t EXP_P_PLUS_1_OVER_4[4] = {
    0x0000000000000000ull,
    0x0000000040000000ull,
    0x4000000000000000ull,
    0x3fffffffc0000000ull,
};

inline constexpr std::uint64_t EXP_P_MINUS_1_OVER_2[4] = {
    0xffffffffffffffffull,
    0x000000007fffffffull,
    0x8000000000000000ull,
    0x7fffffff80000000ull,
};

inline constexpr std::uint64_t EXP_P_MINUS_2[4] = {
    0xfffffffffffffffdull,
    0x00000000ffffffffull,
    0x0000000000000000ull,
    0xffffffff00000001ull,
};

struct affine
{
    fe x{};
    fe y{};
    bool inf = true;
};

HEDLEY_ALWAYS_INLINE
int fe_cmp(const fe &a, const fe &b) noexcept
{
    for (int i = 3; i >= 0; --i)
    {
        if (a.d[i] < b.d[i])
            return -1;
        if (a.d[i] > b.d[i])
            return 1;
    }
    return 0;
}

HEDLEY_ALWAYS_INLINE
int limbs_cmp(const std::uint64_t *a, const std::uint64_t *b, int n) noexcept
{
    for (int i = n - 1; i >= 0; --i)
    {
        if (a[i] < b[i])
            return -1;
        if (a[i] > b[i])
            return 1;
    }
    return 0;
}

HEDLEY_ALWAYS_INLINE
bool fe_is_zero(const fe &a) noexcept
{
    return (a.d[0] | a.d[1] | a.d[2] | a.d[3]) == 0;
}

/// @brief Conditionally subtract `P` when `a` is in `[P, 2P)`.
HEDLEY_ALWAYS_INLINE
fe fe_sub_p_if_ge(fe a) noexcept
{
    if (fe_cmp(a, fe{{P[0], P[1], P[2], P[3]}}) >= 0)
    {
        unsigned borrow = 0;
        for (int i = 0; i < 4; ++i)
        {
            const unsigned __int128 need =
                static_cast<unsigned __int128>(P[i]) + borrow;
            const unsigned __int128 cur = a.d[i];
            if (cur >= need)
            {
                a.d[i] = static_cast<std::uint64_t>(cur - need);
                borrow = 0;
            }
            else
            {
                a.d[i] = static_cast<std::uint64_t>(
                    (cur + (static_cast<unsigned __int128>(1) << 64)) - need);
                borrow = 1;
            }
        }
    }
    return a;
}

HEDLEY_ALWAYS_INLINE
void u512_add_u64(std::uint64_t x[8], int limb, std::uint64_t v) noexcept
{
    unsigned __int128 sum =
        static_cast<unsigned __int128>(x[limb]) + v;
    x[limb] = static_cast<std::uint64_t>(sum);
    unsigned __int128 carry = sum >> 64;
    for (int i = limb + 1; carry && i < 8; ++i)
    {
        sum = static_cast<unsigned __int128>(x[i]) + carry;
        x[i] = static_cast<std::uint64_t>(sum);
        carry = sum >> 64;
    }
}

HEDLEY_ALWAYS_INLINE
void u512_add_shifted(std::uint64_t x[8], const std::uint64_t h[4],
    int bit_shift) noexcept
{
    const int limb = bit_shift >> 6;
    const int bit = bit_shift & 63;
    for (int i = 0; i < 4; ++i)
    {
        if (bit == 0)
            u512_add_u64(x, limb + i, h[i]);
        else
        {
            u512_add_u64(x, limb + i,
                static_cast<std::uint64_t>(h[i] << bit));
            if (limb + i + 1 < 8)
                u512_add_u64(x, limb + i + 1, h[i] >> (64 - bit));
        }
    }
}

HEDLEY_ALWAYS_INLINE
bool u512_sub(std::uint64_t x[8], const std::uint64_t y[8]) noexcept
{
    unsigned borrow = 0;
    for (int i = 0; i < 8; ++i)
    {
        const unsigned __int128 need =
            static_cast<unsigned __int128>(y[i]) + borrow;
        const unsigned __int128 cur = x[i];
        if (cur >= need)
        {
            x[i] = static_cast<std::uint64_t>(cur - need);
            borrow = 0;
        }
        else
        {
            x[i] = static_cast<std::uint64_t>(
                (cur + (static_cast<unsigned __int128>(1) << 64)) - need);
            borrow = 1;
        }
    }
    return borrow != 0;
}

HEDLEY_ALWAYS_INLINE
void u512_add_p_at(std::uint64_t x[8], int limb) noexcept
{
    unsigned __int128 carry = 0;
    for (int i = 0; i < 4; ++i)
    {
        carry += static_cast<unsigned __int128>(x[limb + i]) + P[i];
        x[limb + i] = static_cast<std::uint64_t>(carry);
        carry >>= 64;
    }
    for (int i = limb + 4; carry && i < 8; ++i)
    {
        carry += x[i];
        x[i] = static_cast<std::uint64_t>(carry);
        carry >>= 64;
    }
}

/// @brief Fast NIST P-256 reduction of a 512-bit little-endian product.
/// @details `2^256 ≡ 2^224 - 2^192 - 2^96 + 1 (mod p)`. Shifts are applied
///          as full `__int128` limb writes so carry propagation, not a second
///          half-limb store, moves the high bits.
inline fe fe_reduce_wide(const std::uint64_t z[8]) noexcept
{
    __int128 t[10]{};
    auto add = [&](int i, __int128 v) { t[i] += v; };
    for (int i = 0; i < 4; ++i)
    {
        add(i, static_cast<__int128>(z[i]));
        add(i, static_cast<__int128>(z[4 + i]));
    }
    for (int i = 0; i < 4; ++i)
    {
        const __int128 h = static_cast<__int128>(z[4 + i]);
        // +H << 224 and -H << 192 and -H << 96 (carry moves overflow).
        add(i + 3, h << 32);
        add(i + 3, -h);
        add(i + 1, -(h << 32));
    }
    for (int pass = 0; pass < 6; ++pass)
    {
        for (int i = 0; i < 9; ++i)
        {
            const __int128 c = t[i] >> 64;
            t[i] = static_cast<std::uint64_t>(t[i]);
            t[i + 1] += c;
        }
        if (t[4] == 0 && t[5] == 0 && t[6] == 0 && t[7] == 0
            && t[8] == 0 && t[9] == 0)
            break;
        __int128 high[6] = {t[4], t[5], t[6], t[7], t[8], t[9]};
        t[4] = t[5] = t[6] = t[7] = t[8] = t[9] = 0;
        for (int i = 0; i < 6; ++i)
        {
            const __int128 h = high[i];
            if (h == 0)
                continue;
            // h * 2^{256+64*i} ≡ h*2^{64*i} * (1+2^224-2^192-2^96)
            add(i, h);
            add(i + 3, h << 32);
            add(i + 3, -h);
            add(i + 1, -(h << 32));
        }
    }
    for (int i = 0; i < 4; ++i)
    {
        const __int128 c = t[i] >> 64;
        t[i] = static_cast<std::uint64_t>(t[i]);
        t[i + 1] += c;
    }
    for (int guard = 0; guard < 6; ++guard)
    {
        if (t[4] == 0 && t[0] >= 0 && t[1] >= 0 && t[2] >= 0 && t[3] >= 0)
            break;
        if (t[4] != 0 || t[5] != 0)
        {
            __int128 high[2] = {t[4], t[5]};
            t[4] = t[5] = 0;
            for (int i = 0; i < 2; ++i)
            {
                const __int128 h = high[i];
                if (h == 0)
                    continue;
                add(i, h);
                add(i + 3, h << 32);
                add(i + 3, -h);
                add(i + 1, -(h << 32));
            }
            for (int i = 0; i < 4; ++i)
            {
                const __int128 c = t[i] >> 64;
                t[i] = static_cast<std::uint64_t>(t[i]);
                t[i + 1] += c;
            }
            continue;
        }
        add(0, static_cast<__int128>(P[0]));
        add(1, static_cast<__int128>(P[1]));
        add(2, static_cast<__int128>(P[2]));
        add(3, static_cast<__int128>(P[3]));
        for (int i = 0; i < 4; ++i)
        {
            const __int128 c = t[i] >> 64;
            t[i] = static_cast<std::uint64_t>(t[i]);
            t[i + 1] += c;
        }
    }
    fe out{{static_cast<std::uint64_t>(t[0]),
        static_cast<std::uint64_t>(t[1]),
        static_cast<std::uint64_t>(t[2]),
        static_cast<std::uint64_t>(t[3])}};
    out = fe_sub_p_if_ge(out);
    out = fe_sub_p_if_ge(out);
    return out;
}

HEDLEY_ALWAYS_INLINE
fe fe_add(fe a, fe b) noexcept
{
    unsigned __int128 carry = 0;
    fe r{};
    for (int i = 0; i < 4; ++i)
    {
        carry += static_cast<unsigned __int128>(a.d[i]) + b.d[i];
        r.d[i] = static_cast<std::uint64_t>(carry);
        carry >>= 64;
    }
    if (carry)
    {
        std::uint64_t z[8] = {r.d[0], r.d[1], r.d[2], r.d[3], 1, 0, 0, 0};
        return fe_reduce_wide(z);
    }
    return fe_sub_p_if_ge(r);
}

HEDLEY_ALWAYS_INLINE
fe fe_neg(fe a) noexcept
{
    if (fe_is_zero(a))
        return a;
    fe r{};
    unsigned borrow = 0;
    for (int i = 0; i < 4; ++i)
    {
        const unsigned __int128 need =
            static_cast<unsigned __int128>(a.d[i]) + borrow;
        const unsigned __int128 cur = P[i];
        if (cur >= need)
        {
            r.d[i] = static_cast<std::uint64_t>(cur - need);
            borrow = 0;
        }
        else
        {
            r.d[i] = static_cast<std::uint64_t>(
                (cur + (static_cast<unsigned __int128>(1) << 64)) - need);
            borrow = 1;
        }
    }
    return r;
}

HEDLEY_ALWAYS_INLINE
fe fe_sub(fe a, fe b) noexcept
{
    fe r{};
    unsigned borrow = 0;
    for (int i = 0; i < 4; ++i)
    {
        const unsigned __int128 need =
            static_cast<unsigned __int128>(b.d[i]) + borrow;
        const unsigned __int128 cur = a.d[i];
        if (cur >= need)
        {
            r.d[i] = static_cast<std::uint64_t>(cur - need);
            borrow = 0;
        }
        else
        {
            r.d[i] = static_cast<std::uint64_t>(
                (cur + (static_cast<unsigned __int128>(1) << 64)) - need);
            borrow = 1;
        }
    }
    if (borrow)
    {
        unsigned __int128 carry = 0;
        for (int i = 0; i < 4; ++i)
        {
            carry += static_cast<unsigned __int128>(r.d[i]) + P[i];
            r.d[i] = static_cast<std::uint64_t>(carry);
            carry >>= 64;
        }
    }
    return r;
}

HEDLEY_ALWAYS_INLINE
fe fe_mul(fe a, fe b) noexcept
{
    std::uint64_t z[8]{};
    for (int i = 0; i < 4; ++i)
    {
        unsigned __int128 carry = 0;
        for (int j = 0; j < 4; ++j)
        {
            const unsigned __int128 cur =
                static_cast<unsigned __int128>(z[i + j])
                + static_cast<unsigned __int128>(a.d[i]) * b.d[j]
                + carry;
            z[i + j] = static_cast<std::uint64_t>(cur);
            carry = cur >> 64;
        }
        z[i + 4] = static_cast<std::uint64_t>(carry);
    }
    return fe_reduce_wide(z);
}

HEDLEY_ALWAYS_INLINE
fe fe_pow(fe a, const std::uint64_t exp[4]) noexcept
{
    fe result{{1, 0, 0, 0}};
    for (int bit = 255; bit >= 0; --bit)
    {
        result = fe_mul(result, result);
        const unsigned limb = static_cast<unsigned>(bit >> 6);
        const unsigned off = static_cast<unsigned>(bit & 63);
        if ((exp[limb] >> off) & 1u)
            result = fe_mul(result, a);
    }
    return result;
}

HEDLEY_ALWAYS_INLINE
bool fe_is_square(fe a) noexcept
{
    if (fe_is_zero(a))
        return true;
    const fe leg = fe_pow(a, EXP_P_MINUS_1_OVER_2);
    return leg.d[0] == 1 && leg.d[1] == 0 && leg.d[2] == 0 && leg.d[3] == 0;
}

HEDLEY_ALWAYS_INLINE
fe fe_sqrt(fe a) noexcept
{
    return fe_pow(a, EXP_P_PLUS_1_OVER_4);
}

HEDLEY_ALWAYS_INLINE
fe fe_inv(fe a) noexcept
{
    return fe_pow(a, EXP_P_MINUS_2);
}

HEDLEY_ALWAYS_INLINE
fe curve_rhs(fe x) noexcept
{
    const fe x2 = fe_mul(x, x);
    const fe x3 = fe_mul(x2, x);
    const fe three = fe_add(x, fe_add(x, x));
    fe rhs = fe_sub(x3, three);
    fe b{{B[0], B[1], B[2], B[3]}};
    return fe_add(rhs, b);
}

HEDLEY_ALWAYS_INLINE
void fe_from_be(fe &out, const std::uint8_t be[32]) noexcept
{
    for (int i = 0; i < 4; ++i)
    {
        std::uint64_t v = 0;
        for (int j = 0; j < 8; ++j)
            v = (v << 8) | be[i * 8 + j];
        out.d[3 - i] = v;
    }
}

HEDLEY_ALWAYS_INLINE
void fe_to_be(std::uint8_t be[32], const fe &a) noexcept
{
    for (int i = 0; i < 4; ++i)
    {
        const std::uint64_t v = a.d[3 - i];
        for (int j = 0; j < 8; ++j)
            be[i * 8 + j] = static_cast<std::uint8_t>(v >> (56 - 8 * j));
    }
}

HEDLEY_ALWAYS_INLINE
affine generator_point() noexcept
{
    affine g;
    g.inf = false;
    for (int i = 0; i < 4; ++i)
    {
        g.x.d[i] = GX[i];
        g.y.d[i] = GY[i];
    }
    return g;
}

HEDLEY_ALWAYS_INLINE
affine point_neg(affine a) noexcept
{
    if (!a.inf)
        a.y = fe_neg(a.y);
    return a;
}

HEDLEY_ALWAYS_INLINE
fe fe_sqr(fe a) noexcept
{
    return fe_mul(a, a);
}

/// @brief Jacobian coordinates `(X : Y : Z)` for `x=X/Z^2`, `y=Y/Z^3`.
struct jacobian
{
    fe X{};
    fe Y{};
    fe Z{};
};

HEDLEY_ALWAYS_INLINE
bool jac_is_inf(const jacobian &p) noexcept
{
    return fe_is_zero(p.Z);
}

HEDLEY_ALWAYS_INLINE
jacobian jac_from_affine(affine a) noexcept
{
    if (a.inf)
        return {};
    return jacobian{a.x, a.y, fe{{1, 0, 0, 0}}};
}

HEDLEY_ALWAYS_INLINE
affine jac_to_affine(jacobian p) noexcept
{
    if (jac_is_inf(p))
        return {};
    const fe zi = fe_inv(p.Z);
    const fe zi2 = fe_sqr(zi);
    const fe zi3 = fe_mul(zi2, zi);
    affine r;
    r.inf = false;
    r.x = fe_mul(p.X, zi2);
    r.y = fe_mul(p.Y, zi3);
    return r;
}

/// @brief Jacobian doubling for short-Weierstrass with `a = -3` (P-256).
HEDLEY_ALWAYS_INLINE
jacobian jac_double(jacobian p) noexcept
{
    if (jac_is_inf(p) || fe_is_zero(p.Y))
        return {};
    const fe yy = fe_sqr(p.Y);
    const fe yyyy = fe_sqr(yy);
    const fe zz = fe_sqr(p.Z);
    // S = 4*X*Y^2
    const fe s = fe_mul(fe_add(p.X, p.X), fe_add(yy, yy));
    // M = 3*(X+Z^2)*(X-Z^2)
    const fe xz = fe_add(p.X, zz);
    const fe xd = fe_sub(p.X, zz);
    const fe m = fe_mul(fe_add(xd, fe_add(xd, xd)), xz);
    const fe m2 = fe_sqr(m);
    jacobian r;
    r.X = fe_sub(m2, fe_add(s, s));
    // 8*Y^4
    fe eight_yyyy = fe_add(yyyy, yyyy);
    eight_yyyy = fe_add(eight_yyyy, eight_yyyy);
    eight_yyyy = fe_add(eight_yyyy, eight_yyyy);
    r.Y = fe_sub(fe_mul(m, fe_sub(s, r.X)), eight_yyyy);
    // Z3 = 2*Y*Z = (Y+Z)^2 - Y^2 - Z^2
    r.Z = fe_sub(fe_sub(fe_sqr(fe_add(p.Y, p.Z)), yy), zz);
    return r;
}

/// @brief Mixed Jacobian–affine addition.
HEDLEY_ALWAYS_INLINE
jacobian jac_add_mixed(jacobian p, affine q) noexcept
{
    if (q.inf)
        return p;
    if (jac_is_inf(p))
        return jac_from_affine(q);
    const fe z1z1 = fe_sqr(p.Z);
    const fe u2 = fe_mul(q.x, z1z1);
    const fe s2 = fe_mul(q.y, fe_mul(p.Z, z1z1));
    const fe h = fe_sub(u2, p.X);
    const fe r = fe_sub(s2, p.Y);
    if (fe_is_zero(h))
    {
        if (fe_is_zero(r))
            return jac_double(p);
        return {};
    }
    const fe h2 = fe_sqr(h);
    const fe h3 = fe_mul(h2, h);
    const fe v = fe_mul(p.X, h2);
    const fe r2 = fe_sqr(r);
    jacobian out;
    out.X = fe_sub(fe_sub(r2, h3), fe_add(v, v));
    out.Y = fe_sub(fe_mul(r, fe_sub(v, out.X)), fe_mul(p.Y, h3));
    out.Z = fe_mul(h, p.Z);
    return out;
}

HEDLEY_ALWAYS_INLINE
affine point_add(affine a, affine b) noexcept
{
    if (a.inf)
        return b;
    if (b.inf)
        return a;
    // One mixed add + one inversion beats two affine inversions when both
    // inputs are finite; keep the classical path for the equal-x branch.
    if (fe_cmp(a.x, b.x) == 0)
    {
        if (fe_cmp(a.y, b.y) != 0 || fe_is_zero(a.y))
            return {};
        return jac_to_affine(jac_double(jac_from_affine(a)));
    }
    return jac_to_affine(jac_add_mixed(jac_from_affine(a), b));
}

HEDLEY_ALWAYS_INLINE
affine point_sub(affine a, affine b) noexcept
{
    return point_add(a, point_neg(b));
}

HEDLEY_ALWAYS_INLINE
affine point_scalarmul_limbs(affine g, const std::uint64_t k[4]) noexcept
{
    if (g.inf)
        return {};
    // 4-bit window with an affine odd/even table. Doublings stay in Jacobian
    // so the ladder uses one field inversion (final map) instead of one per add.
    affine table[16]{};
    table[1] = g;
    const affine two = jac_to_affine(jac_double(jac_from_affine(g)));
    table[2] = two;
    for (int i = 3; i < 16; ++i)
        table[i] = jac_to_affine(jac_add_mixed(jac_from_affine(table[i - 1]), g));
    jacobian r{};
    bool started = false;
    for (int bit = 255; bit >= 0; bit -= 4)
    {
        if (started)
        {
            r = jac_double(r);
            r = jac_double(r);
            r = jac_double(r);
            r = jac_double(r);
        }
        unsigned nibble = 0;
        for (int j = 0; j < 4; ++j)
        {
            const int b = bit - j;
            if (b < 0)
                break;
            const unsigned limb = static_cast<unsigned>(b >> 6);
            const unsigned off = static_cast<unsigned>(b & 63);
            nibble <<= 1;
            if ((k[limb] >> off) & 1u)
                nibble |= 1u;
        }
        if (nibble != 0)
        {
            r = jac_add_mixed(r, table[nibble]);
            started = true;
        }
    }
    return jac_to_affine(r);
}

HEDLEY_ALWAYS_INLINE
affine point_scalarmul(affine g, unsigned __int128 k) noexcept
{
    std::uint64_t limbs[4] = {
        static_cast<std::uint64_t>(k),
        static_cast<std::uint64_t>(k >> 64),
        0, 0};
    return point_scalarmul_limbs(g, limbs);
}

/// @brief Expand `src` to at least 256 bits via SHA-256(src ‖ counter).
HEDLEY_ALWAYS_INLINE
void stretch_seed(std::uint8_t *dst, std::size_t dst_n,
    const void *src, std::size_t src_n) noexcept
{
    std::size_t off = 0;
    std::uint32_t counter = 0;
    while (off < dst_n)
    {
        class SHA256 h;
        h.add(src, src_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);
        const std::size_t take =
            dst_n - off < sizeof(block) ? dst_n - off : sizeof(block);
        std::memcpy(dst + off, block, take);
        off += take;
        ++counter;
    }
}

HEDLEY_ALWAYS_INLINE
void hash_to_curve(const std::uint8_t *data, std::size_t n, affine &out) noexcept
{
    std::uint8_t be[32]{};
    const std::size_t take = n < 32 ? n : 32;
    std::memcpy(be, data, take);
    for (std::size_t i = 32; i < n; ++i)
        be[i & 31] ^= data[i];
    fe x;
    fe_from_be(x, be);
    fe p{{P[0], P[1], P[2], P[3]}};
    if (fe_cmp(x, p) >= 0)
        x = fe_sub(x, p);
    const fe one{{1, 0, 0, 0}};
    for (;;)
    {
        const fe rhs = curve_rhs(x);
        if (fe_is_square(rhs))
        {
            fe y = fe_sqrt(rhs);
            if ((y.d[0] & 1u) != 0)
                y = fe_neg(y);
            out.inf = false;
            out.x = x;
            out.y = y;
            return;
        }
        x = fe_add(x, one);
    }
}

HEDLEY_ALWAYS_INLINE
bool decode_compressed(const std::uint8_t enc[33], affine &out) noexcept
{
    bool zeros = true;
    for (int i = 0; i < 33; ++i)
        zeros = zeros && enc[i] == 0;
    if (zeros)
    {
        out = affine{};
        return true;
    }
    if (enc[0] != 0x02 && enc[0] != 0x03)
        return false;
    fe x;
    fe_from_be(x, enc + 1);
    fe p{{P[0], P[1], P[2], P[3]}};
    if (fe_cmp(x, p) >= 0)
        return false;
    const fe rhs = curve_rhs(x);
    if (!fe_is_square(rhs))
        return false;
    fe y = fe_sqrt(rhs);
    const bool odd = (y.d[0] & 1u) != 0;
    if (odd != (enc[0] == 0x03))
        y = fe_neg(y);
    if (fe_cmp(fe_mul(y, y), rhs) != 0)
        return false;
    out.inf = false;
    out.x = x;
    out.y = y;
    return true;
}

/// @brief Decode a stored leaf/share encoding. Invalid claimings throw.
HEDLEY_ALWAYS_INLINE
affine decode_strict(const std::uint8_t *bytes, std::size_t n)
{
    std::uint8_t enc[33]{};
    const std::size_t take = n < 33 ? n : 33;
    std::memcpy(enc, bytes, take);
    affine out;
    if (!decode_compressed(enc, out))
        throw std::invalid_argument("p256: not a P-256 point");
    return out;
}

HEDLEY_ALWAYS_INLINE
void encode_point(std::uint8_t enc[33], affine p) noexcept
{
    std::memset(enc, 0, 33);
    if (p.inf)
        return;
    enc[0] = (p.y.d[0] & 1u) ? 0x03 : 0x02;
    fe_to_be(enc + 1, p.x);
}

}  // namespace p256_detail

/// @brief NIST P-256 point, SEC1-compressed in 33 bytes.
class p256
{
  public:
    static constexpr std::size_t encoded_size = 33;
    static constexpr bool dpf_point_group = true;
    /// @brief Leaf PRG bytes are mapped through `from_seed` before storage.
    static constexpr bool dpf_curve_point = true;

    /// @brief The identity (point at infinity).
    HEDLEY_ALWAYS_INLINE
    constexpr p256() noexcept = default;

    /// @brief `scalar` times the generator. A negative scalar negates the point.
    /// @tparam T integral scalar, at most 128 bits
    /// @param scalar the multiplier of the generator
    template <typename T, typename = std::enable_if_t<std::is_integral_v<T>>>
    explicit p256(T scalar)
    {
        bool neg = false;
        unsigned __int128 mag = 0;
        if constexpr (std::is_signed_v<T>)
        {
            if (scalar < 0)
            {
                neg = true;
                using U = std::make_unsigned_t<T>;
                mag = static_cast<U>(0) - static_cast<U>(scalar);
            }
            else
            {
                mag = static_cast<std::make_unsigned_t<T>>(scalar);
            }
        }
        else
        {
            mag = static_cast<unsigned __int128>(scalar);
        }
        auto point = p256_detail::point_scalarmul(
            p256_detail::generator_point(), mag);
        if (neg)
            point = p256_detail::point_neg(point);
        p256_detail::encode_point(enc_, point);
    }

    /// @brief The conventional generator.
    /// @return the generator
    HEDLEY_ALWAYS_INLINE
    static p256 generator()
    {
        p256 g;
        p256_detail::encode_point(g.enc_, p256_detail::generator_point());
        return g;
    }

    /// @brief Decode a 33-byte SEC1 compressed point. 33 zero bytes are the identity.
    /// @param in the compressed encoding
    /// @return the point
    static p256 from_compressed(const std::uint8_t (&in)[encoded_size])
    {
        p256_detail::affine point;
        if (!p256_detail::decode_compressed(in, point))
            throw std::invalid_argument("p256::from_compressed: not a P-256 point");
        p256 out;
        p256_detail::encode_point(out.enc_, point);
        return out;
    }

    /// @brief Decode a 33-byte SEC1 compressed point.
    /// @param in the compressed encoding
    /// @return the point
    static p256 from_compressed(const std::array<std::uint8_t, encoded_size> &in)
    {
        std::uint8_t raw[encoded_size];
        std::memcpy(raw, in.data(), encoded_size);
        return from_compressed(raw);
    }

    /// @brief Map a PRG block to a curve point.
    /// @details Stretches the seed to at least 256 bits, then hashes onto the
    ///          curve. Leaf and arithmetic paths require a valid encoding;
    ///          use this for masks, not for decoding peer points.
    /// @param bytes the PRG output
    /// @param n the number of bytes available
    /// @return the curve point
    static p256 from_seed(const void * bytes, std::size_t n) noexcept
    {
        std::uint8_t wide[64]{};
        p256_detail::stretch_seed(wide, sizeof(wide), bytes, n);
        p256_detail::affine point;
        p256_detail::hash_to_curve(wide, sizeof(wide), point);
        p256 out;
        p256_detail::encode_point(out.enc_, point);
        return out;
    }

    /// @brief The 33-byte encoding.
    /// @return the encoding
    HEDLEY_ALWAYS_INLINE
    const std::uint8_t *bytes() const noexcept { return enc_; }

    /// @brief True for the identity.
    /// @return whether this is the identity
    HEDLEY_ALWAYS_INLINE
    bool is_identity() const noexcept
    {
        std::uint8_t z = 0;
        for (std::size_t i = 0; i < encoded_size; ++i)
            z |= enc_[i];
        return z == 0;
    }

    /// @brief Point addition.
    /// @param a left addend
    /// @param b right addend
    /// @return `a + b` on the curve
    friend p256 operator+(p256 a, p256 b)
    {
        const auto pa = p256_detail::decode_strict(a.enc_, encoded_size);
        const auto pb = p256_detail::decode_strict(b.enc_, encoded_size);
        p256 out;
        p256_detail::encode_point(out.enc_, p256_detail::point_add(pa, pb));
        return out;
    }

    /// @brief Point negation.
    /// @param a the point to negate
    /// @return `-a`
    friend p256 operator-(p256 a)
    {
        const auto pa = p256_detail::decode_strict(a.enc_, encoded_size);
        p256 out;
        p256_detail::encode_point(out.enc_, p256_detail::point_neg(pa));
        return out;
    }

    /// @brief Point subtraction.
    /// @param a minuend
    /// @param b subtrahend
    /// @return `a - b` on the curve
    friend p256 operator-(p256 a, p256 b)
    {
        return a + (-b);
    }

    /// @brief Encoding equality.
    /// @param a left point
    /// @param b right point
    /// @return `true` when the encodings match
    friend bool operator==(p256 a, p256 b) noexcept
    {
        return std::memcmp(a.enc_, b.enc_, encoded_size) == 0;
    }

    /// @brief Encoding inequality.
    /// @param a left point
    /// @param b right point
    /// @return `true` when the encodings differ
    friend bool operator!=(p256 a, p256 b) noexcept
    {
        return !(a == b);
    }

    /// @brief Write the 33-byte encoding in hexadecimal.
    /// @param os the output stream
    /// @param a the point to write
    /// @return `os`
    friend std::ostream & operator<<(std::ostream &os, p256 a)
    {
        const auto flags = os.flags();
        os << std::hex;
        for (std::size_t i = 0; i < encoded_size; ++i)
        {
            const unsigned byte = a.enc_[i];
            os << "0123456789abcdef"[byte >> 4] << "0123456789abcdef"[byte & 0xf];
        }
        os.flags(flags);
        return os;
    }

  private:
    std::uint8_t enc_[encoded_size]{};
};

namespace leaf_arithmetic
{

namespace detail
{

template <typename Leaf>
HEDLEY_ALWAYS_INLINE
Leaf p256_leaf_combine(const Leaf &a, const Leaf &b, bool subtract)
{
    unsigned char ab[sizeof(Leaf)];
    unsigned char bb[sizeof(Leaf)];
    unsigned char cb[sizeof(Leaf)];
    std::memcpy(ab, &a, sizeof(Leaf));
    std::memcpy(bb, &b, sizeof(Leaf));
    const auto pa = p256_detail::decode_strict(ab, sizeof(Leaf));
    const auto pb = p256_detail::decode_strict(bb, sizeof(Leaf));
    const auto pc = subtract ? p256_detail::point_sub(pa, pb)
                             : p256_detail::point_add(pa, pb);
    std::memset(cb, 0, sizeof(cb));
    p256_detail::encode_point(cb, pc);
    Leaf out{};
    std::memcpy(&out, cb, sizeof(Leaf));
    return out;
}

}  // namespace detail

template <typename NodeT, std::size_t N>
struct add_t<p256, std::array<NodeT, N>>
{
    auto operator()(const std::array<NodeT, N> &a,
        const std::array<NodeT, N> &b) const
    {
        return detail::p256_leaf_combine(a, b, false);
    }
};

template <typename NodeT, std::size_t N>
struct subtract_t<p256, std::array<NodeT, N>>
{
    auto operator()(const std::array<NodeT, N> &a,
        const std::array<NodeT, N> &b) const
    {
        return detail::p256_leaf_combine(a, b, true);
    }
};

template <typename NodeT, std::size_t N>
struct multiply_t<p256, std::array<NodeT, N>>
{
    template <typename Leaf>
    Leaf operator()(const Leaf &, p256) const
    {
        static_assert(sizeof(NodeT) == 0,
            "p256 has no point product; p256{k} is k times the generator");
        return Leaf{};
    }
};

}  // namespace leaf_arithmetic

/// @brief Sample a uniform curve point by rejection-sampling a scalar in `[0, n)`.
/// @return a curve point
template <>
HEDLEY_NO_THROW
inline auto uniform_sample<p256>() noexcept
{
    std::uint64_t k[4]{};
    for (;;)
    {
        k[0] = uniform_sample<std::uint64_t>();
        k[1] = uniform_sample<std::uint64_t>();
        k[2] = uniform_sample<std::uint64_t>();
        k[3] = uniform_sample<std::uint64_t>();
        if (p256_detail::limbs_cmp(k, p256_detail::N, 4) < 0)
            break;
    }
    const auto point = p256_detail::point_scalarmul_limbs(
        p256_detail::generator_point(), k);
    if (point.inf)
        return p256{};
    unsigned char enc[p256::encoded_size];
    p256_detail::encode_point(enc, point);
    return p256::from_compressed(enc);
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_P256_HPP__