// Dispatch microbenchmark — PERF_PLAN B1 and B2.
//
// B1 asks what a single ThreadPool dispatch costs. B2 asks whether a worker
// actually sleeps per item, because the whole spin-window hypothesis (B5)
// depends on the answer: if workers are not sleeping, a spin window buys
// nothing and drops off the list.
//
// Both are answered here without touching the library. Sleeping is inferred
// from ru_nvcsw — a thread blocking on a condition variable books a voluntary
// context switch — so `vcsw/task` near 1.0 means a sleep per dispatch and near
// 0 means the worker never went to sleep at all.
//
// Three modes, because "the cost of a dispatch" is three different numbers:
//
//   latency — one task in flight, pool idle in between. The worker is asleep
//             at every submission, so this is dispatch cost *including* a
//             wake. Worst case, and the case a spin window would attack.
//
//   batch   — submit K no-op tasks flat out, then drain. The worker is never
//             idle, so this is the amortised floor: queue and heap operations
//             with no wake at all. Reports the producer-side submit() cost
//             separately from end-to-end throughput.
//
//   steady  — the task resubmits its successor, one in flight, each doing
//             --work-us of work. This is what a KPN node actually does:
//             fire_once processes a token and resubmits. On ThreadPool(1) the
//             worker resubmits to its own queue; on ThreadPool(4) round-robin
//             hands the task to a *different* worker, which may be asleep.
//             That difference is the fanout cost wide-4 pays ~5x per item.
//
// Usage: ./bench_dispatch [--threads=1,2,4] [--mode=latency,batch,steady]
//                         [--tasks=200000] [--work-us=0] [--reps=5] [--warmup=1]

#include <kpn/kpn.hpp>

#include "bench_env.hpp"

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdio>
#include <cstdlib>
#include <mutex>
#include <string>
#include <vector>

using namespace kpn;
using sclock = std::chrono::steady_clock;

struct Opts {
    std::vector<int> threads {1, 2, 4};
    std::vector<std::string> modes {"latency", "batch", "steady"};
    long tasks   = 200000;
    int  work_us = 0;
    int  reps    = 5;
    int  warmup  = 1;
};

static Opts g_opts;

static void busy_us(int us) {
    if (us <= 0) return;
    auto end = sclock::now() + std::chrono::microseconds(us);
    while (sclock::now() < end);
}

struct Sample {
    double ns_per_dispatch  = 0;   // end-to-end, minus the work payload
    double submit_ns        = 0;   // producer side only (batch mode)
    double vcsw_per_task    = 0;   // B2: sleeps per dispatch
    double ivcsw_per_task   = 0;
};

// ── latency: one task at a time, worker asleep between submissions ────────────

static Sample run_latency(int threads, long tasks) {
    ThreadPool pool(threads);
    pool.start();

    std::mutex mx;
    std::condition_variable cv;
    bool done = false;

    bench::RusageDelta ru; ru.start();
    auto t0 = sclock::now();
    for (long i = 0; i < tasks; ++i) {
        { std::lock_guard lk(mx); done = false; }
        pool.submit([&] {
            busy_us(g_opts.work_us);
            { std::lock_guard lk(mx); done = true; }
            cv.notify_one();
        });
        std::unique_lock lk(mx);
        cv.wait(lk, [&] { return done; });
    }
    auto t1 = sclock::now();
    Sample s;
    long iv = 0, vc = 0;
    ru.finish(iv, vc);
    pool.stop();

    double elapsed_ns = std::chrono::duration<double, std::nano>(t1 - t0).count();
    s.ns_per_dispatch = elapsed_ns / tasks - g_opts.work_us * 1000.0;
    // The requesting thread blocks once per task too, so it books a voluntary
    // switch of its own; halve to attribute per side rather than per process.
    s.vcsw_per_task  = static_cast<double>(vc) / tasks / 2.0;
    s.ivcsw_per_task = static_cast<double>(iv) / tasks;
    return s;
}

// ── batch: submit flat out, drain once. No wake in the steady state ───────────

static Sample run_batch(int threads, long tasks) {
    ThreadPool pool(threads);
    pool.start();

    std::atomic<long> ran{0};

    bench::RusageDelta ru; ru.start();
    auto t0 = sclock::now();
    for (long i = 0; i < tasks; ++i)
        pool.submit([&] {
            busy_us(g_opts.work_us);
            ran.fetch_add(1, std::memory_order_relaxed);
        });
    auto t_submitted = sclock::now();
    pool.drain();
    auto t1 = sclock::now();
    Sample s;
    long iv = 0, vc = 0;
    ru.finish(iv, vc);
    pool.stop();

    if (ran.load() != tasks)
        std::fprintf(stderr, "WARNING: batch ran %ld of %ld tasks\n",
                     ran.load(), tasks);

    double elapsed_ns = std::chrono::duration<double, std::nano>(t1 - t0).count();
    s.ns_per_dispatch = elapsed_ns / tasks - g_opts.work_us * 1000.0;
    s.submit_ns = std::chrono::duration<double, std::nano>(t_submitted - t0).count() / tasks;
    s.vcsw_per_task  = static_cast<double>(vc) / tasks;
    s.ivcsw_per_task = static_cast<double>(iv) / tasks;
    return s;
}

// ── steady: the task resubmits its successor, as fire_once does ──────────────

static Sample run_steady(int threads, long tasks) {
    ThreadPool pool(threads);
    pool.start();

    std::mutex mx;
    std::condition_variable cv;
    std::atomic<long> count{0};
    bool  finished = false;
    // Recursive submission: hold the chain in a std::function so the task can
    // resubmit itself. Captured by reference; it outlives the drain below.
    //
    // The counter is atomic rather than mutex-guarded so that this loop
    // measures the pool's dispatch path and not a lock of the benchmark's own.
    std::function<void()> step = [&] {
        busy_us(g_opts.work_us);
        long n = count.fetch_add(1, std::memory_order_relaxed) + 1;
        if (n < tasks) {
            pool.submit(step);
        } else {
            { std::lock_guard lk(mx); finished = true; }
            cv.notify_one();
        }
    };

    bench::RusageDelta ru; ru.start();
    auto t0 = sclock::now();
    pool.submit(step);
    {
        std::unique_lock lk(mx);
        cv.wait(lk, [&] { return finished; });
    }
    auto t1 = sclock::now();
    Sample s;
    long iv = 0, vc = 0;
    ru.finish(iv, vc);
    pool.stop();

    double elapsed_ns = std::chrono::duration<double, std::nano>(t1 - t0).count();
    s.ns_per_dispatch = elapsed_ns / tasks - g_opts.work_us * 1000.0;
    s.vcsw_per_task  = static_cast<double>(vc) / tasks;
    s.ivcsw_per_task = static_cast<double>(iv) / tasks;
    return s;
}

// ── driver ────────────────────────────────────────────────────────────────────

static void run_row(const std::string& mode, int threads, long tasks) {
    auto once = [&] {
        if (mode == "latency") return run_latency(threads, tasks);
        if (mode == "batch")   return run_batch(threads, tasks);
        return run_steady(threads, tasks);
    };

    for (int i = 0; i < g_opts.warmup; ++i) (void)once();

    std::vector<double> ns, sub, vcsw, ivcsw;
    for (int i = 0; i < g_opts.reps; ++i) {
        Sample s = once();
        ns.push_back(s.ns_per_dispatch);
        sub.push_back(s.submit_ns);
        vcsw.push_back(s.vcsw_per_task);
        ivcsw.push_back(s.ivcsw_per_task);
    }

    const double med = bench::percentile(ns, 0.5);
    const double q1  = bench::percentile(ns, 0.25);
    const double q3  = bench::percentile(ns, 0.75);
    const double iqr = med > 0 ? 100.0 * (q3 - q1) / med : 0.0;
    const double sleeps = bench::percentile(vcsw, 0.5);

    std::fprintf(stderr, "%-9s %-8d %-8d %-10ld %-12.0f %-7.1f %-11.0f %-10.2f %-10.2f\n",
                 mode.c_str(), threads, g_opts.work_us, tasks, med, iqr,
                 bench::percentile(sub, 0.5), sleeps,
                 bench::percentile(ivcsw, 0.5));
    // Column names deliberately match bench_pipeline's key columns so that
    // scripts/bench_repro_check.py can gate this benchmark too.
    std::printf("%s,%d,%d,%d,%ld,%d,%.1f,%.2f,%.1f,%.3f,%.3f\n",
                mode.c_str(), threads, g_opts.work_us, threads, tasks,
                g_opts.reps, med, iqr, bench::percentile(sub, 0.5),
                sleeps, bench::percentile(ivcsw, 0.5));
    std::fflush(stdout);
}

static std::vector<int> parse_int_list(const char* s) {
    std::vector<int> out;
    const char* p = s;
    while (*p) {
        char* end = nullptr;
        long v = std::strtol(p, &end, 10);
        if (end == p) break;
        out.push_back(static_cast<int>(v));
        p = end;
        while (*p == ',' || *p == ' ') ++p;
    }
    return out;
}

static std::vector<std::string> parse_word_list(const std::string& s) {
    std::vector<std::string> out;
    std::size_t pos = 0;
    while (pos <= s.size()) {
        std::size_t c = s.find(',', pos);
        if (c == std::string::npos) c = s.size();
        if (c > pos) out.push_back(s.substr(pos, c - pos));
        pos = c + 1;
    }
    return out;
}

static void usage() {
    std::fprintf(stderr,
        "usage: bench_dispatch [options]\n"
        "  --threads=1,2,4                pool sizes\n"
        "  --mode=latency,batch,steady    which measurements to run\n"
        "  --tasks=200000                 dispatches per repetition\n"
        "  --work-us=0                    payload per task\n"
        "  --reps=5 --warmup=1\n");
}

int main(int argc, char** argv) {
    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        auto eq = a.find('=');
        std::string key = a.substr(0, eq);
        std::string val = eq == std::string::npos ? "" : a.substr(eq + 1);

        if      (key == "--help" || key == "-h") { usage(); return 0; }
        else if (key == "--threads") g_opts.threads = parse_int_list(val.c_str());
        else if (key == "--mode")    g_opts.modes   = parse_word_list(val);
        else if (key == "--tasks")   g_opts.tasks   = std::atol(val.c_str());
        else if (key == "--work-us") g_opts.work_us = std::atoi(val.c_str());
        else if (key == "--reps")    g_opts.reps    = std::atoi(val.c_str());
        else if (key == "--warmup")  g_opts.warmup  = std::atoi(val.c_str());
        else { std::fprintf(stderr, "unknown option: %s\n", a.c_str()); usage(); return 2; }
    }
    if (g_opts.reps < 1) g_opts.reps = 1;
    if (g_opts.warmup < 0) g_opts.warmup = 0;

    char cfg[160];
    std::snprintf(cfg, sizeof cfg, "tasks=%ld work_us=%d reps=%d warmup=%d",
                  g_opts.tasks, g_opts.work_us, g_opts.reps, g_opts.warmup);
    bench::print_environment(cfg);

    std::fprintf(stderr, "\n%-9s %-8s %-8s %-10s %-12s %-7s %-11s %-10s %-10s\n",
                 "mode", "threads", "work_us", "tasks", "ns/dispatch", "iqr%",
                 "submit_ns", "vcsw/task", "ivcsw/task");
    std::fprintf(stderr, "%s\n", std::string(96, '-').c_str());
    std::printf("topology,size,work_us,threads,items,reps,ns_per_dispatch,"
                "iqr_pct,submit_ns,vcsw_per_task,ivcsw_per_task\n");

    // latency is a round trip per task, so it is far slower per dispatch than
    // the other modes; scale it down rather than run for minutes.
    for (const auto& mode : g_opts.modes)
        for (int t : g_opts.threads) {
            long tasks = mode == "latency"
                       ? std::max(2000L, g_opts.tasks / 20)
                       : g_opts.tasks;
            run_row(mode, t, tasks);
        }
}