#import <Foundation/Foundation.h>
#import <Metal/Metal.h>
#import <MetalPerformanceShaders/MetalPerformanceShaders.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;
}

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 int cmp_double(const void *a, const void *b) {
    double da = *(const double *)a;
    double db = *(const double *)b;
    return (da > db) - (da < db);
}

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 run_op_suite(const char *label, const char *path, int nocache,
                         read_op_t *ops, size_t n_ops, size_t max_buf,
                         int n_threads, int reps) {
    size_t total_bytes = 0;
    for (size_t i = 0; i < n_ops; i++) total_bytes += ops[i].size;

    double gbps[16];
    for (int r = 0; r < reps; r++) {
        int fd = open(path, O_RDONLY);
        if (fd < 0) { perror("open"); exit(1); }
        if (nocache) {
            if (fcntl(fd, F_NOCACHE, 1) < 0) { perror("F_NOCACHE"); exit(1); }
        }

        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);
        gbps[r] = ((double)total_bytes / 1e9) / dt;
    }
    qsort(gbps, reps, sizeof(double), cmp_double);
    double med = gbps[reps / 2];
    printf("%-48s | th=%d | %6.1f MB | med %6.2f GB/s (min %.2f, max %.2f)\n",
           label, n_threads, (double)total_bytes / 1e6, med, gbps[0], gbps[reps - 1]);
    fflush(stdout);
}

static void run_metal_thermal_test(int duration_sec) {
    @autoreleasepool {
        id<MTLDevice> dev = MTLCreateSystemDefaultDevice();
        if (!dev) return;
        id<MTLCommandQueue> q = [dev newCommandQueue];
        int M = 4096, K = 2560, N = 2560;
        size_t bytesA = (size_t)M * K * 2;
        size_t bytesB = (size_t)K * N * 2;
        size_t bytesC = (size_t)M * N * 2;
        id<MTLBuffer> bufA = [dev newBufferWithLength:bytesA options:MTLResourceStorageModeShared];
        id<MTLBuffer> bufB = [dev newBufferWithLength:bytesB options:MTLResourceStorageModeShared];
        id<MTLBuffer> bufC = [dev newBufferWithLength:bytesC options:MTLResourceStorageModeShared];
        memset([bufA contents], 0x3c, bytesA);
        memset([bufB contents], 0x3c, bytesB);

        MPSMatrixDescriptor *descA = [MPSMatrixDescriptor matrixDescriptorWithRows:M columns:K rowBytes:K*2 dataType:MPSDataTypeFloat16];
        MPSMatrixDescriptor *descB = [MPSMatrixDescriptor matrixDescriptorWithRows:K columns:N rowBytes:N*2 dataType:MPSDataTypeFloat16];
        MPSMatrixDescriptor *descC = [MPSMatrixDescriptor matrixDescriptorWithRows:M columns:N rowBytes:N*2 dataType:MPSDataTypeFloat16];
        MPSMatrix *matA = [[MPSMatrix alloc] initWithBuffer:bufA descriptor:descA];
        MPSMatrix *matB = [[MPSMatrix alloc] initWithBuffer:bufB descriptor:descB];
        MPSMatrix *matC = [[MPSMatrix alloc] initWithBuffer:bufC descriptor:descC];
        MPSMatrixMultiplication *gemm = [[MPSMatrixMultiplication alloc] initWithDevice:dev transposeLeft:NO transposeRight:NO resultRows:M resultColumns:N interiorColumns:K alpha:1.0 beta:0.0];

        printf("\n=== Metal GPU 60s Sustained Load / Thermal Sag Test (M=4096, K=2560, N=2560 FP16) ===\n");
        fflush(stdout);

        double start = now_sec();
        double window_start = start;
        long window_gemms = 0;
        int window_idx = 0;
        int batch_cmd = 12;

        while (1) {
            id<MTLCommandBuffer> cb = [q commandBuffer];
            for (int b = 0; b < batch_cmd; b++) {
                [gemm encodeToCommandBuffer:cb leftMatrix:matA rightMatrix:matB resultMatrix:matC];
            }
            [cb commit];
            [cb waitUntilCompleted];
            window_gemms += batch_cmd;

            double now = now_sec();
            if (now - window_start >= 5.0) {
                double dt = now - window_start;
                double flops = 2.0 * (double)M * (double)N * (double)K * (double)window_gemms;
                double tflops = (flops / dt) / 1e12;
                window_idx++;
                printf("  t=%2ds..%2ds | %4ld GEMMs | %6.2f TFLOP/s\n",
                       (window_idx - 1) * 5, window_idx * 5, window_gemms, tflops);
                fflush(stdout);
                window_start = now;
                window_gemms = 0;
                if (now - start >= (double)duration_sec) break;
            }
        }
    }
}

int main(int argc, char **argv) {
    const char *path = "/tmp/flash_next_ssd_bench2.bin";
    const size_t file_size = 4ULL * 1024ULL * 1024ULL * 1024ULL; // 4 GiB
    printf("=== Phase 0 Hardware Microbench v2 (Disjoint Offsets + QD Curve + 60s Thermal) ===\n");

    int wfd = open(path, O_CREAT | O_TRUNC | O_RDWR, 0600);
    if (wfd < 0) { perror("open write"); return 1; }
    if (fcntl(wfd, F_NOCACHE, 1) < 0) { perror("F_NOCACHE write"); return 1; }

    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;
    for (size_t off = 0; off < file_size; off += chunk) {
        if (write(wfd, wbuf, chunk) != (ssize_t)chunk) { perror("write"); return 1; }
    }
    fsync(wfd);
    close(wfd);

    const int reps = 3;
    int threads_list[] = {1, 2, 4, 8};

    // Disjoint random permutation of 512 KiB blocks across 4 GiB (8192 total blocks, sample first 2048 = 1 GiB)
    {
        size_t blk_size = 512 * 1024;
        size_t total_blks = file_size / blk_size; // 8192
        size_t *perm = malloc(total_blks * sizeof(size_t));
        for (size_t i = 0; i < total_blks; i++) perm[i] = i;
        for (size_t i = total_blks - 1; i > 0; i--) {
            size_t j = xorshift32() % (i + 1);
            size_t tmp = perm[i]; perm[i] = perm[j]; perm[j] = tmp;
        }
        size_t n_ops = 2048; // 1.07 GB, strictly disjoint blocks
        read_op_t *ops = malloc(n_ops * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops; i++) {
            ops[i].offset = (off_t)(perm[i] * blk_size);
            ops[i].size = blk_size;
        }
        for (int ti = 0; ti < 4; ti++) {
            run_op_suite("Disjoint 512 KiB random (std GGUF, F_NOCACHE=1)", path, 1, ops, n_ops, blk_size, threads_list[ti], reps);
        }
        // Control: populate page cache once, then measure F_NOCACHE=0 cached RAM pread
        run_op_suite("Control: 512 KiB cached warmup (F_NOCACHE=0)", path, 0, ops, n_ops, blk_size, 4, 1);
        run_op_suite("Control: 512 KiB cached RAM (F_NOCACHE=0)", path, 0, ops, n_ops, blk_size, 1, reps);
        run_op_suite("Control: 512 KiB cached RAM (F_NOCACHE=0)", path, 0, ops, n_ops, blk_size, 4, reps);
        free(ops);
        free(perm);
    }

    // Re-write file with F_NOCACHE=1 to evict any pages cached by the control test!
    wfd = open(path, O_CREAT | O_TRUNC | O_RDWR, 0600);
    fcntl(wfd, F_NOCACHE, 1);
    for (size_t off = 0; off < file_size; off += chunk) {
        if (write(wfd, wbuf, chunk) != (ssize_t)chunk) { perror("write"); return 1; }
    }
    fsync(wfd);
    close(wfd);

    // Disjoint random permutation of 1.5 MiB blocks across 4 GiB (2730 total blocks, sample first 700 = 1.1 GB)
    {
        size_t blk_size = 1536 * 1024;
        size_t total_blks = file_size / blk_size;
        size_t *perm = malloc(total_blks * sizeof(size_t));
        for (size_t i = 0; i < total_blks; i++) perm[i] = i;
        for (size_t i = total_blks - 1; i > 0; i--) {
            size_t j = xorshift32() % (i + 1);
            size_t tmp = perm[i]; perm[i] = perm[j]; perm[j] = tmp;
        }
        size_t n_ops = 700; // 1.10 GB, strictly disjoint blocks
        read_op_t *ops = malloc(n_ops * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops; i++) {
            ops[i].offset = (off_t)(perm[i] * blk_size);
            ops[i].size = blk_size;
        }
        for (int ti = 0; ti < 4; ti++) {
            run_op_suite("Disjoint 1.5 MiB random (experts' copy, NOCACHE)", path, 1, ops, n_ops, blk_size, threads_list[ti], reps);
        }
        free(ops);
        free(perm);
    }

    unlink(path);
    free(wbuf);

    run_metal_thermal_test(60);
    return 0;
}