#import <Foundation/Foundation.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <unistd.h>

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

static double get_disk0_mib(void) {
    FILE *fp = popen("iostat -Id", "r");
    if (!fp) return 0.0;
    char line[512], last[512] = {0};
    while (fgets(line, sizeof(line), fp)) {
        if (strlen(line) > 10) strcpy(last, line);
    }
    pclose(fp);
    double kbt = 0, mib = 0;
    long long xfrs = 0;
    if (sscanf(last, " %lf %lld %lf", &kbt, &xfrs, &mib) == 3) {
        return mib;
    }
    return 0.0;
}

typedef struct {
    off_t offset;
    size_t size;
} read_op_t;

typedef struct {
    int fd;
    const read_op_t *ops;
    size_t start_idx;
    size_t end_idx;
    size_t max_buf;
} worker_arg_t;

static void *pread_worker(void *arg) {
    worker_arg_t *w = (worker_arg_t *)arg;
    void *buf = NULL;
    posix_memalign(&buf, 16384, w->max_buf);
    volatile uint8_t sink = 0;
    for (size_t i = w->start_idx; i < w->end_idx; i++) {
        ssize_t n = pread(w->fd, buf, w->ops[i].size, w->ops[i].offset);
        if (n > 0) {
            sink ^= ((uint8_t *)buf)[0];
        }
    }
    free(buf);
    return (void *)(uintptr_t)sink;
}

static uint64_t rng_state = 0x853c49e6748fea9bULL;
static uint32_t xorshift32(void) {
    uint64_t x = rng_state;
    x ^= x >> 12;
    x ^= x << 25;
    x ^= x >> 27;
    rng_state = x;
    return (uint32_t)((x * 0x2545F4914F6CDD1DULL) >> 32);
}

static void write_fresh_file(const char *path, size_t file_size, void *wbuf, size_t chunk) {
    unlink(path);
    int wfd = open(path, O_CREAT | O_TRUNC | O_RDWR, 0600);
    if (wfd < 0) { perror("open write"); exit(1); }
    if (fcntl(wfd, F_NOCACHE, 1) < 0) { perror("F_NOCACHE write"); exit(1); }
    for (size_t off = 0; off < file_size; off += chunk) {
        if (write(wfd, wbuf, chunk) != (ssize_t)chunk) { perror("write"); exit(1); }
    }
    fsync(wfd);
    close(wfd);
}

static void run_verified_test(const char *label, const char *path,
                              read_op_t *ops, size_t n_ops, size_t max_buf,
                              int n_threads) {
    size_t total_bytes = 0;
    for (size_t i = 0; i < n_ops; i++) total_bytes += ops[i].size;

    int fd = open(path, O_RDONLY);
    if (fd < 0) { perror("open"); exit(1); }
    if (fcntl(fd, F_NOCACHE, 1) < 0) { perror("F_NOCACHE"); exit(1); }

    double mib0 = get_disk0_mib();
    pthread_t th[16];
    worker_arg_t args[16];
    double t0 = now_sec();
    for (int t = 0; t < n_threads; t++) {
        args[t].fd = fd;
        args[t].ops = ops;
        args[t].start_idx = (n_ops * t) / n_threads;
        args[t].end_idx = (n_ops * (t + 1)) / n_threads;
        args[t].max_buf = max_buf;
        pthread_create(&th[t], NULL, pread_worker, &args[t]);
    }
    for (int t = 0; t < n_threads; t++) {
        pthread_join(th[t], NULL);
    }
    double dt = now_sec() - t0;
    close(fd);
    double mib1 = get_disk0_mib();

    double req_mib = (double)total_bytes / (1024.0 * 1024.0);
    double disk_mib = mib1 - mib0;
    double nvme_pct = (disk_mib / req_mib) * 100.0;
    double gbps = ((double)total_bytes / 1e9) / dt;
    printf("%-46s | th=%d | %6.1f MiB | %6.2f GB/s | disk0=%6.1f MiB (%5.1f%%)\n",
           label, n_threads, req_mib, gbps, disk_mib, nvme_pct);
    fflush(stdout);
}

int main(void) {
    const char *path = "/tmp/flash_next_ssd_bench3.bin";
    const size_t slice_size = 1024ULL * 1024ULL * 1024ULL; // 1 GiB per slice
    const size_t num_slices = 12; // 12 GiB total so every single test reads a virgin 1 GiB region!
    const size_t file_size = slice_size * num_slices;

    size_t chunk = 4 * 1024 * 1024;
    void *wbuf = NULL;
    posix_memalign(&wbuf, 16384, chunk);
    for (size_t i = 0; i < chunk / 8; i++) ((uint64_t *)wbuf)[i] = i * 0x9e3779b97f4a7c15ULL;

    printf("=== Phase 0 Hardware Microbench v3 (Hardware iostat-verified Virgin Slices) ===\n");
    double m0 = get_disk0_mib();
    double tw0 = now_sec();
    write_fresh_file(path, file_size, wbuf, chunk);
    double tw = now_sec() - tw0;
    double m1 = get_disk0_mib();
    printf("Wrote %.1f GiB in %.2f s (%.2f GB/s, disk0 delta=%.1f MiB)\n\n",
           (double)file_size / (1024.0*1024.0*1024.0), tw, ((double)file_size / 1e9) / tw, m1 - m0);
    fflush(stdout);

    int threads_list[] = {1, 2, 4, 8};
    size_t cur_slice = 0;

    // 1. Sequential 4 MiB reads (th=1 and th=4) on virgin 1 GiB slices
    for (int ti = 0; ti < 2; ti++) {
        int th = (ti == 0) ? 1 : 4;
        off_t base = (off_t)(cur_slice++ * slice_size);
        size_t n_ops = 256;
        read_op_t *ops = malloc(n_ops * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops; i++) {
            ops[i].offset = base + (off_t)(i * 4 * 1024 * 1024);
            ops[i].size = 4 * 1024 * 1024;
        }
        run_verified_test("1. Seq 4 MiB reads (SSD ceiling)", path, ops, n_ops, 4*1024*1024, th);
        free(ops);
    }

    // 2. Standard GGUF 512 KiB disjoint random reads (th=1, 2, 4, 8) on virgin 1 GiB slices
    for (int ti = 0; ti < 4; ti++) {
        int th = threads_list[ti];
        off_t base = (off_t)(cur_slice++ * slice_size);
        size_t blk_size = 512 * 1024;
        size_t n_ops = slice_size / blk_size; // 2048 blocks = 1024 MiB
        size_t *perm = malloc(n_ops * sizeof(size_t));
        for (size_t i = 0; i < n_ops; i++) perm[i] = i;
        for (size_t i = n_ops - 1; i > 0; i--) {
            size_t j = xorshift32() % (i + 1);
            size_t tmp = perm[i]; perm[i] = perm[j]; perm[j] = tmp;
        }
        read_op_t *ops = malloc(n_ops * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops; i++) {
            ops[i].offset = base + (off_t)(perm[i] * blk_size);
            ops[i].size = blk_size;
        }
        run_verified_test("2. 512 KiB random (std GGUF 3 reads/exp)", path, ops, n_ops, blk_size, th);
        free(ops);
        free(perm);
    }

    // 3. Contiguous .experts copy 1.5 MiB disjoint random reads (th=1, 2, 4, 8) on virgin 1 GiB slices
    for (int ti = 0; ti < 4; ti++) {
        int th = threads_list[ti];
        off_t base = (off_t)(cur_slice++ * slice_size);
        size_t blk_size = 1536 * 1024;
        size_t n_ops = slice_size / blk_size; // 682 blocks = 1023 MiB
        size_t *perm = malloc(n_ops * sizeof(size_t));
        for (size_t i = 0; i < n_ops; i++) perm[i] = i;
        for (size_t i = n_ops - 1; i > 0; i--) {
            size_t j = xorshift32() % (i + 1);
            size_t tmp = perm[i]; perm[i] = perm[j]; perm[j] = tmp;
        }
        read_op_t *ops = malloc(n_ops * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops; i++) {
            ops[i].offset = base + (off_t)(perm[i] * blk_size);
            ops[i].size = blk_size;
        }
        run_verified_test("3. 1.5 MiB random (experts' copy 1 read/exp)", path, ops, n_ops, blk_size, th);
        free(ops);
        free(perm);
    }

    // 4. Prefill strided ascending order (120 of 224 experts/layer, 1.5 MiB/exp) on virgin 1 GiB slices (th=1, 4)
    for (int ti = 0; ti < 2; ti++) {
        int th = (ti == 0) ? 1 : 4;
        off_t base = (off_t)(cur_slice++ * slice_size);
        size_t blk_size = 1536 * 1024;
        size_t total_blks = slice_size / blk_size; // 682
        size_t n_ops = 360; // ~540 MiB strided ascending across the 1 GiB slice
        read_op_t *ops = malloc(n_ops * sizeof(read_op_t));
        size_t idx = 0;
        for (size_t i = 0; i < n_ops; i++) {
            size_t blk = (i * 189) / 100; // ascending stride ~1.89
            if (blk >= total_blks) blk = total_blks - 1;
            ops[i].offset = base + (off_t)(blk * blk_size);
            ops[i].size = blk_size;
        }
        run_verified_test("4. Prefill strided ascending 1.5 MiB/exp", path, ops, n_ops, blk_size, th);
        free(ops);
    }

    unlink(path);
    free(wbuf);
    return 0;
}