/*
 * Cycle 3 M4''' — concurrent CPU(NEON MC) + QPU(V3D MC) throughput.
 * Same pthread/barrier pattern as bench_concurrent{,_lpf}.c.
 * License: BSD-2-Clause.
 */
#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_put_regular8_h_neon(
    uint8_t *dst, ptrdiff_t dst_stride,
    const uint8_t *src, ptrdiff_t src_stride,
    int h, int mx, int my);

#define SRC_W 16
#define DST_W 8
#define SRC_H 8
#define DST_H 8
#define SRC_BYTES (SRC_H * SRC_W)
#define DST_BYTES (DST_H * DST_W)

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 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

typedef struct {
    int worker_id, affinity_core;
    uint64_t blocks_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 * SRC_BYTES);
    uint8_t *work   = malloc((size_t) NEON_BATCH * SRC_BYTES);
    uint8_t *dsts   = malloc((size_t) NEON_BATCH * DST_BYTES);
    int     *mxs    = malloc(NEON_BATCH * sizeof(int));
    for (int i = 0; i < NEON_BATCH; i++) {
        for (int j = 0; j < SRC_BYTES; j++)
            master[(size_t)i * SRC_BYTES + j] = (uint8_t)(xs_step(&s) & 0xff);
        mxs[i] = (int)(xs_step(&s) & 15);
    }

    pthread_barrier_wait(&g_start);
    double t0 = now_s();
    uint64_t done = 0;
    while (!g_stop) {
        memcpy(work, master, (size_t) NEON_BATCH * SRC_BYTES);
        for (int i = 0; i < NEON_BATCH; i++)
            ff_vp9_put_regular8_h_neon(
                dsts + (size_t)i * DST_BYTES, DST_W,
                work + (size_t)i * SRC_BYTES + 3, SRC_W,
                DST_H, mxs[i], 0);
        done += NEON_BATCH;
    }
    a->elapsed_s = now_s() - t0;
    a->blocks_done = done;
    free(master); free(work); free(dsts); free(mxs);
    return NULL;
}

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

typedef struct {
    int affinity_core, n_blocks;
    uint64_t blocks_done;
    double elapsed_s;
} qpu_args;

typedef struct {
    uint32_t n_blocks, dst_stride_u8, src_stride_u8, _pad;
} 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_blocks = a->n_blocks;
    size_t meta_bytes = (size_t) n_blocks * 4 * sizeof(uint32_t);
    size_t src_bytes  = (size_t) n_blocks * SRC_BYTES;
    size_t dst_bytes  = (size_t) n_blocks * DST_BYTES;

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

    uint64_t s = 0xfeedfacecafebabeULL;
    uint8_t *master = malloc(src_bytes);
    for (size_t i = 0; i < src_bytes; i++) master[i] = (uint8_t)(xs_step(&s) & 0xff);
    memcpy(buf_src.mapped, master, src_bytes);

    uint32_t *meta = buf_meta.mapped;
    assert(DST_W >= 8); assert(SRC_W >= 15);
    for (int i = 0; i < n_blocks; i++) {
        meta[4*i + 0] = (uint32_t)((size_t)i * DST_BYTES);   /* dst_off */
        meta[4*i + 1] = (uint32_t)((size_t)i * SRC_BYTES);   /* src_off (RAW, no +3) */
        meta[4*i + 2] = (uint32_t)(xs_step(&s) & 15);        /* mx */
        meta[4*i + 3] = 0;
    }

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

    const uint32_t bpw = 32;
    uint32_t gc = (uint32_t)((n_blocks + bpw - 1) / bpw);
    push_consts pc = { .n_blocks = (uint32_t) n_blocks,
                       .dst_stride_u8 = DST_W,
                       .src_stride_u8 = SRC_W };

    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);

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

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

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;
}

enum mode { MODE_NEON, MODE_QPU, MODE_MIXED };

int main(int argc, char **argv)
{
    enum mode mode = MODE_NEON;
    int n_neon = 4, qpu_core = 3, qpu_n_blocks = 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-blocks",   required_argument, 0, 'b'},
        {"duration",     required_argument, 0, 'd'},
        {0,0,0,0}
    };
    for (int c; (c = getopt_long(argc, argv, "m:n:c:b: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 'b': qpu_n_blocks = 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 + 1;

    printf("=== M4''' concurrent MC bench ===\n");
    printf("  mode: %s, neon: %d, qpu: core %d / %d blocks, %.1fs\n",
           mode == MODE_NEON ? "neon-only" : mode == MODE_QPU ? "qpu-only" : "mixed",
           has_neon ? n_neon : 0,
           has_qpu ? qpu_core : -1,
           has_qpu ? qpu_n_blocks : 0,
           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_blocks = qpu_n_blocks };
        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 = 0; double max_e = 0;
    if (has_neon) {
        printf("NEON per-thread:\n");
        for (int i = 0; i < n_neon; i++) {
            double mbs = n_args[i].blocks_done / n_args[i].elapsed_s / 1e6;
            printf("  core %d: %.3f Mblock/s\n", n_args[i].affinity_core, mbs);
            total += n_args[i].blocks_done;
            if (n_args[i].elapsed_s > max_e) max_e = n_args[i].elapsed_s;
        }
    }
    if (has_qpu) {
        double mbs = q_args.blocks_done / q_args.elapsed_s / 1e6;
        printf("QPU (core %d): %.3f Mblock/s\n", q_args.affinity_core, mbs);
        total += q_args.blocks_done;
        if (q_args.elapsed_s > max_e) max_e = q_args.elapsed_s;
    }

    double total_mbs = total / max_e / 1e6;
    printf("\n=== AGGREGATE ===\n");
    printf("  Mblock/s        : %.3f\n", total_mbs);
    printf("  30fps@1080p floor: 0.972 Mblock/s — %.1fx margin\n",
           total_mbs / 0.972);

    pthread_barrier_destroy(&g_start);
    return 0;
}