#include "detector.h"
#include "dma_util.h"
#include "rknn_api.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <linux/dma-heap.h>
#include <linux/dma-buf.h>
#include <math.h>
#include <rga/im2d.h>
#include <rga/rga.h>

#define DBG(...) fprintf(stderr, __VA_ARGS__)

struct Detector {
    rknn_context ctx;
    int model_w;
    int model_h;
    int model_type;
    rknn_input_output_num io_num;
    rknn_tensor_attr *input_attrs;
    rknn_tensor_attr *output_attrs;
    float conf_thresh;
    float nms_thresh;

    int rgb_fd;
    void *rgb_map;
    size_t rgb_size;
    int rgb_initialized;
};

enum {
    MODEL_YOLOV5 = 0,
    MODEL_YOLOV8 = 1,
};

Detector *detector_create(const char *model_path, float conf_thresh, float nms_thresh) {
    Detector *det = calloc(1, sizeof(Detector));
    if (!det) return NULL;

    det->conf_thresh = conf_thresh;
    det->nms_thresh = nms_thresh;
    det->ctx = 0;

    FILE *fp = fopen(model_path, "rb");
    if (!fp) { DBG("det: cannot open model %s\n", model_path); goto fail; }
    fseek(fp, 0, SEEK_END);
    long model_size = ftell(fp);
    rewind(fp);
    void *model_data = malloc(model_size);
    if (!model_data) { fclose(fp); goto fail; }
    if (fread(model_data, 1, model_size, fp) != (size_t)model_size) {
        fclose(fp); free(model_data); goto fail;
    }
    fclose(fp);

    int ret = rknn_init(&det->ctx, model_data, model_size, 0, NULL);
    free(model_data);
    if (ret < 0) { DBG("det: rknn_init failed %d\n", ret); goto fail; }

    ret = rknn_query(det->ctx, RKNN_QUERY_IN_OUT_NUM, &det->io_num, sizeof(det->io_num));
    if (ret < 0) { DBG("det: query io num failed %d\n", ret); goto fail; }
    DBG("det: %d inputs, %d outputs\n", det->io_num.n_input, det->io_num.n_output);

    det->input_attrs = calloc(det->io_num.n_input, sizeof(rknn_tensor_attr));
    det->output_attrs = calloc(det->io_num.n_output, sizeof(rknn_tensor_attr));
    if (!det->input_attrs || !det->output_attrs) goto fail;

    for (uint32_t i = 0; i < det->io_num.n_input; i++) {
        det->input_attrs[i].index = i;
        rknn_query(det->ctx, RKNN_QUERY_INPUT_ATTR, &det->input_attrs[i], sizeof(rknn_tensor_attr));
        DBG("det: input[%d] %s n_dims=%d dims=[%d,%d,%d,%d] fmt=%d type=%d\n",
            i, det->input_attrs[i].name, det->input_attrs[i].n_dims,
            det->input_attrs[i].dims[0], det->input_attrs[i].dims[1],
            det->input_attrs[i].dims[2], det->input_attrs[i].dims[3],
            det->input_attrs[i].fmt, det->input_attrs[i].type);
    }
    for (uint32_t i = 0; i < det->io_num.n_output; i++) {
        det->output_attrs[i].index = i;
        rknn_query(det->ctx, RKNN_QUERY_OUTPUT_ATTR, &det->output_attrs[i], sizeof(rknn_tensor_attr));
        DBG("det: output[%d] %s n_dims=%d dims=[%d,%d,%d,%d] fmt=%d type=%d\n",
            i, det->output_attrs[i].name, det->output_attrs[i].n_dims,
            det->output_attrs[i].dims[0], det->output_attrs[i].dims[1],
            det->output_attrs[i].dims[2], det->output_attrs[i].dims[3],
            det->output_attrs[i].fmt, det->output_attrs[i].type);
    }

    det->model_w = det->input_attrs[0].dims[3];
    det->model_h = det->input_attrs[0].dims[2];

    det->model_type = MODEL_YOLOV8;
    if (det->io_num.n_output == 3) {
        int all_anchor = 1;
        for (uint32_t i = 0; i < det->io_num.n_output; i++) {
            if (det->output_attrs[i].n_dims != 3 ||
                !(det->output_attrs[i].dims[2] == 85 ||
                  det->output_attrs[i].dims[2] == 56)) {
                all_anchor = 0;
                break;
            }
        }
        if (all_anchor) det->model_type = MODEL_YOLOV5;
    }
    DBG("det: model input %dx%d type=%s\n", det->model_w, det->model_h,
        det->model_type == MODEL_YOLOV5 ? "yolov5" : "yolov8");

    det->rgb_size = det->model_w * det->model_h * 3;
    if (dma_heap_alloc("reserved", det->rgb_size, &det->rgb_fd) < 0) {
        DBG("det: failed to alloc rgb buffer\n");
        goto fail;
    }
    det->rgb_map = mmap(NULL, det->rgb_size, PROT_READ | PROT_WRITE, MAP_SHARED, det->rgb_fd, 0);
    if (det->rgb_map == MAP_FAILED) {
        DBG("det: mmap rgb buffer failed\n");
        close(det->rgb_fd); det->rgb_fd = -1;
        goto fail;
    }
    det->rgb_initialized = 1;

    DBG("det: created successfully, model %dx%d\n", det->model_w, det->model_h);
    return det;

fail:
    detector_destroy(det);
    return NULL;
}

void detector_destroy(Detector *det) {
    if (!det) return;
    if (det->rgb_initialized && det->rgb_map && det->rgb_map != MAP_FAILED)
        munmap(det->rgb_map, det->rgb_size);
    if (det->rgb_fd >= 0) close(det->rgb_fd);
    if (det->ctx) rknn_destroy(det->ctx);
    free(det->input_attrs);
    free(det->output_attrs);
    free(det);
}

int detector_get_input_width(Detector *det) { return det->model_w; }
int detector_get_input_height(Detector *det) { return det->model_h; }

static int rga_scale_nv12_to_rgb(int src_fd, int src_w, int src_h,
                                  int dst_fd, int dst_w, int dst_h,
                                  void *rgb_map, size_t rgb_size) {
    memset(rgb_map, 0, rgb_size);

    float scale = fminf((float)dst_w / src_w, (float)dst_h / src_h);
    int content_w = (int)(src_w * scale);
    int content_h = (int)(src_h * scale);
    int pad_x = (dst_w - content_w) / 2;
    int pad_y = (dst_h - content_h) / 2;

    rga_buffer_t src = wrapbuffer_fd_t(src_fd, src_w, src_h, src_w, src_h, RK_FORMAT_YCbCr_420_SP);
    rga_buffer_t dst = wrapbuffer_fd_t(dst_fd, dst_w, dst_h, dst_w, dst_h, RK_FORMAT_RGB_888);

    im_rect src_rect = {0, 0, src_w, src_h};
    im_rect dst_rect = {pad_x, pad_y, content_w, content_h};
    rga_buffer_t pat;
    memset(&pat, 0, sizeof(pat));
    im_rect prect;
    memset(&prect, 0, sizeof(prect));

    int ret = improcess(src, dst, pat, src_rect, dst_rect, prect, IM_SYNC);
    if (ret != IM_STATUS_SUCCESS) {
        DBG("det: RGA scale failed: %d\n", ret);
        return -1;
    }
    return 0;
}

static inline float sigmoid(float x) {
    return 1.0f / (1.0f + expf(-x));
}

static int yolov5_decode(Detector *det, rknn_output *outputs,
                          DetResult *results, int max_results) {
    int num = 0;
    int num_classes = -1;

    for (uint32_t o = 0; o < det->io_num.n_output; o++) {
        rknn_tensor_attr *oa = &det->output_attrs[o];
        if (oa->n_dims != 3) continue;

        int C = oa->dims[2];
        if (C != 85 && C != 56) continue;
        if (num_classes < 0) num_classes = C - 5;
    }
    if (num_classes <= 0) return 0;

    int anchors[][3][2] = {
        {{10,13}, {16,30}, {33,23}},
        {{30,61}, {62,45}, {59,119}},
        {{116,90}, {156,198}, {373,326}},
    };
    int strides[] = {8, 16, 32};

    for (uint32_t o = 0; o < det->io_num.n_output && o < 3; o++) {
        rknn_tensor_attr *oa = &det->output_attrs[o];
        if (oa->n_dims != 3) continue;

        int C = oa->dims[2];
        if (C != 85 && C != 56) continue;

        int num_anchors = oa->dims[1];
        int grid_size = num_anchors / 3;
        int grid_h = grid_size, grid_w = grid_size;
        for (int g = (int)sqrtf(grid_size); g >= 1; g--) {
            if (grid_size % g == 0) { grid_h = g; grid_w = grid_size / g; break; }
        }

        float *data = (float *)outputs[o].buf;
        int stride = strides[o];

        for (int a = 0; a < 3 && num < max_results; a++) {
            for (int gy = 0; gy < grid_h && num < max_results; gy++) {
                for (int gx = 0; gx < grid_w && num < max_results; gx++) {
                    int idx = a * grid_h * grid_w + gy * grid_w + gx;
                    float *p = data + idx * C;

                    float obj_conf = sigmoid(p[4]);
                    if (obj_conf < det->conf_thresh) continue;

                    float max_cls = 0;
                    int best_cls = 0;
                    for (int c = 0; c < num_classes; c++) {
                        float s = sigmoid(p[5 + c]);
                        if (s > max_cls) { max_cls = s; best_cls = c; }
                    }

                    float score = obj_conf * max_cls;
                    if (score < det->conf_thresh) continue;

                    float cx = (sigmoid(p[0]) * 2.0f - 0.5f + gx) * stride;
                    float cy = (sigmoid(p[1]) * 2.0f - 0.5f + gy) * stride;
                    float bw = powf(sigmoid(p[2]) * 2.0f, 2) * anchors[o][a][0];
                    float bh = powf(sigmoid(p[3]) * 2.0f, 2) * anchors[o][a][1];

                    results[num].x1 = cx - bw / 2;
                    results[num].y1 = cy - bh / 2;
                    results[num].x2 = cx + bw / 2;
                    results[num].y2 = cy + bh / 2;
                    results[num].score = score;
                    results[num].label = best_cls;
                    num++;
                }
            }
        }
    }
    return num;
}

static int yolov8_decode(Detector *det, rknn_output *outputs,
                          DetResult *results, int max_results) {
    int num = 0;

    for (uint32_t o = 0; o < det->io_num.n_output; o++) {
        rknn_tensor_attr *oa = &det->output_attrs[o];
        int n_dims = oa->n_dims;
        uint32_t *dims = oa->dims;

        if (n_dims != 3) continue;

        int C, H, W;
        int num_classes;

        if (dims[1] == 84 || dims[1] == 56) {
            C = dims[1]; H = 1; W = dims[2];
            num_classes = C - 4;
        } else if (dims[2] == 84 || dims[2] == 56) {
            C = dims[2]; H = dims[1]; W = 1;
            num_classes = C - 4;
        } else {
            H = dims[1]; W = dims[2];
            C = (n_dims > 3) ? dims[3] : 1;
            num_classes = C - 4;
        }

        if (num_classes <= 0) continue;

        float *data = (float *)outputs[o].buf;

        int num_candidates;
        int ch_stride;
        if (dims[1] == 84 || dims[1] == 56) {
            num_candidates = dims[2];
            ch_stride = 1;
        } else if (dims[2] == 84 || dims[2] == 56) {
            num_candidates = dims[1];
            ch_stride = 1;
        } else {
            num_candidates = H * W;
            ch_stride = num_candidates;
        }

        for (int i = 0; i < num_candidates && num < max_results; i++) {
            float cx, cy, bw, bh;
            float max_score = 0;
            int best_label = -1;

            if (dims[1] == 84 || dims[1] == 56) {
                cx = data[0 * num_candidates + i];
                cy = data[1 * num_candidates + i];
                bw = data[2 * num_candidates + i];
                bh = data[3 * num_candidates + i];
                for (int c = 0; c < num_classes; c++) {
                    float s = data[(4 + c) * num_candidates + i];
                    if (s > max_score) { max_score = s; best_label = c; }
                }
            } else if (dims[2] == 84 || dims[2] == 56) {
                cx = data[i * C + 0];
                cy = data[i * C + 1];
                bw = data[i * C + 2];
                bh = data[i * C + 3];
                for (int c = 0; c < num_classes; c++) {
                    float s = data[i * C + 4 + c];
                    if (s > max_score) { max_score = s; best_label = c; }
                }
            } else {
                int idx = i;
                cx = data[0 * ch_stride + idx];
                cy = data[1 * ch_stride + idx];
                bw = data[2 * ch_stride + idx];
                bh = data[3 * ch_stride + idx];
                for (int c = 0; c < num_classes; c++) {
                    float s = data[(4 + c) * ch_stride + idx];
                    if (s > max_score) { max_score = s; best_label = c; }
                }
            }

            if (max_score > det->conf_thresh && best_label >= 0) {
                float x1 = cx - bw / 2;
                float y1 = cy - bh / 2;
                float x2 = cx + bw / 2;
                float y2 = cy + bh / 2;
                results[num].x1 = x1;
                results[num].y1 = y1;
                results[num].x2 = x2;
                results[num].y2 = y2;
                results[num].score = max_score;
                results[num].label = best_label;
                num++;
            }
        }
    }
    return num;
}

static int cmp_score_desc(const void *a, const void *b) {
    float sa = ((const DetResult *)a)->score;
    float sb = ((const DetResult *)b)->score;
    if (sb > sa) return 1;
    if (sb < sa) return -1;
    return 0;
}

static int nms(DetResult *boxes, int num_boxes, float iou_thresh) {
    if (num_boxes == 0) return 0;
    int *keep = malloc(num_boxes * sizeof(int));
    int *removed = calloc(num_boxes, sizeof(int));
    int kept = 0;

    for (int i = 0; i < num_boxes; i++) {
        if (removed[i]) continue;
        keep[kept++] = i;
        for (int j = i + 1; j < num_boxes; j++) {
            if (removed[j]) continue;
            float x1 = fmaxf(boxes[i].x1, boxes[j].x1);
            float y1 = fmaxf(boxes[i].y1, boxes[j].y1);
            float x2 = fminf(boxes[i].x2, boxes[j].x2);
            float y2 = fminf(boxes[i].y2, boxes[j].y2);
            float inter = fmaxf(0, x2 - x1) * fmaxf(0, y2 - y1);
            float a1 = (boxes[i].x2 - boxes[i].x1) * (boxes[i].y2 - boxes[i].y1);
            float a2 = (boxes[j].x2 - boxes[j].x1) * (boxes[j].y2 - boxes[j].y1);
            float iou = inter / (a1 + a2 - inter + 1e-6f);
            if (iou > iou_thresh) removed[j] = 1;
        }
    }

    for (int i = 0; i < kept; i++)
        boxes[i] = boxes[keep[i]];

    free(keep);
    free(removed);
    return kept;
}

int detector_detect(Detector *det, int src_nv12_fd, int src_w, int src_h,
                     DetResult *results, int *num_results) {
    *num_results = 0;

    if (!det || !det->rgb_initialized) return -1;

    int mw = det->model_w;
    int mh = det->model_h;

    if (rga_scale_nv12_to_rgb(src_nv12_fd, src_w, src_h,
                                det->rgb_fd, mw, mh,
                                det->rgb_map, det->rgb_size) < 0)
        return -1;

    /* RGA wrote the RGB buffer via DMA; invalidate CPU cache before
       rknn_inputs_set reads it, otherwise detection sees stale pixels. */
    {
        struct dma_buf_sync sync = { .flags = DMA_BUF_SYNC_START | DMA_BUF_SYNC_READ };
        ioctl(det->rgb_fd, DMA_BUF_IOCTL_SYNC, &sync);
    }

    rknn_input inputs[1];
    memset(inputs, 0, sizeof(inputs));
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].pass_through = 0;
    inputs[0].buf = det->rgb_map;
    inputs[0].size = det->rgb_size;

    int ret = rknn_inputs_set(det->ctx, 1, inputs);
    if (ret < 0) { DBG("det: rknn_inputs_set failed %d\n", ret); return -1; }

    ret = rknn_run(det->ctx, NULL);
    if (ret < 0) { DBG("det: rknn_run failed %d\n", ret); return -1; }

    rknn_output outputs[det->io_num.n_output];
    memset(outputs, 0, sizeof(outputs));
    for (uint32_t i = 0; i < det->io_num.n_output; i++) {
        outputs[i].want_float = 1;
        outputs[i].is_prealloc = 0;
        outputs[i].index = i;
    }

    ret = rknn_outputs_get(det->ctx, det->io_num.n_output, outputs, NULL);
    if (ret < 0) { DBG("det: rknn_outputs_get failed %d\n", ret); return -1; }

    int num_raw;
    if (det->model_type == MODEL_YOLOV5)
        num_raw = yolov5_decode(det, outputs, results, MAX_DETECTIONS);
    else
        num_raw = yolov8_decode(det, outputs, results, MAX_DETECTIONS);
    qsort(results, num_raw, sizeof(DetResult), cmp_score_desc);
    int num_kept = nms(results, num_raw, det->nms_thresh);
    *num_results = num_kept;

    rknn_outputs_release(det->ctx, det->io_num.n_output, outputs);
    return 0;
}

static const uint8_t BOX_COLORS[][3] = {
    {255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0},
    {255, 0, 255}, {0, 255, 255}, {128, 255, 0}, {255, 128, 0},
    {0, 128, 255}, {255, 0, 128},
};

void detector_draw_one_result(DetResult *r, void *y_plane, int out_w, int out_h,
                               int tile_x, int tile_y, int tile_w, int tile_h,
                               int src_w, int src_h) {
    if (!r || !y_plane) return;

    int mw, mh;
    mw = (tile_w > 0) ? tile_w : out_w;
    mh = (tile_h > 0) ? tile_h : out_h;

    float scale_x = (float)mw / src_w;
    float scale_y = (float)mh / src_h;

    int dx1 = tile_x + (int)(r->x1 * scale_x);
    int dy1 = tile_y + (int)(r->y1 * scale_y);
    int dx2 = tile_x + (int)(r->x2 * scale_x);
    int dy2 = tile_y + (int)(r->y2 * scale_y);

    if (dx1 < 0) dx1 = 0;
    if (dy1 < 0) dy1 = 0;
    if (dx2 >= out_w) dx2 = out_w - 1;
    if (dy2 >= out_h) dy2 = out_h - 1;
    if (dx1 >= dx2 || dy1 >= dy2) return;

    uint8_t *y = (uint8_t *)y_plane;
    uint8_t *uv = y + out_w * out_h;
    int color_idx = r->label % 10;
    uint8_t R = BOX_COLORS[color_idx][0];
    uint8_t G = BOX_COLORS[color_idx][1];
    uint8_t B = BOX_COLORS[color_idx][2];
    uint8_t luma = (uint8_t)(0.299f * R + 0.587f * G + 0.114f * B);
    uint8_t Cb = (uint8_t)(128.0f - 0.169f * R - 0.331f * G + 0.500f * B);
    uint8_t Cr = (uint8_t)(128.0f + 0.500f * R - 0.419f * G - 0.081f * B);

    for (int x = dx1; x <= dx2; x++) {
        y[dy1 * out_w + x] = luma;
        y[dy2 * out_w + x] = luma;
    }
    for (int yc = dy1; yc <= dy2; yc++) {
        y[yc * out_w + dx1] = luma;
        y[yc * out_w + dx2] = luma;
    }

    int uv_x1 = dx1 & ~1;
    int uv_x2 = (dx2 + 1) & ~1;
    int uv_y1 = dy1 & ~1;
    int uv_y2 = (dy2 + 1) & ~1;
    for (int ux = uv_x1; ux < uv_x2; ux += 2) {
        uv[uv_y1 * out_w + ux] = Cb;
        uv[uv_y1 * out_w + ux + 1] = Cr;
        uv[uv_y2 * out_w + ux] = Cb;
        uv[uv_y2 * out_w + ux + 1] = Cr;
    }
    for (int uy = uv_y1; uy <= uv_y2; uy += 2) {
        uv[uy * out_w + uv_x1] = Cb;
        uv[uy * out_w + uv_x1 + 1] = Cr;
        if (uv_x2 < out_w) {
            uv[uy * out_w + uv_x2] = Cb;
            uv[uy * out_w + uv_x2 + 1] = Cr;
        }
    }
}