/// @file dpf/prg_chacha.hpp
/// @brief ChaCha stream PRG. Same 128-bit block interface as `aes128`.
/// @details `eval(seed, pos)` is the `pos`-th 16-byte chunk of ChaCha
///          keystream (RFC 8439). The 256-bit key is the 128-bit seed
///          followed by the fixed domain separator `"dpf-chacha-prg\0\0"`.
///          The nonce is zero. The ChaCha block counter is `pos / 4`, and
///          the chunk inside that block is `pos % 4`.
///
///          `chacha20` is the RFC round count. `chacha12` and `chacha8` are
///          the same construction with fewer rounds. `chacha<R>` accepts any
///          positive even round count.
/// @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_PRG_CHACHA_HPP__
#define LIBDPF_INCLUDE_DPF_PRG_CHACHA_HPP__

#include <array>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <stdexcept>

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

#include "dpf/prg_count.hpp"

namespace dpf
{

namespace prg
{

namespace chacha_detail
{

inline constexpr std::uint32_t zero_nonce[3] = {0, 0, 0};

/// @brief ASCII `"dpf-chacha-prg"` plus two zero bytes. Public second half of the key.
inline constexpr std::uint8_t domain[16] = {
    'd', 'p', 'f', '-', 'c', 'h', 'a', 'c',
    'h', 'a', '-', 'p', 'r', 'g', 0, 0
};

HEDLEY_PURE
HEDLEY_NON_NULL(1)
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr std::uint32_t load_le32(const std::uint8_t * p) noexcept
{
    return static_cast<std::uint32_t>(p[0])
        | (static_cast<std::uint32_t>(p[1]) << 8)
        | (static_cast<std::uint32_t>(p[2]) << 16)
        | (static_cast<std::uint32_t>(p[3]) << 24);
}

HEDLEY_NON_NULL(1)
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
void store_le32(std::uint8_t * p, std::uint32_t w) noexcept
{
    p[0] = static_cast<std::uint8_t>(w);
    p[1] = static_cast<std::uint8_t>(w >> 8);
    p[2] = static_cast<std::uint8_t>(w >> 16);
    p[3] = static_cast<std::uint8_t>(w >> 24);
}

HEDLEY_PURE
HEDLEY_NON_NULL(1)
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
simde__m128i load_block(const std::uint8_t * p) noexcept
{
    simde__m128i out;
    std::memcpy(&out, p, sizeof(out));
    return out;
}

/// @brief 128-bit seed in the low half, `domain` in the high half, both little-endian.
/// @param seed the PRG seed
/// @param key the `key`
HEDLEY_NON_NULL(2)
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
void seed_key(simde__m128i seed, std::uint32_t key[8]) noexcept
{
    std::uint8_t raw[16];
    std::memcpy(raw, &seed, sizeof(raw));
    for (int i = 0; i < 4; ++i)
    {
        key[i] = load_le32(raw + 4 * i);
        key[4 + i] = load_le32(domain + 4 * i);
    }
}

template <int N>
HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr std::uint32_t rotl(std::uint32_t x) noexcept
{
    static_assert(N > 0 && N < 32, "ChaCha rotation is between 1 and 31");
    return (x << N) | (x >> (32 - N));
}

HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
void quarter(std::uint32_t & a, std::uint32_t & b,
    std::uint32_t & c, std::uint32_t & d) noexcept
{
    a += b; d ^= a; d = rotl<16>(d);
    c += d; b ^= c; b = rotl<12>(b);
    a += b; d ^= a; d = rotl<8>(d);
    c += d; b ^= c; b = rotl<7>(b);
}

template <int N>
HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
simde__m128i rotl_epi32(simde__m128i v) noexcept
{
    static_assert(N > 0 && N < 32, "ChaCha rotation is between 1 and 31");
    return simde_mm_or_si128(simde_mm_slli_epi32(v, N),
        simde_mm_srli_epi32(v, 32 - N));
}

HEDLEY_NON_NULL(1)
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
void quarter(simde__m128i x[], int a, int b, int c, int d) noexcept
{
    x[a] = simde_mm_add_epi32(x[a], x[b]);
    x[d] = rotl_epi32<16>(simde_mm_xor_si128(x[d], x[a]));
    x[c] = simde_mm_add_epi32(x[c], x[d]);
    x[b] = rotl_epi32<12>(simde_mm_xor_si128(x[b], x[c]));
    x[a] = simde_mm_add_epi32(x[a], x[b]);
    x[d] = rotl_epi32<8>(simde_mm_xor_si128(x[d], x[a]));
    x[c] = simde_mm_add_epi32(x[c], x[d]);
    x[b] = rotl_epi32<7>(simde_mm_xor_si128(x[b], x[c]));
}

HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
std::uint32_t epi32_lane(simde__m128i v, int lane) noexcept
{
    // Shuffle control is an immediate, so each lane is its own case.
    switch (lane)
    {
    case 1: v = simde_mm_shuffle_epi32(v, 0x01); break;
    case 2: v = simde_mm_shuffle_epi32(v, 0x02); break;
    case 3: v = simde_mm_shuffle_epi32(v, 0x03); break;
    default: break;
    }
    return static_cast<std::uint32_t>(simde_mm_cvtsi128_si32(v));
}

/// @name ChaCha blocks
/// @tparam Rounds ChaCha round count. Must be positive and even
/// @param key the ChaCha key words
/// @param counter the ChaCha block counter
/// @param out the output buffer
/// @{

/// @brief One ChaCha block.
/// @details `key` is 8 little-endian words. `nonce` is 3 words.
/// @param key the ChaCha key words
/// @param counter the ChaCha block counter
/// @param nonce the ChaCha nonce
/// @param out the output buffer
template <unsigned Rounds>
HEDLEY_NON_NULL(1, 3, 4)
HEDLEY_NO_THROW
void block(const std::uint32_t key[8], std::uint32_t counter,
    const std::uint32_t nonce[3], std::uint8_t out[64]) noexcept
{
    static_assert(Rounds >= 2 && Rounds % 2 == 0,
        "ChaCha rounds must be a positive even number");

    std::uint32_t s[16] = {
        0x61707865u, 0x3320646eu, 0x79622d32u, 0x6b206574u,
        key[0], key[1], key[2], key[3],
        key[4], key[5], key[6], key[7],
        counter, nonce[0], nonce[1], nonce[2]
    };
    std::uint32_t orig[16];
    std::memcpy(orig, s, sizeof(orig));

HEDLEY_PRAGMA(GCC unroll 16)
    for (unsigned r = 0; r < Rounds; r += 2)
    {
        quarter(s[0], s[4], s[8], s[12]);
        quarter(s[1], s[5], s[9], s[13]);
        quarter(s[2], s[6], s[10], s[14]);
        quarter(s[3], s[7], s[11], s[15]);
        quarter(s[0], s[5], s[10], s[15]);
        quarter(s[1], s[6], s[11], s[12]);
        quarter(s[2], s[7], s[8], s[13]);
        quarter(s[3], s[4], s[9], s[14]);
    }

    for (int i = 0; i < 16; ++i)
    {
        store_le32(out + 4 * i, s[i] + orig[i]);
    }
}

/// @brief Four independent ChaCha blocks.
/// @details Lane `i` uses `key[i]` and `counter[i]`. Nonce is zero. Each `out[i]` receives 64 bytes.
template <unsigned Rounds>
HEDLEY_NON_NULL(1, 2, 3)
HEDLEY_NO_THROW
void block4(const std::uint32_t key[][8], const std::uint32_t counter[4],
    std::uint8_t out[][64]) noexcept
{
    static_assert(Rounds >= 2 && Rounds % 2 == 0,
        "ChaCha rounds must be a positive even number");

    simde__m128i x[16];
    x[0] = simde_mm_set1_epi32(static_cast<int>(0x61707865u));
    x[1] = simde_mm_set1_epi32(static_cast<int>(0x3320646eu));
    x[2] = simde_mm_set1_epi32(static_cast<int>(0x79622d32u));
    x[3] = simde_mm_set1_epi32(static_cast<int>(0x6b206574u));
    for (int w = 0; w < 8; ++w)
    {
        x[4 + w] = simde_mm_set_epi32(
            static_cast<int>(key[3][w]),
            static_cast<int>(key[2][w]),
            static_cast<int>(key[1][w]),
            static_cast<int>(key[0][w]));
    }
    x[12] = simde_mm_set_epi32(
        static_cast<int>(counter[3]),
        static_cast<int>(counter[2]),
        static_cast<int>(counter[1]),
        static_cast<int>(counter[0]));
    x[13] = simde_mm_setzero_si128();
    x[14] = simde_mm_setzero_si128();
    x[15] = simde_mm_setzero_si128();

    simde__m128i orig[16];
    for (int i = 0; i < 16; ++i)
    {
        orig[i] = x[i];
    }

HEDLEY_PRAGMA(GCC unroll 16)
    for (unsigned r = 0; r < Rounds; r += 2)
    {
        quarter(x, 0, 4, 8, 12);
        quarter(x, 1, 5, 9, 13);
        quarter(x, 2, 6, 10, 14);
        quarter(x, 3, 7, 11, 15);
        quarter(x, 0, 5, 10, 15);
        quarter(x, 1, 6, 11, 12);
        quarter(x, 2, 7, 8, 13);
        quarter(x, 3, 4, 9, 14);
    }

    for (int lane = 0; lane < 4; ++lane)
    {
        for (int w = 0; w < 16; ++w)
        {
            std::uint32_t sum = epi32_lane(x[w], lane) + epi32_lane(orig[w], lane);
            store_le32(out[lane] + 4 * w, sum);
        }
    }
}

/// @}

}  // namespace chacha_detail

/// @brief ChaCha stream PRG with `Rounds` rounds (20 is RFC 8439).
/// @tparam Rounds ChaCha round count. Must be positive and even
template <unsigned Rounds = 20>
struct chacha final
{
    static_assert(Rounds >= 2 && Rounds % 2 == 0,
        "ChaCha rounds must be a positive even number");

    using block_type = simde__m128i;
    static constexpr unsigned rounds = Rounds;

    /// @brief One 128-bit lane of a ChaCha block.
    /// @param seed the PRG seed
    /// @param pos lane index; the block counter is `pos >> 2`
    /// @return that lane
    /// \complexity One ChaCha block: `Rounds/2` double-rounds, eight quarter-rounds each (`chacha_detail::block`).
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static block_type eval(block_type seed, psnip_uint32_t pos) noexcept
    {
        note_eval(1, primitive::chacha);
        std::uint32_t key[8];
        chacha_detail::seed_key(seed, key);
        std::uint8_t buf[64];
        chacha_detail::block<Rounds>(key, pos >> 2,
            chacha_detail::zero_nonce, buf);
        return chacha_detail::load_block(buf + 16 * (pos & 3u));
    }

    /// @brief Positions 0 and 1, one ChaCha block (the first 32 keystream bytes).
    /// @param seed the PRG seed
    /// @return Positions 0 and 1, one ChaCha block (the first 32 keystream bytes)
    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    static auto eval01(block_type seed) noexcept
    {
        note_eval(2, primitive::chacha);
        std::uint32_t key[8];
        chacha_detail::seed_key(seed, key);
        std::uint8_t buf[64];
        chacha_detail::block<Rounds>(key, 0, chacha_detail::zero_nonce, buf);
HEDLEY_PRAGMA(GCC diagnostic push)
HEDLEY_PRAGMA(GCC diagnostic ignored "-Wignored-attributes")
        return std::array<block_type, 2>{
            chacha_detail::load_block(buf),
            chacha_detail::load_block(buf + 16)
        };
HEDLEY_PRAGMA(GCC diagnostic pop)
    }

    /// @brief `count` blocks starting at lane `pos`. `output` is unused when `count` is 0.
    /// @param seed the PRG seed
    /// @param output the destination. Unused when `count` is 0
    /// @param count the number of blocks
    /// @param pos the 0-based index
    /// @throws std::invalid_argument if `pos + count` wraps `uint32_t`.
    HEDLEY_ALWAYS_INLINE
    static void eval(block_type seed, block_type * HEDLEY_RESTRICT output,
        psnip_uint32_t count, psnip_uint32_t pos = 0)
    {
        if (HEDLEY_UNLIKELY(count == 0))
        {
            return;
        }
        if (count > 1 &&
            pos > static_cast<psnip_uint32_t>(~static_cast<psnip_uint32_t>(0)) - (count - 1u))
        {
            throw std::invalid_argument("prg lane index is out of range");
        }
        if (count == 1)
        {
            output[0] = eval(seed, pos);
            return;
        }
        if (count == 2 && pos == 0)
        {
            auto kids = eval01(seed);
            output[0] = kids[0];
            output[1] = kids[1];
            return;
        }

        note_eval(count, primitive::chacha);
        std::uint32_t key[8];
        chacha_detail::seed_key(seed, key);
        psnip_uint32_t i = 0;

        // `pos` may begin mid-block. Those chunks share one ChaCha block.
        if ((pos & 3u) != 0u)
        {
            std::uint8_t buf[64];
            chacha_detail::block<Rounds>(key, pos >> 2,
                chacha_detail::zero_nonce, buf);
            while (i < count && ((pos + i) & 3u) != 0u)
            {
                output[i] = chacha_detail::load_block(
                    buf + 16 * ((pos + i) & 3u));
                ++i;
            }
        }

        // Four consecutive counters cover 16 output blocks.
        while (i + 16u <= count)
        {
            std::uint32_t base = (pos + i) >> 2;
            std::uint32_t keys[4][8];
            std::uint32_t counters[4];
            for (int lane = 0; lane < 4; ++lane)
            {
                std::memcpy(keys[lane], key, sizeof(key));
                counters[lane] = base + static_cast<std::uint32_t>(lane);
            }
            std::uint8_t buf[4][64];
            chacha_detail::block4<Rounds>(keys, counters, buf);
            for (int lane = 0; lane < 4; ++lane)
            {
                for (int chunk = 0; chunk < 4; ++chunk)
                {
                    output[i++] = chacha_detail::load_block(
                        buf[lane] + 16 * chunk);
                }
            }
        }

        while (i + 4u <= count)
        {
            std::uint8_t buf[64];
            chacha_detail::block<Rounds>(key, (pos + i) >> 2,
                chacha_detail::zero_nonce, buf);
            for (int chunk = 0; chunk < 4; ++chunk)
            {
                output[i++] = chacha_detail::load_block(buf + 16 * chunk);
            }
        }

        if (i < count)
        {
            std::uint8_t buf[64];
            chacha_detail::block<Rounds>(key, (pos + i) >> 2,
                chacha_detail::zero_nonce, buf);
            unsigned chunk = 0;
            while (i < count)
            {
                output[i++] = chacha_detail::load_block(buf + 16 * chunk);
                ++chunk;
            }
        }
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2, 3)
    static void eval01_x4(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT left,
        block_type * HEDLEY_RESTRICT right) noexcept
    {
        std::uint32_t keys[4][8];
        std::uint32_t counters[4] = {0, 0, 0, 0};
        for (int lane = 0; lane < 4; ++lane)
        {
            chacha_detail::seed_key(seeds[lane], keys[lane]);
        }
        std::uint8_t buf[4][64];
        chacha_detail::block4<Rounds>(keys, counters, buf);
        for (int lane = 0; lane < 4; ++lane)
        {
            left[lane] = chacha_detail::load_block(buf[lane]);
            right[lane] = chacha_detail::load_block(buf[lane] + 16);
        }
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2)
    static void eval_x4(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT output, psnip_uint32_t pos = 0) noexcept
    {
        std::uint32_t keys[4][8];
        std::uint32_t ctr = pos >> 2;
        std::uint32_t counters[4] = {ctr, ctr, ctr, ctr};
        for (int lane = 0; lane < 4; ++lane)
        {
            chacha_detail::seed_key(seeds[lane], keys[lane]);
        }
        std::uint8_t buf[4][64];
        chacha_detail::block4<Rounds>(keys, counters, buf);
        unsigned chunk = pos & 3u;
        for (int lane = 0; lane < 4; ++lane)
        {
            output[lane] = chacha_detail::load_block(buf[lane] + 16 * chunk);
        }
    }

    HEDLEY_NO_THROW
    HEDLEY_ALWAYS_INLINE
    HEDLEY_NON_NULL(1, 2)
    static void eval_x8(const block_type * HEDLEY_RESTRICT seeds,
        block_type * HEDLEY_RESTRICT output, psnip_uint32_t pos = 0) noexcept
    {
        eval_x4(seeds, output, pos);
        eval_x4(seeds + 4, output + 4, pos);
    }

    /// @brief Raw-bit subtractive share of `T` for party `Party` (see `prg.hpp`).
    /// @tparam T value type
    /// @tparam Party party index, `0` or `1`
    /// @param seed the PRG seed
    /// @param pos the 0-based index
    /// @return Raw-bit subtractive share of `T` for party `Party` (see `prg.hpp`)
    /// @see `prg.hpp`
    template <typename T, std::size_t Party>
    HEDLEY_NO_THROW
    static auto expand(block_type seed, psnip_uint32_t pos = 0) noexcept;
};  // struct chacha

/// @brief RFC 8439 ChaCha20.
using chacha20 = chacha<20>;

/// @brief ChaCha12. Same keying as `chacha20`, 12 rounds.
using chacha12 = chacha<12>;

/// @brief ChaCha8. Same keying as `chacha20`, 8 rounds.
using chacha8 = chacha<8>;

}  // namespace prg

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PRG_CHACHA_HPP__