/*
 * vmbench - CPU, memory, scheduling and IO benchmarks for a Windows guest.
 *
 * Deliberately has no GPU component: this guest has no display adapter, so every
 * number here is a CPU, memory, scheduling or IO proxy for interactive and
 * CPU-bound workloads.
 *
 * Build: x86_64-w64-mingw32-gcc -O2 -o vmbench.exe vmbench.c
 *
 * Timing is QueryPerformanceCounter throughout. rdtsc is deliberately avoided:
 * TSC behaviour is one of the things the fidelity work changes, so timing with it
 * would measure the instrument.
 */

#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

static double qpc_freq;

static double now(void)
{
    LARGE_INTEGER c;
    QueryPerformanceCounter(&c);
    return (double)c.QuadPart / qpc_freq;
}

static int cmp_double(const void *a, const void *b)
{
    double x = *(const double *)a, y = *(const double *)b;
    return (x > y) - (x < y);
}

static double pct(double *sorted, size_t n, double p)
{
    double idx = p * (double)(n - 1);
    size_t lo = (size_t)idx;
    if (lo + 1 >= n) return sorted[n - 1];
    double frac = idx - (double)lo;
    return sorted[lo] + frac * (sorted[lo + 1] - sorted[lo]);
}

static void pin(int cpu)
{
    if (cpu >= 0) SetThreadAffinityMask(GetCurrentThread(), (DWORD_PTR)1 << cpu);
    SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_HIGHEST);
}

/* xorshift, used to build permutations without pulling in a real PRNG */
static uint64_t rng_s = 0x243F6A8885A308D3ull;
static uint64_t rng(void)
{
    rng_s ^= rng_s << 13; rng_s ^= rng_s >> 7; rng_s ^= rng_s << 17;
    return rng_s;
}

/* ---------------------------------------------------------------- cpu ----- */
/*
 * Two numbers, because latency-sensitive workloads need both:
 *   latency   - a serial dependency chain, one op per iteration, nothing to
 *               overlap. This is what a single-threaded hot path looks like.
 *   throughput- four independent chains the core can pipeline.
 * The volatile sink stops the optimiser deleting the whole loop.
 */
static volatile uint64_t sink_u64;
static volatile double sink_f64;

static double cpu_int_latency(uint64_t iters)
{
    uint64_t x = 1;
    double t0 = now();
    for (uint64_t i = 0; i < iters; i++) {
        x = x * 6364136223846793005ull + 1442695040888963407ull;
        x ^= x >> 29;
    }
    double t1 = now();
    sink_u64 = x;
    return (double)iters / (t1 - t0) / 1e6; /* Mops/s */
}

static double cpu_int_throughput(uint64_t iters)
{
    uint64_t a = 1, b = 2, c = 3, d = 4;
    double t0 = now();
    for (uint64_t i = 0; i < iters; i++) {
        a = a * 6364136223846793005ull + 1;
        b = b * 6364136223846793005ull + 2;
        c = c * 6364136223846793005ull + 3;
        d = d * 6364136223846793005ull + 4;
    }
    double t1 = now();
    sink_u64 = a ^ b ^ c ^ d;
    return (double)(iters * 4) / (t1 - t0) / 1e6;
}

static double cpu_fp(uint64_t iters)
{
    double a = 1.0000001, b = 1.0000002, c = 1.0000003, d = 1.0000004;
    double t0 = now();
    for (uint64_t i = 0; i < iters; i++) {
        a = a * 1.0000001 + 0.0000001;
        b = b * 1.0000002 + 0.0000002;
        c = c * 1.0000003 + 0.0000003;
        d = d * 1.0000004 + 0.0000004;
    }
    double t1 = now();
    sink_f64 = a + b + c + d;
    return (double)(iters * 4) / (t1 - t0) / 1e6;
}

/* ------------------------------------------------------------ cpu_mt ----- */
struct mt_arg { int cpu; uint64_t iters; double mops; };

static DWORD WINAPI mt_worker(LPVOID p)
{
    struct mt_arg *a = (struct mt_arg *)p;
    pin(a->cpu);
    a->mops = cpu_int_throughput(a->iters);
    return 0;
}

static void bench_cpu_mt(int nthreads, uint64_t iters)
{
    struct mt_arg *args = calloc(nthreads, sizeof(*args));
    HANDLE *th = calloc(nthreads, sizeof(*th));
    for (int i = 0; i < nthreads; i++) {
        args[i].cpu = i;
        args[i].iters = iters;
        th[i] = CreateThread(NULL, 0, mt_worker, &args[i], 0, NULL);
    }
    WaitForMultipleObjects(nthreads, th, TRUE, INFINITE);
    double total = 0;
    for (int i = 0; i < nthreads; i++) { total += args[i].mops; CloseHandle(th[i]); }
    printf("cpu_mt threads=%d aggregate_Mops=%.1f per_thread_avg=%.1f\n",
           nthreads, total, total / nthreads);
    free(args); free(th);
}

/* ------------------------------------------------------------- memlat ---- */
/*
 * Pointer chase over a random single cycle, one node per 64-byte line so each
 * step is exactly one cache miss at the level being measured. A strided walk
 * would be prefetched and would measure bandwidth, not latency.
 */
struct node { struct node *next; char pad[56]; };

static double memlat(size_t bytes, uint64_t steps)
{
    size_t n = bytes / sizeof(struct node);
    if (n < 16) n = 16;
    struct node *buf = _aligned_malloc(n * sizeof(struct node), 64);
    if (!buf) return -1;
    memset(buf, 0, n * sizeof(struct node));

    size_t *perm = malloc(n * sizeof(size_t));
    for (size_t i = 0; i < n; i++) perm[i] = i;
    for (size_t i = n - 1; i > 0; i--) {          /* Fisher-Yates */
        size_t j = (size_t)(rng() % (i + 1));
        size_t t = perm[i]; perm[i] = perm[j]; perm[j] = t;
    }
    for (size_t i = 0; i < n; i++)                /* link into one cycle */
        buf[perm[i]].next = &buf[perm[(i + 1) % n]];
    free(perm);

    struct node *p = buf;
    for (uint64_t i = 0; i < n * 2; i++) p = p->next;   /* warm */

    double t0 = now();
    for (uint64_t i = 0; i < steps; i++) p = p->next;
    double t1 = now();
    sink_u64 = (uint64_t)(uintptr_t)p;

    _aligned_free(buf);
    return (t1 - t0) / (double)steps * 1e9;   /* ns per access */
}

/* ------------------------------------------------------------- membw ----- */
static double membw_read(size_t bytes)
{
    size_t n = bytes / sizeof(uint64_t);
    uint64_t *b = _aligned_malloc(n * sizeof(uint64_t), 64);
    if (!b) return -1;
    for (size_t i = 0; i < n; i++) b[i] = i;
    uint64_t s0 = 0, s1 = 0, s2 = 0, s3 = 0;
    double t0 = now();
    for (size_t i = 0; i + 3 < n; i += 4) {
        s0 += b[i]; s1 += b[i+1]; s2 += b[i+2]; s3 += b[i+3];
    }
    double t1 = now();
    sink_u64 = s0 + s1 + s2 + s3;
    _aligned_free(b);
    return (double)bytes / (t1 - t0) / 1e9;   /* GB/s */
}

static double membw_copy(size_t bytes)
{
    size_t half = bytes / 2;
    char *a = _aligned_malloc(half, 64), *b = _aligned_malloc(half, 64);
    if (!a || !b) return -1;
    memset(a, 1, half); memset(b, 2, half);
    memcpy(b, a, half);                        /* warm */
    double t0 = now();
    memcpy(b, a, half);
    double t1 = now();
    sink_u64 = (uint64_t)b[0];
    _aligned_free(a); _aligned_free(b);
    return (double)(half * 2) / (t1 - t0) / 1e9;
}

/* --------------------------------------------------------------- c2c ----- */
/*
 * Core-to-core round trip. On a 7950X this is the number that separates a
 * same-CCD pair from a cross-CCD pair, which is the whole 16-vs-32 vCPU
 * argument, so it gets measured directly rather than assumed.
 */
static volatile LONG c2c_flag;
static int c2c_cpu_b;
static uint64_t c2c_iters;

static DWORD WINAPI c2c_worker(LPVOID p)
{
    (void)p;
    pin(c2c_cpu_b);
    for (uint64_t i = 0; i < c2c_iters; i++) {
        while (InterlockedCompareExchange(&c2c_flag, 2, 1) != 1) YieldProcessor();
    }
    return 0;
}

static double c2c(int cpu_a, int cpu_b, uint64_t iters)
{
    c2c_cpu_b = cpu_b; c2c_iters = iters; c2c_flag = 0;
    HANDLE h = CreateThread(NULL, 0, c2c_worker, NULL, 0, NULL);
    pin(cpu_a);
    Sleep(50);
    double t0 = now();
    for (uint64_t i = 0; i < iters; i++) {
        InterlockedExchange(&c2c_flag, 1);
        while (InterlockedCompareExchange(&c2c_flag, 0, 2) != 2) YieldProcessor();
    }
    double t1 = now();
    WaitForSingleObject(h, 5000);
    CloseHandle(h);
    return (t1 - t0) / (double)iters * 1e9;   /* ns round trip */
}

/* ------------------------------------------------------------ jitter ----- */
/*
 * The scheduling-tail proxy, and the most important number in the suite.
 *
 * Spin reading the clock as fast as possible and record every gap between
 * consecutive reads. On an idle, well-behaved system the gaps are tiny and
 * uniform. Any gap far above the floor means this thread was NOT RUNNING -
 * preempted by the host, stalled behind a VMEXIT, or waiting on the emulator
 * thread. Mean throughput averages all of that away; the tail is what shows up
 * as hitching in an interactive session, so the output is percentiles, not an
 * average.
 */
static void jitter(int cpu, double seconds, size_t max_samples)
{
    pin(cpu);
    double *gaps = malloc(max_samples * sizeof(double));
    if (!gaps) { printf("jitter alloc failed\n"); return; }

    double t_end = now() + seconds;
    double prev = now();
    size_t n = 0;
    while (n < max_samples) {
        double t = now();
        double d = (t - prev) * 1e6;          /* microseconds */
        prev = t;
        gaps[n++] = d;
        if (t > t_end) break;
    }

    qsort(gaps, n, sizeof(double), cmp_double);
    printf("jitter cpu=%d samples=%zu floor_us=%.3f p50=%.3f p99=%.3f p99.9=%.3f p99.99=%.3f max_us=%.1f\n",
           cpu, n, gaps[0], pct(gaps, n, 0.50), pct(gaps, n, 0.99),
           pct(gaps, n, 0.999), pct(gaps, n, 0.9999), gaps[n - 1]);

    /* count of stalls over thresholds - these are the visible hitches */
    size_t o10 = 0, o100 = 0, o1000 = 0;
    for (size_t i = 0; i < n; i++) {
        if (gaps[i] > 10.0)   o10++;
        if (gaps[i] > 100.0)  o100++;
        if (gaps[i] > 1000.0) o1000++;
    }
    printf("jitter cpu=%d stalls_over_10us=%zu over_100us=%zu over_1ms=%zu\n",
           cpu, o10, o100, o1000);
    free(gaps);
}

/* ------------------------------------------------------------- diskio ---- */
/*
 * Unbuffered file IO, so the guest's page cache is out of the way and what gets
 * measured is the virtual disk path: guest driver -> QEMU device model ->
 * emulator/IO thread -> host file. 4K random reads are the number that matters
 * for scattered small-block access; sequential is the easy case.
 *
 * FILE_FLAG_NO_BUFFERING requires sector-aligned buffers, offsets and lengths.
 */
#define IO_FILE   "C:\\Users\\User\\vmbench_io.tmp"
#define IO_SIZE   (512u << 20)      /* 512 MiB */
#define IO_BLK    4096

static void diskio(void)
{
    HANDLE h = CreateFileA(IO_FILE, GENERIC_READ | GENERIC_WRITE,
                           0, NULL, CREATE_ALWAYS,
                           FILE_FLAG_NO_BUFFERING | FILE_FLAG_WRITE_THROUGH, NULL);
    if (h == INVALID_HANDLE_VALUE) { printf("diskio open failed %lu\n", GetLastError()); return; }

    void *buf = _aligned_malloc(1u << 20, 4096);
    memset(buf, 0xA5, 1u << 20);

    /* lay the file down, and measure sequential write in the same pass */
    DWORD got;
    double t0 = now();
    for (size_t off = 0; off < IO_SIZE; off += (1u << 20))
        WriteFile(h, buf, 1u << 20, &got, NULL);
    FlushFileBuffers(h);
    double t1 = now();
    printf("diskio seq_write_MBs=%.1f\n", (double)IO_SIZE / (t1 - t0) / 1e6);

    /* sequential read */
    LARGE_INTEGER z = {0}; SetFilePointerEx(h, z, NULL, FILE_BEGIN);
    t0 = now();
    for (size_t off = 0; off < IO_SIZE; off += (1u << 20))
        ReadFile(h, buf, 1u << 20, &got, NULL);
    t1 = now();
    printf("diskio seq_read_MBs=%.1f\n", (double)IO_SIZE / (t1 - t0) / 1e6);

    /* 4K random read - the one that shows device-model latency */
    const int N = 20000;
    size_t blocks = IO_SIZE / IO_BLK;
    t0 = now();
    for (int i = 0; i < N; i++) {
        LARGE_INTEGER p;
        p.QuadPart = (LONGLONG)((rng() % blocks) * IO_BLK);
        SetFilePointerEx(h, p, NULL, FILE_BEGIN);
        ReadFile(h, buf, IO_BLK, &got, NULL);
    }
    t1 = now();
    double us = (t1 - t0) / N * 1e6;
    printf("diskio rand4k_IOPS=%.0f rand4k_lat_us=%.1f\n", N / (t1 - t0), us);

    _aligned_free(buf);
    CloseHandle(h);
    DeleteFileA(IO_FILE);
}

/* -------------------------------------------------------------- clock ---- */
/*
 * How expensive is asking the time? Software that polls the clock in a tight
 * loop calls QPC thousands of times a second, so if the guest's clock source
 * traps to the hypervisor it shows up as a flat tax on everything and as a
 * longer scheduling tail.
 *
 * Bare metal Windows serves QueryPerformanceCounter from the TSC in tens of ns.
 * If this reports hundreds of ns or microseconds, Windows has fallen back to a
 * clock that costs a VMEXIT per read.
 */
static void bench_clock(void)
{
    pin(0);
    const uint64_t N = 2000000;
    LARGE_INTEGER c;
    volatile uint64_t acc = 0;

    for (uint64_t i = 0; i < 10000; i++) { QueryPerformanceCounter(&c); acc += c.QuadPart; }

    double t0 = now();
    for (uint64_t i = 0; i < N; i++) { QueryPerformanceCounter(&c); acc += (uint64_t)c.QuadPart; }
    double t1 = now();
    double qpc_ns = (t1 - t0) / (double)N * 1e9;

    t0 = now();
    for (uint64_t i = 0; i < N; i++) acc += __rdtsc();
    t1 = now();
    double tsc_ns = (t1 - t0) / (double)N * 1e9;

    t0 = now();
    for (uint64_t i = 0; i < N; i++) acc += GetTickCount64();
    t1 = now();
    double gtc_ns = (t1 - t0) / (double)N * 1e9;

    sink_u64 = acc;
    printf("clock qpc_ns=%.1f rdtsc_ns=%.1f gettickcount_ns=%.1f qpc_freq=%.0f\n",
           qpc_ns, tsc_ns, gtc_ns, qpc_freq);
}

/* -------------------------------------------------------------- main ----- */
static void usage(void)
{
    printf("usage: vmbench <diskio|clock|cpu|cpu_mt N|memlat|membw|c2c A B|jitter CPU SECS|all>\n");
}

int main(int argc, char **argv)
{
    LARGE_INTEGER f;
    QueryPerformanceFrequency(&f);
    qpc_freq = (double)f.QuadPart;

    SYSTEM_INFO si;
    GetSystemInfo(&si);
    printf("# vmbench  qpc_freq=%.0f logical_cpus=%lu groups=%u\n",
           qpc_freq, (unsigned long)si.dwNumberOfProcessors,
           (unsigned)GetActiveProcessorGroupCount());

    SetPriorityClass(GetCurrentProcess(), HIGH_PRIORITY_CLASS);

    if (argc < 2) { usage(); return 1; }
    const char *cmd = argv[1];
    int all = strcmp(cmd, "all") == 0;

    if (all || !strcmp(cmd, "clock")) {
        bench_clock();
    }

    if (all || !strcmp(cmd, "diskio")) {
        diskio();
    }

    if (all || !strcmp(cmd, "cpu")) {
        pin(0);
        double best_l = 0, best_t = 0, best_f = 0;
        for (int r = 0; r < 5; r++) {
            double l = cpu_int_latency(200000000ull);
            double t = cpu_int_throughput(100000000ull);
            double fp = cpu_fp(100000000ull);
            if (l > best_l) best_l = l;
            if (t > best_t) best_t = t;
            if (fp > best_f) best_f = fp;
        }
        printf("cpu int_latency_Mops=%.1f int_throughput_Mops=%.1f fp_Mops=%.1f\n",
               best_l, best_t, best_f);
    }

    if (all || !strcmp(cmd, "cpu_mt")) {
        int n = (!all && argc > 2) ? atoi(argv[2]) : (int)si.dwNumberOfProcessors;
        bench_cpu_mt(n, 100000000ull);
        if (all) { bench_cpu_mt(1, 100000000ull); bench_cpu_mt(8, 100000000ull); }
    }

    if (all || !strcmp(cmd, "memlat")) {
        pin(0);
        size_t sizes[] = { 32u<<10, 512u<<10, 8u<<20, 64u<<20, 256u<<20 };
        const char *names[] = { "L1_32K", "L2_512K", "L3_8M", "DRAM_64M", "DRAM_256M" };
        for (int i = 0; i < 5; i++) {
            double best = 1e9;
            for (int r = 0; r < 3; r++) {
                double v = memlat(sizes[i], 20000000ull);
                if (v > 0 && v < best) best = v;
            }
            printf("memlat %s ns=%.2f\n", names[i], best);
        }
    }

    if (all || !strcmp(cmd, "membw")) {
        pin(0);
        double br = 0, bc = 0;
        for (int r = 0; r < 3; r++) {
            double x = membw_read(256u<<20);  if (x > br) br = x;
            double y = membw_copy(256u<<20);  if (y > bc) bc = y;
        }
        printf("membw read_GBs=%.2f copy_GBs=%.2f\n", br, bc);
    }

    if (all || !strcmp(cmd, "c2c")) {
        if (!all && argc > 3) {
            printf("c2c %s->%s ns=%.1f\n", argv[2], argv[3],
                   c2c(atoi(argv[2]), atoi(argv[3]), 200000));
        } else {
            int nc = (int)si.dwNumberOfProcessors;
            printf("c2c 0->1_ns=%.1f\n", c2c(0, 1, 200000));
            if (nc > 8)  printf("c2c 0->8_ns=%.1f\n",  c2c(0, 8, 200000));
            if (nc > 16) printf("c2c 0->16_ns=%.1f\n", c2c(0, 16, 200000));
            if (nc > 24) printf("c2c 0->24_ns=%.1f\n", c2c(0, 24, 200000));
        }
    }

    if (all || !strcmp(cmd, "jitter")) {
        int cpu   = (!all && argc > 2) ? atoi(argv[2]) : 0;
        double se = (!all && argc > 3) ? atof(argv[3]) : 10.0;
        jitter(cpu, se, 20000000);
    }

    return 0;
}