/// @file dpf/sfss.hpp
/// @brief Streaming DPF (SDPF) compiler: Song et al., USENIX Security 2026
///        (ePrint 2025/2304 §4.2). Setup once for a fixed point; Enc streams
///        short weight ciphertexts; Eval(key, x, ct) shares f(x)·m.
/// @details Leaf is `dpf::vec<fp61, 2>` = (β, r). Enc/Eval use a seed-keyed
///          AES-MMO PRF on the r-lane shares (subtractive convention).
///          Window Enc telescopes PRF masks so a sum of ciphertexts needs
///          only two PRF calls at Eval (Appendix C.2).
/// @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_SFSS_HPP__
#define LIBDPF_INCLUDE_DPF_SFSS_HPP__

#include <cstddef>
#include <cstdint>
#include <stdexcept>
#include <utility>

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

#include "dpf/eval_point.hpp"
#include "dpf/fp61.hpp"
#include "dpf/incremental.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/random.hpp"
#include "dpf/secret_share.hpp"
#include "dpf/vec.hpp"

namespace dpf
{

/// @brief Leaf payload for SDPF: lane 0 is β, lane 1 is the PRF-key sample r.
using sdpf_leaf = vec<fp61, 2>;

/// @brief Streaming encryption state (message counter).
struct sfss_state
{
    std::uint64_t ctr = 1;
};

/// @brief Client encryption key: subtractive shares of r at α, and β.
struct sdpf_enc_key
{
    fp61 h0{};
    fp61 h1{};
    fp61 beta{1};
};

/// @brief One streaming ciphertext.
struct sfss_ct
{
    std::uint64_t j = 0;
    fp61 c{};
};

/// @brief PRF F(k, j) → fp61 via fixed-key AES-MMO with seed from k.
/// @details Domain-separated from the tree PRG by the high qword tag.
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
fp61 sfss_prf(fp61 key, std::uint64_t j) noexcept
{
    const simde__m128i seed = simde_mm_set_epi64x(
        static_cast<std::int64_t>(0x5346535350524631ull),  // "SFSSPRF1"
        static_cast<std::int64_t>(key.raw()));
    const auto pos = static_cast<psnip_uint32_t>(j
        ^ static_cast<std::uint64_t>(j >> 32));
    const simde__m128i block = prg::aes128::eval(seed, pos);
    return fp61::from_seed(&block, sizeof(block));
}

namespace detail
{
namespace sfss_impl
{

HEDLEY_ALWAYS_INLINE
fp61 mask_diff(fp61 h0, fp61 h1, std::uint64_t j) noexcept
{
    return sfss_prf(h0, j) - sfss_prf(h1, j);
}

template <typename Share>
HEDLEY_ALWAYS_INLINE
sdpf_leaf leaf_of(const Share & y)
{
    if constexpr (std::is_same_v<std::decay_t<Share>, sdpf_leaf>)
        return y;
    else
        return y.raw();
}

}  // namespace sfss_impl
}  // namespace detail

/// @brief Gen SDPF keys for f_{α,β}. Returns (k0, k1, ke, st).
/// @tparam InteriorPRG tree interior PRG
/// @tparam ExteriorPRG leaf PRG
/// @param alpha the secret point
/// @param beta the point payload (default 1)
/// @throws std::invalid_argument if beta is 0
template <typename InteriorPRG = prg::aes128,
          typename ExteriorPRG = prg::aes128,
          typename InputT>
auto make_sdpf(InputT alpha, fp61 beta = fp61{1})
{
    if (beta.raw() == 0)
        throw std::invalid_argument("make_sdpf: beta must be nonzero");

    sdpf_leaf leaf{};
    leaf[0] = beta;
    leaf[1] = uniform_sample<fp61>();

    auto keys = make_dpf<InteriorPRG, ExteriorPRG>(alpha, leaf);
    auto & k0 = keys.first;
    auto & k1 = keys.second;

    const auto y0 = detail::sfss_impl::leaf_of(*eval_point(k0, alpha));
    const auto y1 = detail::sfss_impl::leaf_of(*eval_point(k1, alpha));

    sdpf_enc_key ke;
    ke.h0 = y0[1];
    ke.h1 = y1[1];
    ke.beta = beta;
    return std::make_tuple(std::move(k0), std::move(k1), ke, sfss_state{});
}

/// @brief Enc(st, ke, m) → ciphertext for stream index st.ctr; advances ctr.
/// @details c = m − (F(h0,j) − F(h1,j)) · β⁻¹  (β = 1 skips the inverse).
inline sfss_ct sfss_enc(sfss_state & st, const sdpf_enc_key & ke, fp61 m)
{
    if (st.ctr == 0)
        throw std::overflow_error("sfss_enc: counter wrapped");
    const std::uint64_t j = st.ctr++;
    const fp61 mask = detail::sfss_impl::mask_diff(ke.h0, ke.h1, j);
    fp61 c;
    if (ke.beta.raw() == 1)
        c = m - mask;
    else
        c = m - mask * detail::shamir_field<fp61>::inv(ke.beta);
    return sfss_ct{j, c};
}

/// @brief Window Enc: c = m − mask(j)·β⁻¹ + mask(j+1)·β⁻¹ (telescoping).
inline sfss_ct sfss_enc_window(sfss_state & st, const sdpf_enc_key & ke, fp61 m)
{
    if (st.ctr == 0 || st.ctr == UINT64_MAX)
        throw std::overflow_error("sfss_enc_window: counter wrapped");
    const std::uint64_t j = st.ctr++;
    fp61 mask_j = detail::sfss_impl::mask_diff(ke.h0, ke.h1, j);
    fp61 mask_n = detail::sfss_impl::mask_diff(ke.h0, ke.h1, j + 1);
    if (ke.beta.raw() != 1)
    {
        const fp61 inv_b = detail::shamir_field<fp61>::inv(ke.beta);
        mask_j = mask_j * inv_b;
        mask_n = mask_n * inv_b;
    }
    return sfss_ct{j, m - mask_j + mask_n};
}

/// @brief Eval(kb, x, ct) → subtractive share of f_{α,β}(x) · m.
/// @tparam Key a party key from `make_sdpf`
template <typename Key, typename InputT>
fp61 eval_sdpf(const Key & key, InputT x, const sfss_ct & ct)
{
    const auto leaf = detail::sfss_impl::leaf_of(*eval_point(key, x));
    const fp61 e = leaf[0];
    const fp61 h = leaf[1];
    return e * ct.c + sfss_prf(h, ct.j);
}

/// @brief Window Eval after summing window ciphertexts into `c_agg`.
/// @details s = e · c_agg + F(h, j_lo) − F(h, j_hi + 1).
///          For ciphertexts from `sfss_enc_window` with consecutive counters
///          in `[j_lo, j_hi]`, reconstruct(s0, s1) = f(x) · Σ m_j.
template <typename Key, typename InputT>
fp61 eval_sdpf_window(const Key & key, InputT x, fp61 c_agg,
    std::uint64_t j_lo, std::uint64_t j_hi)
{
    if (j_hi < j_lo)
        throw std::invalid_argument("eval_sdpf_window: empty window");
    const auto leaf = detail::sfss_impl::leaf_of(*eval_point(key, x));
    const fp61 e = leaf[0];
    const fp61 h = leaf[1];
    return e * c_agg + sfss_prf(h, j_lo) - sfss_prf(h, j_hi + 1);
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_SFSS_HPP__