/*
 * Cycle 2 M4'' — concurrent CPU(NEON LPF) + QPU(V3D LPF) throughput.
 *
 * Same pthread/barrier/timer pattern as bench_concurrent.c, but the
 * NEON worker calls ff_vp9_loop_filter_h_4_8_neon (per edge) and the
 * QPU worker dispatches v3d_lpf_h_4_8.spv.
 *
 * License: BSD-2-Clause; links FFmpeg NEON snapshot (LGPL-2.1+).
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stddef.h>
#include <time.h>
#include <getopt.h>
#include <pthread.h>
#include <sched.h>
#include <assert.h>
#include <vulkan/vulkan.h>

#include "v3d_runner.h"

extern void ff_vp9_loop_filter_h_4_8_neon(
    uint8_t *dst, ptrdiff_t stride, int E, int I, int H);

/* --- RNG / edge gen (mirrors bench_neon_lpf.c) ------------------- */

#define EDGE_STRIDE 8
#define EDGE_BYTES  64

static inline uint64_t xs_step(uint64_t *s) {
    uint64_t x = *s; x ^= x << 13; x ^= x >> 7; x ^= x << 17; return *s = x;
}
static uint64_t xs_init(uint64_t s) { return s ? s : 0xa57edbeef5717ULL; }

static void gen_edge_pixels(uint8_t *buf, uint64_t *s) {
    int a = (int)(xs_step(s) % 200) + 20;
    int b = (int)(xs_step(s) % 200) + 20;
    int n = (int)(xs_step(s) % 30);
    for (int r = 0; r < 8; r++)
        for (int c = 0; c < 8; c++) {
            int base = (c < 4) ? a : b;
            int noise = ((int)(xs_step(s) % (2*n + 1))) - n;
            int v = base + noise;
            buf[r*EDGE_STRIDE + c] = (uint8_t)(v < 0 ? 0 : v > 255 ? 255 : v);
        }
}
static void gen_thresholds(int *E, int *I, int *H, uint64_t *s) {
    *E = (int)(xs_step(s) % 81);
    *I = (int)(xs_step(s) % 41);
    *H = (int)(xs_step(s) % 11);
}
static double now_s(void) {
    struct timespec t; clock_gettime(CLOCK_MONOTONIC_RAW, &t);
    return t.tv_sec + t.tv_nsec * 1e-9;
}

static volatile int g_stop = 0;
static pthread_barrier_t g_start;

/* --- NEON worker ------------------------------------------------- */

#define NEON_BATCH 8192   /* edges held in memory per worker */

typedef struct {
    int worker_id, affinity_core;
    uint64_t edges_done;
    double elapsed_s;
} neon_args;

static void *neon_worker(void *p)
{
    neon_args *a = p;
    cpu_set_t cs; CPU_ZERO(&cs); CPU_SET(a->affinity_core, &cs);
    pthread_setaffinity_np(pthread_self(), sizeof(cs), &cs);

    uint64_t s = xs_init((uint64_t) a->worker_id * 0xc01dbeefULL);
    uint8_t *master = malloc((size_t) NEON_BATCH * EDGE_BYTES);
    uint8_t *work   = malloc((size_t) NEON_BATCH * EDGE_BYTES);
    int *Es = malloc(NEON_BATCH * sizeof(int));
    int *Is = malloc(NEON_BATCH * sizeof(int));
    int *Hs = malloc(NEON_BATCH * sizeof(int));
    for (int i = 0; i < NEON_BATCH; i++) {
        gen_edge_pixels(master + (size_t)i * EDGE_BYTES, &s);
        gen_thresholds(&Es[i], &Is[i], &Hs[i], &s);
    }

    pthread_barrier_wait(&g_start);
    double t0 = now_s();
    uint64_t done = 0;
    while (!g_stop) {
        memcpy(work, master, (size_t) NEON_BATCH * EDGE_BYTES);
        for (int i = 0; i < NEON_BATCH; i++)
            ff_vp9_loop_filter_h_4_8_neon(work + (size_t)i * EDGE_BYTES + 4,
                                          EDGE_STRIDE, Es[i], Is[i], Hs[i]);
        done += NEON_BATCH;
    }
    a->elapsed_s = now_s() - t0;
    a->edges_done = done;
    free(master); free(work); free(Es); free(Is); free(Hs);
    return NULL;
}

/* --- QPU worker ------------------------------------------------- */

typedef struct {
    int affinity_core;
    int n_edges;
    uint64_t edges_done;
    double elapsed_s;
} qpu_args;

typedef struct {
    uint32_t n_edges, dst_stride_u8, _pad0, _pad1;
} push_consts;

static void *qpu_worker(void *p)
{
    qpu_args *a = p;
    cpu_set_t cs; CPU_ZERO(&cs); CPU_SET(a->affinity_core, &cs);
    pthread_setaffinity_np(pthread_self(), sizeof(cs), &cs);

    v3d_runner *r = v3d_runner_create();
    if (!r) return NULL;

    int n_edges = a->n_edges;
    size_t dst_bytes  = (size_t) n_edges * EDGE_BYTES;
    size_t meta_bytes = (size_t) n_edges * 4 * sizeof(uint32_t);

    v3d_buffer buf_meta = {0}, buf_dst = {0};
    v3d_runner_create_buffer(r, meta_bytes, &buf_meta);
    v3d_runner_create_buffer(r, dst_bytes,  &buf_dst);

    uint64_t s = 0xfeedfacecafebabeULL;
    uint8_t *master = malloc(dst_bytes);
    for (int i = 0; i < n_edges; i++) gen_edge_pixels(master + (size_t)i * EDGE_BYTES, &s);

    uint32_t *meta = buf_meta.mapped;
    assert(EDGE_STRIDE >= 4);
    for (int i = 0; i < n_edges; i++) {
        uint32_t mx = (uint32_t)((size_t)i * EDGE_BYTES + 4);
        assert(mx >= 4);
        int E, I, H; gen_thresholds(&E, &I, &H, &s);
        meta[4*i + 0] = mx;
        meta[4*i + 1] = (uint32_t) E;
        meta[4*i + 2] = (uint32_t) I;
        meta[4*i + 3] = (uint32_t) H;
    }
    memcpy(buf_dst.mapped, master, dst_bytes);

    v3d_pipeline pipe = {0};
    v3d_runner_create_pipeline(r, "v3d_lpf_h_4_8.spv", 2, sizeof(push_consts), &pipe);
    v3d_buffer bufs[2] = { buf_meta, buf_dst };
    v3d_runner_bind_buffers(r, &pipe, bufs, 2);

    const uint32_t edges_per_wg = 32;
    uint32_t gc = (uint32_t)((n_edges + edges_per_wg - 1) / edges_per_wg);
    push_consts pc = { .n_edges = (uint32_t) n_edges,
                       .dst_stride_u8 = EDGE_STRIDE };

    VkCommandBuffer cb = v3d_runner_alloc_cmdbuf(r);
    VkCommandBufferBeginInfo cbbi = { .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
    vkBeginCommandBuffer(cb, &cbbi);
    vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.pipeline);
    vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE,
                            pipe.layout, 0, 1, &pipe.desc_set, 0, NULL);
    vkCmdPushConstants(cb, pipe.layout, VK_SHADER_STAGE_COMPUTE_BIT,
                       0, sizeof(pc), &pc);
    vkCmdDispatch(cb, gc, 1, 1);
    vkEndCommandBuffer(cb);

    for (int i = 0; i < 5; i++) v3d_runner_submit_wait(r, cb);   /* warm */

    pthread_barrier_wait(&g_start);
    double t0 = now_s();
    uint64_t done = 0;
    while (!g_stop) {
        memcpy(buf_dst.mapped, master, dst_bytes);
        v3d_runner_submit_wait(r, cb);
        done += n_edges;
    }
    a->elapsed_s = now_s() - t0;
    a->edges_done = done;

    free(master);
    v3d_runner_destroy_pipeline(r, &pipe);
    v3d_runner_destroy_buffer(r, &buf_dst);
    v3d_runner_destroy_buffer(r, &buf_meta);
    v3d_runner_destroy(r);
    return NULL;
}

/* --- Timer ------------------------------------------------------ */

typedef struct { double duration_s; } timer_args;
static void *timer_thread(void *p) {
    timer_args *a = p;
    pthread_barrier_wait(&g_start);
    double end = now_s() + a->duration_s;
    while (now_s() < end) {
        struct timespec ts = {0, 1000000}; nanosleep(&ts, NULL);
    }
    g_stop = 1;
    return NULL;
}

/* --- Main ------------------------------------------------------- */

enum mode { MODE_NEON, MODE_QPU, MODE_MIXED };

int main(int argc, char **argv)
{
    enum mode mode = MODE_NEON;
    int n_neon = 4;
    int qpu_core = 3;
    int qpu_n_edges = 65536;
    double duration = 8.0;

    static struct option opts[] = {
        {"mode",          required_argument, 0, 'm'},
        {"neon-threads",  required_argument, 0, 'n'},
        {"qpu-core",      required_argument, 0, 'c'},
        {"qpu-edges",     required_argument, 0, 'e'},
        {"duration",      required_argument, 0, 'd'},
        {0,0,0,0}
    };
    for (int c; (c = getopt_long(argc, argv, "m:n:c:e:d:", opts, 0)) != -1;) {
        switch (c) {
        case 'm':
            if      (!strcmp(optarg, "neon-only")) mode = MODE_NEON;
            else if (!strcmp(optarg, "qpu-only"))  mode = MODE_QPU;
            else if (!strcmp(optarg, "mixed"))     mode = MODE_MIXED;
            else { fprintf(stderr, "bad mode\n"); return 2; }
            break;
        case 'n': n_neon = atoi(optarg); break;
        case 'c': qpu_core = atoi(optarg); break;
        case 'e': qpu_n_edges = atoi(optarg); break;
        case 'd': duration = atof(optarg); break;
        default: return 2;
        }
    }
    int has_qpu  = (mode == MODE_QPU || mode == MODE_MIXED);
    int has_neon = (mode == MODE_NEON || mode == MODE_MIXED);
    int n_workers = (has_neon ? n_neon : 0) + (has_qpu ? 1 : 0);
    int barrier_count = n_workers + 1 /* timer */ + 1 /* main */;

    printf("=== M4'' concurrent LPF bench ===\n");
    printf("  mode:         %s\n", mode == MODE_NEON ? "neon-only" : mode == MODE_QPU ? "qpu-only" : "mixed");
    printf("  neon threads: %d (cores 0..%d)\n", has_neon ? n_neon : 0, has_neon ? n_neon - 1 : -1);
    printf("  qpu host:     core %d, %d edges/dispatch\n",
           has_qpu ? qpu_core : -1, has_qpu ? qpu_n_edges : 0);
    printf("  duration:     %.1f s\n\n", duration);

    pthread_barrier_init(&g_start, NULL, barrier_count);

    pthread_t timer_tid; timer_args ta = { .duration_s = duration };
    pthread_create(&timer_tid, NULL, timer_thread, &ta);

    pthread_t neon_tids[16] = {0};
    neon_args n_args[16] = {0};
    if (has_neon) {
        for (int i = 0; i < n_neon; i++) {
            n_args[i] = (neon_args){ .worker_id = i, .affinity_core = i };
            pthread_create(&neon_tids[i], NULL, neon_worker, &n_args[i]);
        }
    }
    pthread_t qpu_tid = 0;
    qpu_args q_args = {0};
    if (has_qpu) {
        q_args = (qpu_args){ .affinity_core = qpu_core, .n_edges = qpu_n_edges };
        pthread_create(&qpu_tid, NULL, qpu_worker, &q_args);
    }

    pthread_barrier_wait(&g_start);

    pthread_join(timer_tid, NULL);
    if (has_neon) for (int i = 0; i < n_neon; i++) pthread_join(neon_tids[i], NULL);
    if (has_qpu)  pthread_join(qpu_tid, NULL);

    uint64_t total_edges = 0; double max_elapsed = 0;

    if (has_neon) {
        printf("NEON per-thread:\n");
        for (int i = 0; i < n_neon; i++) {
            double mes = n_args[i].edges_done / n_args[i].elapsed_s / 1e6;
            printf("  core %d: %.3f Medge/s  (%llu edges / %.3f s)\n",
                   n_args[i].affinity_core, mes,
                   (unsigned long long) n_args[i].edges_done, n_args[i].elapsed_s);
            total_edges += n_args[i].edges_done;
            if (n_args[i].elapsed_s > max_elapsed) max_elapsed = n_args[i].elapsed_s;
        }
    }
    if (has_qpu) {
        double mes = q_args.edges_done / q_args.elapsed_s / 1e6;
        printf("QPU (host core %d): %.3f Medge/s  (%llu edges / %.3f s)\n",
               q_args.affinity_core, mes,
               (unsigned long long) q_args.edges_done, q_args.elapsed_s);
        total_edges += q_args.edges_done;
        if (q_args.elapsed_s > max_elapsed) max_elapsed = q_args.elapsed_s;
    }

    double total_mes = total_edges / max_elapsed / 1e6;
    printf("\n=== AGGREGATE ===\n");
    printf("  total edges    : %llu\n", (unsigned long long) total_edges);
    printf("  wall-clock     : %.3f s\n", max_elapsed);
    printf("  Medge/s        : %.3f\n", total_mes);

    pthread_barrier_destroy(&g_start);
    return 0;
}