#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <linux/dma-heap.h>
#include <math.h>

#include "detector.h"

static int dma_heap_alloc(size_t size, int *fd) {
    char path[64];
    snprintf(path, 64, "/dev/dma_heap/reserved");
    int heap_fd = open(path, O_RDWR | O_CLOEXEC);
    if (heap_fd < 0) {
        heap_fd = open("/dev/dma_heap/system-uncached", O_RDWR | O_CLOEXEC);
        if (heap_fd < 0) return -1;
    }
    struct dma_heap_allocation_data d = { .len = size, .fd_flags = O_RDWR | O_CLOEXEC };
    int ret = ioctl(heap_fd, DMA_HEAP_IOCTL_ALLOC, &d);
    close(heap_fd);
    if (ret < 0) return -1;
    *fd = d.fd;
    return 0;
}

static double get_time_ms(void) {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
}

static int bench_model(const char *model_path, const char *model_name,
                       int nv12_fd, int src_w, int src_h,
                       int warmup, int iterations) {
    printf("\n=== Benchmark: %s (%s) ===\n", model_name, model_path);

    Detector *det = detector_create(model_path, 0.25f, 0.45f);
    if (!det) {
        printf("  FAILED to create detector\n");
        return -1;
    }

    int mw = detector_get_input_width(det);
    int mh = detector_get_input_height(det);
    printf("  Model input: %dx%d\n", mw, mh);

    DetResult results[MAX_DETECTIONS];
    int num_results = 0;

    for (int i = 0; i < warmup; i++) {
        detector_detect(det, nv12_fd, src_w, src_h, results, &num_results);
    }

    double t_start = get_time_ms();
    int total_detections = 0;
    int valid_runs = 0;

    for (int i = 0; i < iterations; i++) {
        num_results = 0;
        int ret = detector_detect(det, nv12_fd, src_w, src_h, results, &num_results);
        if (ret == 0) {
            total_detections += num_results;
            valid_runs++;
        }
    }

    double t_end = get_time_ms();
    double elapsed = t_end - t_start;
    double avg_ms = valid_runs > 0 ? elapsed / valid_runs : 0;
    double fps = avg_ms > 0 ? 1000.0 / avg_ms : 0;
    int avg_det = valid_runs > 0 ? total_detections / valid_runs : 0;

    printf("  Warmup: %d, Iterations: %d (valid: %d)\n", warmup, iterations, valid_runs);
    printf("  Total time: %.1f ms\n", elapsed);
    printf("  Avg inference: %.2f ms/frame\n", avg_ms);
    printf("  Avg FPS: %.1f\n", fps);
    printf("  Avg detections: %d\n", avg_det);

    if (valid_runs > 0 && num_results > 0) {
        printf("  Last frame detections:\n");
        for (int i = 0; i < num_results && i < 10; i++) {
            printf("    [%d] label=%d score=%.3f box=[%.1f,%.1f,%.1f,%.1f]\n",
                   i, results[i].label, results[i].score,
                   results[i].x1, results[i].y1, results[i].x2, results[i].y2);
        }
        if (num_results > 10)
            printf("    ... and %d more\n", num_results - 10);
    }

    detector_destroy(det);
    return 0;
}

int main(int argc, char **argv) {
    const char *model5 = "models/yolov5s.rknn";
    const char *model8 = "models/yolov8s.rknn";
    int src_w = 1920, src_h = 1080;
    int warmup = 5, iterations = 50;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "-w") == 0 && i + 1 < argc) { src_w = atoi(argv[++i]); }
        else if (strcmp(argv[i], "-h") == 0 && i + 1 < argc) { src_h = atoi(argv[++i]); }
        else if (strcmp(argv[i], "-n") == 0 && i + 1 < argc) { iterations = atoi(argv[++i]); }
        else if (strcmp(argv[i], "-m5") == 0 && i + 1 < argc) { model5 = argv[++i]; }
        else if (strcmp(argv[i], "-m8") == 0 && i + 1 < argc) { model8 = argv[++i]; }
    }

    printf("Detector Benchmark: YOLOv5s vs YOLOv8s on RK3588 NPU\n");
    printf("Source: %dx%d NV12 (synthetic), Warmup: %d, Iterations: %d\n",
           src_w, src_h, warmup, iterations);

    size_t nv12_size = src_w * src_h * 3 / 2;
    int nv12_fd = -1;
    if (dma_heap_alloc(nv12_size, &nv12_fd) < 0) {
        printf("ERROR: failed to allocate NV12 buffer\n");
        return 1;
    }

    void *nv12_map = mmap(NULL, nv12_size, PROT_READ | PROT_WRITE,
                           MAP_SHARED, nv12_fd, 0);
    if (nv12_map == MAP_FAILED) {
        printf("ERROR: mmap failed\n");
        close(nv12_fd);
        return 1;
    }

    memset(nv12_map, 0x80, nv12_size);
    uint8_t *y_plane = (uint8_t *)nv12_map;
    for (int row = 0; row < src_h; row++) {
        for (int col = 0; col < src_w; col++) {
            int val = ((row + col) % 200) + 30;
            y_plane[row * src_w + col] = (uint8_t)val;
        }
    }

    printf("NV12 buffer ready (%zu bytes)\n", nv12_size);

    int ret5 = bench_model(model5, "YOLOv5s", nv12_fd, src_w, src_h, warmup, iterations);
    int ret8 = bench_model(model8, "YOLOv8s", nv12_fd, src_w, src_h, warmup, iterations);

    printf("\n=== Summary ===\n");
    if (ret5 == 0 && ret8 == 0) {
        printf("Both models loaded and ran successfully.\n");
        printf("See above for per-model latency and FPS comparison.\n");
    } else {
        if (ret5 < 0) printf("YOLOv5s: FAILED\n");
        if (ret8 < 0) printf("YOLOv8s: FAILED\n");
    }

    munmap(nv12_map, nv12_size);
    close(nv12_fd);
    return 0;
}