/// @file dpf/thread_work.hpp
/// @brief Lend a thread's randomness source to another thread for one task,
///        and charge the task's costs back to the thread that asked for it.
/// @details A kernel posted to a compute pool runs on a pool thread while the
///          party thread waits for it (`protocol::invoke_kernel`). The party
///          thread takes `draw_context::capture()`, and the pool thread holds
///          `adopt_draws` around the kernel, so the kernel's `uniform_fill`
///          draws come from the party's source (an installed `experiment`, or
///          any other hook) in program order. `work_meter::now()` before and
///          after the kernel measures its symmetric-key blocks, random bytes,
///          and CPU on the pool thread, and `absorb_work` on the party thread
///          adds them to the party's own counters. `thread_cpu_ns()` is this
///          thread's CPU time plus the CPU time absorbed that way.
#ifndef LIBDPF_INCLUDE_DPF_THREAD_WORK_HPP__
#define LIBDPF_INCLUDE_DPF_THREAD_WORK_HPP__

#include <cstddef>
#include <cstdint>
#include <ctime>

#include "dpf/experiment_note.hpp"
#include "dpf/prg_count.hpp"
#include "dpf/random.hpp"

namespace dpf
{

namespace detail
{

inline thread_local std::uint64_t absorbed_cpu_ns_tls = 0;

inline std::uint64_t own_thread_cpu_ns() noexcept
{
    timespec ts{};
#if defined(CLOCK_THREAD_CPUTIME_ID)
    if (clock_gettime(CLOCK_THREAD_CPUTIME_ID, &ts) == 0)
        return static_cast<std::uint64_t>(ts.tv_sec) * 1000000000ull
            + static_cast<std::uint64_t>(ts.tv_nsec);
#endif
    return 0;
}

}  // namespace detail

/// @brief True when this platform has a per-thread CPU clock.
inline bool have_thread_cpu_clock() noexcept
{
#if defined(CLOCK_THREAD_CPUTIME_ID)
    timespec ts{};
    return clock_gettime(CLOCK_THREAD_CPUTIME_ID, &ts) == 0;
#else
    return false;
#endif
}

/// @brief This thread's CPU time plus the CPU time of tasks it handed off.
inline std::uint64_t thread_cpu_ns() noexcept
{
    return detail::own_thread_cpu_ns() + detail::absorbed_cpu_ns_tls;
}

/// @brief This thread's costs at one instant.
struct work_meter
{
    prg::counts sym{};
    std::uint64_t random_bytes = 0;
    std::uint64_t cpu_ns = 0;

    static work_meter now() noexcept
    {
        work_meter m;
        m.sym = prg::snapshot();
        m.random_bytes = random_bytes_count();
        m.cpu_ns = detail::own_thread_cpu_ns();
        return m;
    }
};

/// @brief Charge this thread for the costs between `before` and `after`,
///        both taken on the thread that did the work.
inline void absorb_work(const work_meter & before, const work_meter & after) noexcept
{
    prg::counts delta{};
    for (std::size_t i = 0; i < delta.size(); ++i)
        delta[i] = after.sym[i] - before.sym[i];
    prg::absorb(delta);
    detail::random_bytes_tls += after.random_bytes - before.random_bytes;
    detail::absorbed_cpu_ns_tls += after.cpu_ns - before.cpu_ns;
}

/// @brief One thread's randomness source and seed sink.
struct draw_context
{
    void (*hook)(void *, std::size_t) = nullptr;
    void * ctx = nullptr;
    detail::experiment_seed_sink * sink = nullptr;

    static draw_context capture() noexcept
    {
        draw_context c;
        c.hook = detail::uniform_bytes_hook;
        c.ctx = detail::uniform_bytes_ctx;
        c.sink = detail::experiment_seed_sink_tls;
        return c;
    }
};

/// @brief Draw from `from` on this thread until the scope ends. The thread
///        that captured `from` must not draw while this scope is open.
class adopt_draws
{
public:
    explicit adopt_draws(const draw_context & from) noexcept
      : saved_(draw_context::capture())
    {
        set(from);
    }
    ~adopt_draws() { set(saved_); }
    adopt_draws(const adopt_draws &) = delete;
    adopt_draws & operator=(const adopt_draws &) = delete;

private:
    static void set(const draw_context & c) noexcept
    {
        detail::uniform_bytes_hook = c.hook;
        detail::uniform_bytes_ctx = c.ctx;
        detail::experiment_seed_sink_tls = c.sink;
    }

    draw_context saved_;
};

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_THREAD_WORK_HPP__