/// @file test/profile/harness.hpp
/// @brief Timing loop shared by the in-process profile drivers.
#ifndef LIBDPF_TEST_PROFILE_HARNESS_HPP__
#define LIBDPF_TEST_PROFILE_HARNESS_HPP__

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstring>
#include <functional>
#include <iostream>
#include <optional>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

#include <chrono>

#include "dpf/prg_count.hpp"

namespace profile
{

struct sample
{
    std::uint64_t sink = 0;
    std::uint64_t out_bytes = 0;
    /// @brief Set when the case instruments the named cost; blank in TSV otherwise.
    std::optional<std::uint64_t> prg_evals;
    std::optional<std::uint64_t> preprocess_bytes;
    std::optional<std::uint64_t> alloc_bytes;
    std::optional<std::uint64_t> logical_bytes;
};

inline sample operator+(sample a, sample b)
{
    a.sink ^= b.sink + 0x9e3779b97f4a7c15ull;
    a.out_bytes += b.out_bytes;
    if (a.prg_evals || b.prg_evals)
        a.prg_evals = a.prg_evals.value_or(0) + b.prg_evals.value_or(0);
    if (a.preprocess_bytes || b.preprocess_bytes)
        a.preprocess_bytes = a.preprocess_bytes.value_or(0)
            + b.preprocess_bytes.value_or(0);
    if (a.alloc_bytes || b.alloc_bytes)
        a.alloc_bytes = a.alloc_bytes.value_or(0) + b.alloc_bytes.value_or(0);
    if (a.logical_bytes || b.logical_bytes)
        a.logical_bytes = a.logical_bytes.value_or(0)
            + b.logical_bytes.value_or(0);
    return a;
}

/// @brief Keep `word` live. `out_bytes` is reported, not timed on its own.
inline sample touch_word(std::uint64_t word, std::uint64_t out_bytes = 0)
{
    sample s;
    s.sink = word;
    s.out_bytes = out_bytes;
    asm volatile("" : "+r"(s.sink)::"memory");
    return s;
}

/// @brief Fold the ends of a buffer and publish a compiler barrier over it.
template <typename Buf>
sample touch_buf(const Buf & buf)
{
    sample s;
    using value_type = typename Buf::value_type;
    s.out_bytes = buf.size() * sizeof(value_type);
    s.sink = buf.size();
    if (buf.size() != 0)
    {
        std::uint64_t a = 0;
        std::uint64_t b = 0;
        const std::size_t n = sizeof(value_type) < sizeof(a) ? sizeof(value_type) : sizeof(a);
        std::memcpy(&a, &buf[0], n);
        std::memcpy(&b, &buf[buf.size() - 1], n);
        s.sink ^= a ^ (b + buf.size());
    }
    if (buf.size() != 0)
        asm volatile("" : "+r"(s.sink) : "r"(buf.data()) : "memory");
    else
        asm volatile("" : "+r"(s.sink)::"memory");
    return s;
}

template <typename T>
sample touch_vec(const std::vector<T> & v)
{
    std::uint64_t h = v.size();
    for (const T & x : v)
    {
        std::uint64_t w = 0;
        if constexpr (std::is_integral_v<T>)
            w = static_cast<std::uint64_t>(x);
        else
        {
            const std::size_t n = sizeof(T) < sizeof(w) ? sizeof(T) : sizeof(w);
            std::memcpy(&w, &x, n);
        }
        h ^= w + 0x9e3779b97f4a7c15ull;
        h *= 0x100000001b3ull;
    }
    return touch_word(h, v.size() * sizeof(T));
}

template <typename T, std::size_t N>
sample touch_arr(const std::array<T, N> & a)
{
    std::uint64_t h = N;
    for (const T & x : a)
    {
        std::uint64_t w = 0;
        if constexpr (std::is_integral_v<T>)
            w = static_cast<std::uint64_t>(x);
        else
            w = x ? 1u : 0u;
        h = (h << 1) ^ w;
    }
    return touch_word(h, N * sizeof(T));
}

/// @brief Reset the PRG counter, run `fn`, and attach the delta to the sample.
template <typename Fn>
sample with_prg(Fn && fn)
{
    dpf::prg::reset_eval_count();
    sample s = std::forward<Fn>(fn)();
    s.prg_evals = dpf::prg::eval_count();
    return s;
}

inline sample with_costs(sample s, std::uint64_t preprocess, std::uint64_t alloc,
    std::uint64_t logical)
{
    s.preprocess_bytes = preprocess;
    s.alloc_bytes = alloc;
    s.logical_bytes = logical;
    return s;
}

inline std::string opt_field(const std::optional<std::uint64_t> & v)
{
    return v ? std::to_string(*v) : std::string{};
}

inline std::uint64_t ticks()
{
#if defined(__x86_64__) || defined(__i386__)
    unsigned lo = 0;
    unsigned hi = 0;
    asm volatile("rdtscp" : "=a"(lo), "=d"(hi)::"rcx");
    return (static_cast<std::uint64_t>(hi) << 32) | lo;
#else
    return 0;
#endif
}

struct work
{
    std::string name;
    std::uint64_t items = 1;
    std::function<sample()> fn;
    /// @brief Top-level group: `eval` or `grotto`.
    std::string family;
    /// @brief Piecewise rerun unit, such as `interval` or `horner`.
    std::string slice;
    /// @brief `std` is the default matrix. `heavy` is opt-in.
    std::string tier = "std";
};

inline work make_work(std::string family, std::string slice, std::string name,
    std::uint64_t items, std::function<sample()> fn, std::string tier = "std")
{
    work w;
    w.family = std::move(family);
    w.slice = std::move(slice);
    w.name = std::move(name);
    w.items = items;
    w.fn = std::move(fn);
    w.tier = std::move(tier);
    return w;
}

struct parsed
{
    std::uint64_t repeat = 1;
    std::uint64_t warmup = 0;
    bool list = false;
    std::vector<std::string> only;
    std::vector<std::string> families;
    std::vector<std::string> slices;
    std::string tier;
};

inline parsed parse_args(int argc, char ** argv, std::uint64_t repeat,
    std::uint64_t warmup, const char * usage)
{
    parsed o;
    o.repeat = repeat;
    o.warmup = warmup;
    for (int i = 1; i < argc; ++i)
    {
        const std::string a = argv[i];
        auto need = [&](const char * flag) {
            if (i + 1 >= argc)
            {
                std::cerr << "missing value for " << flag << "\n";
                std::exit(2);
            }
            return std::string(argv[++i]);
        };
        if (a == "--list")
            o.list = true;
        else if (a == "--repeat")
            o.repeat = std::stoull(need("--repeat"));
        else if (a == "--warmup")
            o.warmup = std::stoull(need("--warmup"));
        else if (a == "--case")
            o.only.push_back(need("--case"));
        else if (a == "--family")
            o.families.push_back(need("--family"));
        else if (a == "--slice")
            o.slices.push_back(need("--slice"));
        else if (a == "--tier")
            o.tier = need("--tier");
        else if (a == "--help" || a == "-h")
        {
            std::cout << usage;
            std::exit(0);
        }
        else
        {
            std::cerr << "unknown argument: " << a << "\n" << usage;
            std::exit(2);
        }
    }
    if (o.repeat == 0)
        o.repeat = 1;
    return o;
}

inline bool contains(const std::vector<std::string> & hay, const std::string & needle)
{
    return std::find(hay.begin(), hay.end(), needle) != hay.end();
}

inline bool matches(const parsed & opt, const work & w)
{
    if (!opt.tier.empty() && opt.tier != "all" && w.tier != opt.tier)
        return false;
    if (!opt.families.empty() && !contains(opt.families, w.family))
        return false;
    if (!opt.slices.empty() && !contains(opt.slices, w.slice))
        return false;
    if (!opt.only.empty() && !contains(opt.only, w.name))
        return false;
    return true;
}

inline int run_works(const parsed & opt, const std::vector<work> & all)
{
    for (const auto & name : opt.only)
    {
        bool found = false;
        for (const auto & w : all)
            found = found || w.name == name;
        if (!found)
        {
            std::cerr << "unknown case: " << name << "\n";
            return 2;
        }
    }

    std::vector<const work *> chosen;
    for (const auto & w : all)
    {
        if (matches(opt, w))
            chosen.push_back(&w);
    }
    if (chosen.empty())
    {
        std::cerr << "no cases match the requested family/slice/tier/case\n";
        return 2;
    }

    if (opt.list)
    {
        std::cout << "family\tslice\tcase\titems\ttier\n";
        for (const work * w : chosen)
            std::cout << w->family << '\t' << w->slice << '\t' << w->name
                      << '\t' << w->items << '\t' << w->tier << '\n';
        return 0;
    }

    std::cout
        << "family\tslice\tcase\titems\trepeat\twarmup\tavg_ns\tmin_ns\tmax_ns\t"
        << "avg_cycles\tper_item_ns\tout_bytes\tsink\t"
        << "prg_evals\tpreprocess_bytes\talloc_bytes\tlogical_bytes\tlayout_waste\n";
    int fails = 0;
    std::uint64_t all_sink = 0;
    for (const work * w : chosen)
    {
        try
        {
            for (std::uint64_t i = 0; i < opt.warmup; ++i)
                all_sink ^= w->fn().sink;

            std::uint64_t total_ns = 0;
            std::uint64_t min_ns = ~std::uint64_t{0};
            std::uint64_t max_ns = 0;
            std::uint64_t total_cycles = 0;
            sample last{};
            for (std::uint64_t i = 0; i < opt.repeat; ++i)
            {
                const auto c0 = ticks();
                const auto t0 = std::chrono::steady_clock::now();
                const sample s = w->fn();
                const auto t1 = std::chrono::steady_clock::now();
                const auto c1 = ticks();
                const auto ns = static_cast<std::uint64_t>(
                    std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count());
                total_ns += ns;
                min_ns = std::min(min_ns, ns);
                max_ns = std::max(max_ns, ns);
                total_cycles += c1 - c0;
                // Last sample, not an xor across repeats: identical samples
                // would cancel and the column would read as zero.
                last = s;
                all_sink ^= s.sink + i;
            }
            const std::uint64_t avg_ns = total_ns / opt.repeat;
            const std::uint64_t avg_cycles = total_cycles / opt.repeat;
            const std::uint64_t per_item = w->items == 0 ? avg_ns : avg_ns / w->items;
            std::string layout_waste;
            if (last.alloc_bytes && last.logical_bytes
                && *last.alloc_bytes >= *last.logical_bytes)
            {
                layout_waste = std::to_string(*last.alloc_bytes - *last.logical_bytes);
            }
            std::cout << w->family << '\t' << w->slice << '\t' << w->name << '\t'
                      << w->items << '\t' << opt.repeat << '\t'
                      << opt.warmup << '\t' << avg_ns << '\t' << min_ns << '\t' << max_ns
                      << '\t' << avg_cycles << '\t' << per_item << '\t' << last.out_bytes
                      << '\t' << last.sink << '\t'
                      << opt_field(last.prg_evals) << '\t'
                      << opt_field(last.preprocess_bytes) << '\t'
                      << opt_field(last.alloc_bytes) << '\t'
                      << opt_field(last.logical_bytes) << '\t'
                      << layout_waste << '\n';
        }
        catch (const std::exception & ex)
        {
            std::cerr << w->name << " failed: " << ex.what() << "\n";
            ++fails;
        }
    }
    std::cout << "sink\t" << all_sink << "\n";
    return fails == 0 ? 0 : 1;
}

}  // namespace profile

#endif  // LIBDPF_TEST_PROFILE_HARNESS_HPP__