/* Multicore hash/RC4 offload benchmark for the Phase 6 CPU budget. */
#include "naut/hash.h"
#include "naut/rc4.h"
#include "naut/worker.h"

#include <poll.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#define JOB_BYTES (1u << 20)
#define JOB_COUNT 64

typedef enum { BENCH_SHA1, BENCH_SHA256, BENCH_RC4 } bench_kind;

typedef struct {
    naut_job base;
    bench_kind kind;
    uint8_t *data;
    uint8_t digest[NAUT_SHA256_LEN];
} bench_job;

static double now_seconds(void) {
    struct timespec time;
    clock_gettime(CLOCK_MONOTONIC, &time);
    return time.tv_sec + time.tv_nsec * 1e-9;
}

static void run_job(naut_job *base) {
    bench_job *job = base->context;
    if (job->kind == BENCH_SHA1) {
        naut_sha1(job->data, JOB_BYTES, job->digest);
    } else if (job->kind == BENCH_SHA256) {
        naut_sha256(job->data, JOB_BYTES, job->digest);
    } else {
        static const uint8_t key[20] = {
            0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
            10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
        };
        naut_rc4 rc4;
        naut_rc4_init(&rc4, key, sizeof key, 1024);
        naut_rc4_xor(&rc4, job->data, JOB_BYTES);
    }
    base->result = NAUT_OK;
}

static double run(naut_worker_pool *pool, bench_job *jobs,
                  bench_kind kind, int rounds) {
    for (int i = 0; i < JOB_COUNT; i++) jobs[i].kind = kind;
    int total_completed = 0;
    double start = now_seconds();
    for (int round = 0; round < rounds; round++) {
        for (int i = 0; i < JOB_COUNT; i++) {
            while (!naut_worker_submit(pool, &jobs[i].base))
                sched_yield();
        }
        int completed = 0;
        while (completed < JOB_COUNT) {
            naut_job *base;
            if (naut_worker_complete(pool, &base)) {
                (void)base;
                completed++;
                total_completed++;
                continue;
            }
            struct pollfd pfd = {
                .fd = naut_worker_eventfd(pool),
                .events = POLLIN,
            };
            if (poll(&pfd, 1, 5000) <= 0) break;
            uint64_t count;
            (void)read(pfd.fd, &count, sizeof count);
        }
        if (completed != JOB_COUNT) break;
    }
    double seconds = now_seconds() - start;
    return ((double)total_completed * JOB_BYTES / 1e9) / seconds;
}

int main(int argc, char **argv) {
    int threads = argc > 1 ? atoi(argv[1]) : 8;
    int rounds = argc > 2 ? atoi(argv[2]) : 16;
    if (threads < 1 || rounds < 1) return 2;
    naut_worker_pool *pool =
        naut_worker_pool_create((uint32_t)threads, 128, -1);
    if (!pool) return 1;

    bench_job *jobs = calloc(JOB_COUNT, sizeof(*jobs));
    uint8_t *slab = aligned_alloc(NAUT_PAGE, JOB_COUNT * JOB_BYTES);
    if (!jobs || !slab) return 1;
    memset(slab, 0xa5, JOB_COUNT * JOB_BYTES);
    for (int i = 0; i < JOB_COUNT; i++) {
        jobs[i].base.run = run_job;
        jobs[i].base.context = &jobs[i];
        jobs[i].data = slab + (size_t)i * JOB_BYTES;
    }

    double sha1 = run(pool, jobs, BENCH_SHA1, rounds);
    double sha256 = run(pool, jobs, BENCH_SHA256, rounds);
    double rc4 = run(pool, jobs, BENCH_RC4, rounds);
    printf("%d workers, %.2f GiB processed per primitive\n",
           threads, (double)JOB_COUNT * rounds * JOB_BYTES / (1u << 30));
    printf("  sha1    %.2f GB/s\n", sha1);
    printf("  sha256  %.2f GB/s\n", sha256);
    printf("  rc4     %.2f GB/s\n", rc4);

    free(slab);
    free(jobs);
    naut_worker_pool_destroy(pool);
    return 0;
}