/// @file dpf/verifiable.hpp
/// @brief Verifiable evaluation tokens and extractable-key helpers.
/// @details VDPF proof fold follows de Castro–Polychroniadou (hash-based
///          correction seeds, 2λ-bit tokens, equality Verify). Extractable
///          checks are public-part equality, ROM-style leaf XOF, and an
///          \(\mathbb{F}_{2^{61}-1}\) weight-1 subset sketch. Phantom tags
///          `dpf::verifiable` / `dpf::extractable` live in placement.hpp.
/// @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_VERIFIABLE_HPP__
#define LIBDPF_INCLUDE_DPF_VERIFIABLE_HPP__

#include <array>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <iterator>
#include <type_traits>
#include <utility>

#include "hedley/hedley.h"
#include "simde/simde/x86/avx2.h"
#include "portable-snippets/exact-int/exact-int.h"

#include "dpf/placement.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/fp61.hpp"
#include "dpf/xor_wrapper.hpp"
#include "dpf/twiddle.hpp"
#include "dpf/utils.hpp"

namespace dpf
{

/// 4λ = 64-byte correction seed (four AES blocks).
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
using cs_block = std::array<simde__m128i, 4>;
/// 2λ = 32-byte proof token (two AES blocks).
using proof_token = std::array<simde__m128i, 2>;
HEDLEY_PRAGMA(GCC diagnostic pop)

namespace detail
{
namespace vdpf
{

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
simde__m128i mmo(simde__m128i seed, psnip_uint32_t pos) noexcept
{
    return prg::aes128::eval(seed, pos);
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
cs_block hash_level_seed(std::size_t level, simde__m128i seed) noexcept
{
    const simde__m128i tagged = simde_mm_xor_si128(seed,
        simde_mm_set_epi64x(static_cast<psnip_int64_t>(0x56),
            static_cast<psnip_int64_t>(level)));
    return cs_block{
        mmo(tagged, 0),
        mmo(tagged, 1),
        mmo(tagged, 2),
        mmo(tagged, 3)};
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
cs_block hash_node(std::size_t level, psnip_uint64_t x_bits,
    simde__m128i seed) noexcept
{
    const simde__m128i tagged = simde_mm_xor_si128(seed,
        simde_mm_set_epi64x(static_cast<psnip_int64_t>(0x5600 | (level & 0xff)),
            static_cast<psnip_int64_t>(x_bits)));
    return cs_block{
        mmo(tagged, 0),
        mmo(tagged, 1),
        mmo(tagged, 2),
        mmo(tagged, 3)};
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
cs_block make_cs(std::size_t level, psnip_uint64_t prefix_bits,
    simde__m128i s0, simde__m128i s1) noexcept
{
    const auto h0v = hash_node(level, prefix_bits, s0);
    const auto h1v = hash_node(level, prefix_bits, s1);
    return cs_block{
        simde_mm_xor_si128(h0v[0], h1v[0]),
        simde_mm_xor_si128(h0v[1], h1v[1]),
        simde_mm_xor_si128(h0v[2], h1v[2]),
        simde_mm_xor_si128(h0v[3], h1v[3])};
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
cs_block correct(cs_block pi_tilde, const cs_block & cs,
    bool t) noexcept
{
    if (!t)
        return pi_tilde;
    return cs_block{
        simde_mm_xor_si128(pi_tilde[0], cs[0]),
        simde_mm_xor_si128(pi_tilde[1], cs[1]),
        simde_mm_xor_si128(pi_tilde[2], cs[2]),
        simde_mm_xor_si128(pi_tilde[3], cs[3])};
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
proof_token h0(const cs_block & in) noexcept
{
    const simde__m128i a = simde_mm_xor_si128(in[0], in[2]);
    const simde__m128i b = simde_mm_xor_si128(in[1], in[3]);
    return proof_token{mmo(a, 0x48), mmo(b, 0x48)};
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_CONST
proof_token xor_proof(proof_token a, proof_token b) noexcept
{
    return proof_token{
        simde_mm_xor_si128(a[0], b[0]),
        simde_mm_xor_si128(a[1], b[1])};
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_CONST
proof_token zero_proof() noexcept
{
    return proof_token{simde_mm_setzero_si128(), simde_mm_setzero_si128()};
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
void fold_node(proof_token & pi, std::size_t level,
    psnip_uint64_t x_bits, simde__m128i seed, const cs_block & cs) noexcept
{
    const bool t = static_cast<bool>(dpf::get_lo_bit(seed));
    const cs_block tilde = hash_node(level, x_bits, seed);
    const cs_block corrected = correct(tilde, cs, t);
    cs_block mixed{
        simde_mm_xor_si128(pi[0], corrected[0]),
        simde_mm_xor_si128(pi[1], corrected[1]),
        corrected[2],
        corrected[3]};
    pi = xor_proof(pi, h0(mixed));
}

HEDLEY_NO_THROW
inline void leaf_xof(simde__m128i seed, simde__m128i * HEDLEY_RESTRICT out,
    psnip_uint32_t count, psnip_uint32_t pos = 0) noexcept
{
    const simde__m128i tagged = simde_mm_xor_si128(seed,
        simde_mm_set_epi64x(0x45, 0));
    for (psnip_uint32_t i = 0; i < count; ++i)
        out[i] = mmo(tagged, pos + i);
}

/// Drop-in exterior PRG for leaf stretch under `dpf::extractable`.
template <typename BasePRG>
struct extractable_leaf_prg
{
    using block_type = typename BasePRG::block_type;
    HEDLEY_NO_THROW
    static void eval(block_type seed, block_type * HEDLEY_RESTRICT out,
        psnip_uint32_t count, psnip_uint32_t pos = 0) noexcept
    {
        leaf_xof(seed, out, count, pos);
    }
};

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
bool proof_equal(const proof_token & a, const proof_token & b) noexcept
{
    return std::memcmp(&a, &b, sizeof(proof_token)) == 0;
}

template <typename KeyT>
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
void init_proof(proof_token & pi, const KeyT & /*key*/) noexcept
{
    // Running proof starts at 0; each fold mixes in corrected leaf digests.
    pi = zero_proof();
}

}  // namespace vdpf
}  // namespace detail

/// @brief A proof token the caller owns, passed into evaluation.
struct prove_ref
{
    /// @brief The token updated by the evaluation.
    proof_token & token;
    /// @brief Bind `t`.
    /// @param t the token to update
    HEDLEY_NO_THROW
    explicit prove_ref(proof_token & t) noexcept : token{t} { }
};

/// @brief Bind `t` as the proof accumulator for one evaluation.
/// @param t the token to update
/// @return a `prove_ref` bound to `t`
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
prove_ref prove(proof_token & t) noexcept
{
    return prove_ref{t};
}

/// @brief Whether two proof tokens are identical.
/// @param a the first token
/// @param b the second token
/// @return `true` when every byte matches
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_PURE
bool verify(const proof_token & a, const proof_token & b) noexcept
{
    return detail::vdpf::proof_equal(a, b);
}

/// @brief Fold two batches of proof tokens and compare them.
/// @tparam Range0 range of `proof_token` for party 0
/// @tparam Range1 range of `proof_token` for party 1
/// @param left party 0 tokens, in evaluation order
/// @param right party 1 tokens, in the same order
/// @return `false` when the ranges differ in length or the folded tokens differ
template <typename Range0, typename Range1>
bool verify_batch(Range0 && left, Range1 && right)
{
    proof_token a = detail::vdpf::zero_proof();
    proof_token b = detail::vdpf::zero_proof();
    auto it0 = std::begin(left);
    auto it1 = std::begin(right);
    const auto end0 = std::end(left);
    const auto end1 = std::end(right);
    for (; it0 != end0 && it1 != end1; ++it0, ++it1)
    {
        a = detail::vdpf::xor_proof(a, *it0);
        b = detail::vdpf::xor_proof(b, *it1);
        a[0] = detail::vdpf::mmo(a[0], 1);
        b[0] = detail::vdpf::mmo(b[0], 1);
    }
    if (it0 != end0 || it1 != end1)
        return false;
    return verify(a, b);
}

/// @brief Whether two keys publish the same correction words, advice, and hash.
/// @tparam KeyT0 key type of party 0
/// @tparam KeyT1 key type of party 1
/// @param k0 party 0 key
/// @param k1 party 1 key
/// @return `false` when a public field differs
template <typename KeyT0, typename KeyT1>
bool same_public_part(const KeyT0 & k0, const KeyT1 & k1)
{
    static_assert(KeyT0::is_verifiable == KeyT1::is_verifiable,
        "same_public_part: mismatched verifiable flags");
    if (std::memcmp(k0.correction_words().data(), k1.correction_words().data(),
            sizeof(typename KeyT0::correction_words_array)) != 0)
        return false;
    if (std::memcmp(k0.correction_advice().data(), k1.correction_advice().data(),
            sizeof(typename KeyT0::correction_advice_array)) != 0)
        return false;
    if constexpr (KeyT0::is_verifiable)
    {
        if (std::memcmp(k0.correction_seeds().data(),
                k1.correction_seeds().data(),
                sizeof(typename KeyT0::correction_seeds_array)) != 0)
            return false;
    }
    return std::memcmp(&k0.common_part_hash(), &k1.common_part_hash(),
               sizeof(digest_type)) == 0;
}

struct sketch_share
{
    fp61 z1{};
    fp61 z2{};
    fp61 z3{};
};

/// @brief Weight-1 subset sketch of payloads `ys` against challenges `rs`.
/// @tparam YRange range of integers convertible to `fp61`
/// @tparam RRange range of challenges, one per payload
/// @param ys the payloads
/// @param rs the challenges
/// @return the three folded moments. A short range stops at the shorter end
template <typename YRange, typename RRange>
sketch_share sketch_fold(YRange && ys, RRange && rs)
{
    sketch_share out{};
    auto iy = std::begin(ys);
    auto ir = std::begin(rs);
    const auto ey = std::end(ys);
    const auto er = std::end(rs);
    for (; iy != ey && ir != er; ++iy, ++ir)
    {
        const fp61 y{*iy};
        const fp61 r{*ir};
        const fp61 r2 = r * r;
        out.z1 = out.z1 + y;
        out.z2 = out.z2 + y * r;
        out.z3 = out.z3 + y * r2;
    }
    return out;
}

/// @brief Whether `s0 - s1` is a weight-1 subset sketch.
/// @param s0 party 0's folded sketch
/// @param s1 party 1's folded sketch
/// @return `true` when `z2² = z1 · z3` after the shares are opened
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
HEDLEY_CONST
bool sketch_verify(sketch_share s0, sketch_share s1) noexcept
{
    const fp61 z1 = s0.z1 - s1.z1;
    const fp61 z2 = s0.z2 - s1.z2;
    const fp61 z3 = s0.z3 - s1.z3;
    return (z2 * z2) == (z1 * z3);
}

template <typename T, typename = void>
struct has_dpf_fp61 : std::false_type
{ };
template <typename T>
struct has_dpf_fp61<T, std::void_t<decltype(std::decay_t<T>::dpf_fp61)>>
    : std::bool_constant<std::decay_t<T>::dpf_fp61>
{ };

template <typename T, typename = void>
struct extractable_codomain_ok
    : std::bool_constant<
          (utils::bitlength_of_v<std::decay_t<T>> >= 128)
          || has_dpf_fp61<T>::value>
{ };
template <typename T>
struct extractable_codomain_ok<xor_wrapper<T>, void>
    : extractable_codomain_ok<T>
{ };
template <typename T>
inline constexpr bool extractable_codomain_ok_v =
    extractable_codomain_ok<std::decay_t<T>>::value;

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_VERIFIABLE_HPP__