#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 int cmp_off(const void *a, const void *b) {
    off_t oa = ((const read_op_t *)a)->offset;
    off_t ob = ((const read_op_t *)b)->offset;
    return (oa > ob) - (oa < ob);
}

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("%-46s | 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_bench(void) {
    @autoreleasepool {
        id<MTLDevice> dev = MTLCreateSystemDefaultDevice();
        if (!dev) {
            printf("Metal device not available\n");
            return;
        }
        printf("\n=== Metal GPU Benchmark ===\n");
        printf("Device: %s | unifiedMemory=%d | recMaxWorkingSet=%.2f GiB\n",
               [[dev name] UTF8String], [dev hasUnifiedMemory],
               (double)[dev recommendedMaxWorkingSetSize] / (1024.0 * 1024.0 * 1024.0));

        id<MTLCommandQueue> q = [dev newCommandQueue];

        // Test 1: FP16 GEMM at Flash-Next shapes (K=2560, N=2560 dense and K=2560, N=1920 fused gate+up/expert)
        int Ms[] = {512, 1024, 4096};
        int Ns[] = {2560, 4096};
        int K = 2560;
        for (int mi = 0; mi < 3; mi++) {
            for (int ni = 0; ni < 2; ni++) {
                int M = Ms[mi];
                int N = Ns[ni];
                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];

                // Warmup
                for (int w = 0; w < 3; w++) {
                    id<MTLCommandBuffer> cb = [q commandBuffer];
                    [gemm encodeToCommandBuffer:cb leftMatrix:matA rightMatrix:matB resultMatrix:matC];
                    [cb commit];
                    [cb waitUntilCompleted];
                }

                int iters = 30;
                int batch_cmd = 8; // 8 GEMMs per command buffer to measure sustained GPU kernel throughput
                double best_tflops = 0.0;
                double tflops_arr[5];
                for (int r = 0; r < 5; r++) {
                    double t0 = now_sec();
                    for (int it = 0; it < iters; it++) {
                        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];
                    }
                    double dt = now_sec() - t0;
                    double flops = 2.0 * (double)M * (double)N * (double)K * (double)(iters * batch_cmd);
                    tflops_arr[r] = (flops / dt) / 1e12;
                }
                qsort(tflops_arr, 5, sizeof(double), cmp_double);
                printf("FP16 GEMM M=%4d K=%4d N=%4d             | med %6.2f TFLOP/s (min %.2f, max %.2f)\n",
                       M, K, N, tflops_arr[2], tflops_arr[0], tflops_arr[4]);
            }
        }
    }
}

int main(int argc, char **argv) {
    const char *path = "/tmp/flash_next_ssd_bench.bin";
    const size_t file_size = 4ULL * 1024ULL * 1024ULL * 1024ULL; // 4 GiB
    printf("=== Phase 0 Hardware Microbench: 13\" MacBook Air M5 ===\n");
    printf("Creating 4.0 GiB uncached benchmark file at %s ...\n", path);
    fflush(stdout);

    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;

    double tw0 = now_sec();
    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);
    double tw = now_sec() - tw0;
    printf("Wrote 4.0 GiB uncached in %.2f s (%.2f GB/s)\n\n", tw, ((double)file_size / 1e9) / tw);
    fflush(stdout);

    const int reps = 5;

    // 1. Sequential 4 MiB reads (1 GiB per rep, rotating across the 4 GiB file)
    {
        size_t n_ops = 256; // 1 GiB
        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)(i * 4 * 1024 * 1024);
            ops[i].size = 4 * 1024 * 1024;
        }
        run_op_suite("1. Seq 4 MiB reads (F_NOCACHE=1, SSD ceiling)", path, 1, ops, n_ops, 4*1024*1024, 1, reps);
        run_op_suite("1b. Seq 4 MiB reads (F_NOCACHE=1, 4 threads)", path, 1, ops, n_ops, 4*1024*1024, 4, reps);
        free(ops);
    }

    // 2. Prefill chunk pattern: 120 missed experts/layer * 16 layers = 1920 experts
    // Layout A: 3 reads/expert of 512 KiB (gate, up, down separate tensors, sorted within tensor)
    // Layout B: 1 read/expert of 1536 KiB (Strata contiguous .experts copy, sorted within layer)
    {
        size_t layers = 8; // 8 layers * 120 experts * 1.5 MiB = 1.44 GiB per rep
        size_t exps_per_layer = 120;
        size_t proj_bytes = 512 * 1024; // 512 KiB (Q2_0 projection size)
        size_t exp_bytes = 3 * proj_bytes; // 1.5 MiB

        // Layout A: 3 projections per expert (2880 ops)
        size_t n_ops_a = layers * exps_per_layer * 3;
        read_op_t *ops_a = malloc(n_ops_a * sizeof(read_op_t));
        size_t idx = 0;
        for (size_t l = 0; l < layers; l++) {
            off_t layer_base = (off_t)(l * 512ULL * 1024ULL * 1024ULL); // 512 MiB per layer region
            for (size_t p = 0; p < 3; p++) {
                off_t tensor_base = layer_base + (off_t)(p * 170ULL * 1024ULL * 1024ULL);
                for (size_t e = 0; e < exps_per_layer; e++) {
                    // Ascending expert IDs (sorted within chunk) with stride ~2-4 experts
                    size_t exp_id = (e * 2 + (xorshift32() & 1)) % 320;
                    ops_a[idx].offset = tensor_base + (off_t)(exp_id * proj_bytes);
                    ops_a[idx].size = proj_bytes;
                    idx++;
                }
            }
        }
        run_op_suite("2a. Prefill 3x512 KiB/exp (std GGUF)", path, 1, ops_a, n_ops_a, proj_bytes, 1, reps);
        run_op_suite("2b. Prefill 3x512 KiB/exp (std GGUF)", path, 1, ops_a, n_ops_a, proj_bytes, 4, reps);
        run_op_suite("2c. Prefill 3x512 KiB/exp (std GGUF)", path, 1, ops_a, n_ops_a, proj_bytes, 8, reps);

        // Layout B: 1 contiguous read of 1.5 MiB per expert (960 ops)
        size_t n_ops_b = layers * exps_per_layer;
        read_op_t *ops_b = malloc(n_ops_b * sizeof(read_op_t));
        idx = 0;
        for (size_t l = 0; l < layers; l++) {
            off_t layer_base = (off_t)(l * 512ULL * 1024ULL * 1024ULL);
            for (size_t e = 0; e < exps_per_layer; e++) {
                size_t exp_id = (e * 2 + (xorshift32() & 1)) % 320;
                ops_b[idx].offset = layer_base + (off_t)(exp_id * exp_bytes);
                ops_b[idx].size = exp_bytes;
                idx++;
            }
        }
        run_op_suite("3a. Prefill 1x1.5 MiB/exp (experts' copy)", path, 1, ops_b, n_ops_b, exp_bytes, 1, reps);
        run_op_suite("3b. Prefill 1x1.5 MiB/exp (experts' copy)", path, 1, ops_b, n_ops_b, exp_bytes, 4, reps);
        run_op_suite("3c. Prefill 1x1.5 MiB/exp (experts' copy)", path, 1, ops_b, n_ops_b, exp_bytes, 8, reps);

        free(ops_a);
        free(ops_b);
    }

    // 3. Decode miss pattern: 6 missed experts/layer * 48 layers = 288 experts (~453 MB)
    // Random expert IDs within each layer
    {
        size_t layers = 48;
        size_t misses_per_layer = 6;
        size_t proj_bytes = 512 * 1024;
        size_t exp_bytes = 3 * proj_bytes;

        // Decode Layout A: 3 x 512 KiB per missed expert (864 random ops)
        size_t n_ops_a = layers * misses_per_layer * 3;
        read_op_t *ops_a = malloc(n_ops_a * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops_a; i++) {
            size_t blk = xorshift32() % (file_size / proj_bytes);
            ops_a[i].offset = (off_t)(blk * proj_bytes);
            ops_a[i].size = proj_bytes;
        }
        run_op_suite("4a. Decode 3x512 KiB random (std GGUF)", path, 1, ops_a, n_ops_a, proj_bytes, 1, reps);
        run_op_suite("4b. Decode 3x512 KiB random (std GGUF)", path, 1, ops_a, n_ops_a, proj_bytes, 4, reps);
        run_op_suite("4c. Decode 3x512 KiB random (std GGUF)", path, 1, ops_a, n_ops_a, proj_bytes, 8, reps);

        // Decode Layout B: 1 x 1.5 MiB per missed expert (288 random ops)
        size_t n_ops_b = layers * misses_per_layer;
        read_op_t *ops_b = malloc(n_ops_b * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops_b; i++) {
            size_t blk = xorshift32() % (file_size / exp_bytes);
            ops_b[i].offset = (off_t)(blk * exp_bytes);
            ops_b[i].size = exp_bytes;
        }
        run_op_suite("5a. Decode 1x1.5 MiB random (experts' copy)", path, 1, ops_b, n_ops_b, exp_bytes, 1, reps);
        run_op_suite("5b. Decode 1x1.5 MiB random (experts' copy)", path, 1, ops_b, n_ops_b, exp_bytes, 4, reps);
        run_op_suite("5c. Decode 1x1.5 MiB random (experts' copy)", path, 1, ops_b, n_ops_b, exp_bytes, 8, reps);

        free(ops_a);
        free(ops_b);
    }

    // 4. PLE n-gram table lookup latency: 16 heads * 320 B row reads per token (512 tokens = 8192 random 4 KiB page reads)
    {
        size_t n_ops = 4096;
        size_t row_page = 4096;
        read_op_t *ops = malloc(n_ops * sizeof(read_op_t));
        for (size_t i = 0; i < n_ops; i++) {
            size_t blk = xorshift32() % (file_size / row_page);
            ops[i].offset = (off_t)(blk * row_page);
            ops[i].size = row_page;
        }
        run_op_suite("6a. PLE 4 KiB random reads (1 thread)", path, 1, ops, n_ops, row_page, 1, reps);
        run_op_suite("6b. PLE 4 KiB random reads (8 threads)", path, 1, ops, n_ops, row_page, 8, reps);
        free(ops);
    }

    unlink(path);
    free(wbuf);

    // 5. Metal GPU FP16 GEMM benchmark
    run_metal_bench();

    return 0;
}