/// @file dpf/random.hpp
/// @brief Entropy source and uniform sampling.
/// @author Ryan Henry <ryan.henry@ucalgary.ca>
/// @copyright Copyright (c) 2019-2024 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_RANDOM_HPP__
#define LIBDPF_INCLUDE_DPF_RANDOM_HPP__

#include <bsd/stdlib.h>

#include <array>
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <exception>
#include <fcntl.h>
#include <mutex>
#include <type_traits>
#include <unistd.h>
#include <utility>

#include "hedley/hedley.h"

#include "dpf/secret_share.hpp"

namespace dpf
{

namespace detail
{

/// @brief Thread-local count of bytes delivered by `uniform_fill`.
inline thread_local std::uint64_t random_bytes_tls = 0;

/// @brief When set, `uniform_fill` copies from this hook and does not read the
/// system RNG. Used to feed the same beaver coins to dealer `make_dpf` and
/// Doerner–Shelat gen. Null in normal use.
inline thread_local void (*uniform_bytes_hook)(void *, std::size_t) = nullptr;

/// @brief State for `uniform_bytes_hook`, set and read by the hook's owner
/// (`experiment` keeps itself here). Travels with the hook when another thread
/// adopts this one's draws (`dpf/thread_work.hpp`).
inline thread_local void * uniform_bytes_ctx = nullptr;

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
bool fill_from_hook(T & buf) noexcept
{
    if (uniform_bytes_hook == nullptr)
    {
        return false;
    }
    uniform_bytes_hook(&buf, sizeof(buf));
    return true;
}

/// @brief `bool` and `enum : bool` (including `dpf::bit`) have only two valid
/// representations. Filling them with a raw entropy byte is undefined.
/// @tparam T value type
/// @return `bool` and `enum : bool` (including `dpf::bit`) have only two valid representations
template <typename T>
HEDLEY_NO_THROW
constexpr bool is_boolean_representation() noexcept
{
    using U = std::remove_cv_t<T>;
    if constexpr (std::is_same_v<U, bool>)
    {
        return true;
    }
    else if constexpr (std::is_enum_v<U>)
    {
        return std::is_same_v<std::underlying_type_t<U>, bool>;
    }
    else
    {
        return false;
    }
}

#if !defined(LIBDPF_USE_ARC4RANDOM)

/// @brief One unbuffered, exclusively locked read of the entropy device.
/// @details Buffering would copy unread bytes into a `fork()` child, so parent and
/// child would repeat the same key material. The lock keeps concurrent
/// `fread` calls off the shared `FILE`.
struct entropy_source
{
#if defined(LIBDPF_USE_DEV_RANDOM)
    static constexpr const char * path = "/dev/random";
    static constexpr const char * open_error = "dpf: cannot open /dev/random\n";
#else
    static constexpr const char * path = "/dev/urandom";
    static constexpr const char * open_error = "dpf: cannot open /dev/urandom\n";
#endif

    FILE * fp = nullptr;
    std::mutex mu;

    entropy_source() = default;
    entropy_source(const entropy_source &) = delete;
    entropy_source & operator=(const entropy_source &) = delete;
    entropy_source(entropy_source &&) = delete;
    entropy_source & operator=(entropy_source &&) = delete;

    HEDLEY_NO_THROW
    ~entropy_source() noexcept
    {
        if (fp != nullptr)
        {
            std::fclose(fp);
        }
    }

    void open_unlocked()
    {
        if (fp != nullptr)
        {
            return;
        }
        fp = std::fopen(path, "rb");
        if (fp == nullptr)
        {
            std::fputs(open_error, stderr);
            std::terminate();
        }
        // Before any read. A buffered FILE duplicates entropy across fork().
        if (std::setvbuf(fp, nullptr, _IONBF, 0) != 0)
        {
            std::fclose(fp);
            fp = nullptr;
            std::fputs("dpf: cannot disable entropy buffering\n", stderr);
            std::terminate();
        }
        int fd = ::fileno(fp);
        if (fd >= 0)
        {
            ::fcntl(fd, F_SETFD, FD_CLOEXEC);
        }
    }

    void read(void * dst, std::size_t n)
    {
        std::lock_guard<std::mutex> lock(mu);
        if (fp == nullptr)
        {
            open_unlocked();
        }
        auto * p = static_cast<unsigned char *>(dst);
        while (n > 0)
        {
            std::size_t got = std::fread(p, 1, n, fp);
            if (got == 0)
            {
                if (std::ferror(fp) && errno == EINTR)
                {
                    std::clearerr(fp);
                    continue;
                }
                std::fputs("dpf: entropy read failed\n", stderr);
                std::terminate();
            }
            p += got;
            n -= got;
        }
    }
};

inline entropy_source & entropy()
{
    static entropy_source source;
    return source;
}

#endif  // !LIBDPF_USE_ARC4RANDOM

}  // namespace detail

/// @brief Zero the thread-local random-byte counter.
HEDLEY_ALWAYS_INLINE
void reset_random_bytes_count() noexcept
{
    detail::random_bytes_tls = 0;
}

/// @brief Bytes filled by `uniform_fill` since the last reset on this thread.
HEDLEY_ALWAYS_INLINE
std::uint64_t random_bytes_count() noexcept
{
    return detail::random_bytes_tls;
}

template <typename T>
HEDLEY_NO_THROW
auto & uniform_fill(T & buf) noexcept  // NOLINT(runtime/references)
{
    static_assert(std::is_trivially_copyable_v<std::remove_cv_t<T>>,
        "uniform_fill requires a trivially copyable type");

    if constexpr (detail::is_boolean_representation<T>())
    {
        unsigned char raw = 0;
        uniform_fill(raw);
        buf = static_cast<T>(static_cast<bool>(raw & 1u));
        return buf;
    }
    else
    {
        if (detail::fill_from_hook(buf))
        {
            detail::random_bytes_tls += sizeof(buf);
            return buf;
        }
#if defined(LIBDPF_USE_ARC4RANDOM)
        arc4random_buf(&buf, sizeof(buf));
#else
        detail::entropy().read(&buf, sizeof(buf));
#endif
        detail::random_bytes_tls += sizeof(buf);
        return buf;
    }
}

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
auto uniform_sample() noexcept
{
    using U = std::remove_cv_t<T>;
    U buf;
    uniform_fill(buf);
    return buf;
}

template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
auto additively_share(T secret) noexcept
{
    using T_ = std::remove_cv_t<std::remove_reference_t<T>>;
    T_ tmp = uniform_sample<T_>();
    T_ other = detail::group_sub(static_cast<T_>(secret), tmp);
    return std::make_pair(
        additive_share<T_, 0>::from_raw(tmp),
        additive_share<T_, 1>::from_raw(other));
}

/// @brief Uniform (3,3)-additive sharing of `secret`.
/// @details Two components are uniform. The third is `secret` minus those
///          two in the share group, so the three shares sum to `secret`.
/// @tparam T value type
/// @param secret the cleartext secret
/// @return shares for parties 0, 1, and 2
template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
auto additively_share3(T secret) noexcept
{
    using T_ = std::remove_cv_t<std::remove_reference_t<T>>;
    const T_ a = uniform_sample<T_>();
    const T_ b = uniform_sample<T_>();
    const T_ c = detail::group_sub(
        detail::group_sub(static_cast<T_>(secret), a), b);
    return std::make_tuple(
        additive3_share<T_, 0>::from_raw(a),
        additive3_share<T_, 1>::from_raw(b),
        additive3_share<T_, 2>::from_raw(c));
}

/// @brief Uniform (2,3)-replicated sharing of `secret`.
/// @details The underlying (3,3) components are a uniform additive split.
///          Each party receives its component and the next party's.
/// @tparam T value type
/// @param secret the cleartext secret
/// @return shares for parties 0, 1, and 2
template <typename T>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
auto share_replicated(T secret) noexcept
{
    using T_ = std::remove_cv_t<std::remove_reference_t<T>>;
    const T_ x0 = uniform_sample<T_>();
    const T_ x1 = uniform_sample<T_>();
    const T_ x2 = detail::group_sub(
        detail::group_sub(static_cast<T_>(secret), x0), x1);
    return make_replicated_shares(x0, x1, x2);
}

namespace shamir
{

/// @brief Uniform `(K,N)` Shamir sharing of `secret`.
/// @details Coefficients of `x, ..., x^{K-1}` are `uniform_sample<T>`. The
///          shares are `deal<T, K, N>`. Threshold 1 draws nothing: every share
///          equals `secret`. `shamir3::share_secret` is the `(2,3)` case on
///          `fp61`, reindexed to points `1`, `2`, and `3`.
/// @tparam K shares required to reconstruct
/// @tparam N shareholders
/// @tparam T field type. Opening needs `detail::shamir_field<T>`
/// @param secret the cleartext secret
/// @return one share per party `0 .. N-1`
/// \complexity O(NK) field operations, plus `K-1` field samples. No messages.
template <typename T, std::size_t K, std::size_t N>
HEDLEY_WARN_UNUSED_RESULT
HEDLEY_NO_THROW
auto share_secret(T secret) noexcept
{
    using T_ = std::remove_cv_t<std::remove_reference_t<T>>;
    constexpr std::size_t degree = access<K, N>::degree;
    std::array<T_, degree> coeff{};
    for (std::size_t i = 0; i < degree; ++i)
        coeff[i] = uniform_sample<T_>();
    return deal<T_, K, N>(static_cast<T_>(secret), coeff);
}

}  // namespace shamir

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_RANDOM_HPP__