#include "decoder.h"
#include "dma_util.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <pthread.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 <gst/gst.h>
#include <gst/app/gstappsink.h>
#include <gst/video/video.h>
#include <gst/allocators/gstdmabuf.h>

#define DEC_NUM_BUF 3

struct Decoder {
    GstElement *pipeline;
    GstElement *appsink;

    /* Zero-copy mode: appsink negotiates video/x-raw(memory:DMABuf) and we
     * hand out dups of the decoder's own dmabuf fds. Pulled samples are kept
     * alive per slot until decoder_release_frame() unrefs them, so GStreamer
     * can never overwrite a buffer the composer is still using. */
    int use_dmabuf;
    GstSample *slot_sample[DEC_NUM_BUF];

    /* Fallback mode: CPU copies each frame into our own dma-heap buffers. */
    int buf_fd[DEC_NUM_BUF];
    void *buf_map[DEC_NUM_BUF];
    size_t buf_size;
    int buf_width;
    int buf_height;
    int buf_stride;
    int buf_stride_v;
    int buf_index;

    int buf_inuse[DEC_NUM_BUF];
    pthread_mutex_t buf_mutex;
    pthread_cond_t buf_cond;

    int codec_type;
    int width;
    int height;
    int stride;
};

static GstElement *build_pipeline(const char *url, int want_dmabuf) {
    char pipeline_str[2048];
    if (want_dmabuf) {
        snprintf(pipeline_str, sizeof(pipeline_str),
                 "rtspsrc name=src location=%s protocols=tcp latency=0 ! "
                 "application/x-rtp,media=video ! parsebin name=pb ! mppvideodec ! "
                 "video/x-raw(memory:DMABuf),format=NV12 ! appsink name=sink "
                 "max-buffers=%d drop=true sync=false ",
                 url, DEC_NUM_BUF);
    } else {
        snprintf(pipeline_str, sizeof(pipeline_str),
                 "rtspsrc name=src location=%s protocols=tcp latency=0 ! "
                 "application/x-rtp,media=video ! parsebin name=pb ! mppvideodec ! "
                 "video/x-raw,format=NV12 ! appsink name=sink "
                 "max-buffers=1 drop=true sync=false ",
                 url);
    }

    GError *err = NULL;
    GstElement *pipeline = gst_parse_launch(pipeline_str, &err);
    if (!pipeline || err) {
        fprintf(stderr, "GStreamer pipeline creation failed: %s\n",
                err ? err->message : "unknown");
        if (err) g_error_free(err);
        if (pipeline) { gst_object_unref(pipeline); pipeline = NULL; }
        return NULL;
    }

    GstElement *appsink = gst_bin_get_by_name(GST_BIN(pipeline), "sink");
    gst_app_sink_set_drop(GST_APP_SINK(appsink), TRUE);
    gst_app_sink_set_max_buffers(GST_APP_SINK(appsink),
                                 want_dmabuf ? DEC_NUM_BUF : 1);
    gst_app_sink_set_emit_signals(GST_APP_SINK(appsink), FALSE);

    /* Set rtspsrc low-latency options */
    GstElement *src = gst_bin_get_by_name(GST_BIN(pipeline), "src");
    if (src) {
        g_object_set(src, "latency", 0, NULL);
        g_object_set(src, "buffer-mode", 0, NULL);
        g_object_set(src, "do-rtcp", FALSE, NULL);
        gst_object_unref(src);
    }
    gst_object_unref(appsink);
    return pipeline;
}

static void destroy_pipeline(Decoder *dec) {
    if (dec->pipeline) {
        gst_element_set_state(dec->pipeline, GST_STATE_NULL);
        gst_object_unref(dec->pipeline);
        dec->pipeline = NULL;
    }
    if (dec->appsink) { gst_object_unref(dec->appsink); dec->appsink = NULL; }
}

/* Fill dec dims from the caps of `sample`; returns ver_stride (UV plane row
 * offset in rows) or -1 on bad caps. */
static int probe_dims(Decoder *dec, GstSample *sample) {
    GstCaps *caps = gst_sample_get_caps(sample);
    GstVideoInfo vinfo;
    if (!caps || !gst_video_info_from_caps(&vinfo, caps)) return -1;
    dec->width = GST_VIDEO_INFO_WIDTH(&vinfo);
    dec->height = GST_VIDEO_INFO_HEIGHT(&vinfo);
    dec->stride = GST_VIDEO_INFO_PLANE_STRIDE(&vinfo, 0);
    int stride_v = GST_VIDEO_INFO_PLANE_STRIDE(&vinfo, 1);
    gsize uv_off = GST_VIDEO_INFO_COMP_OFFSET(&vinfo, 1);
    /* UV plane starts after ver_stride aligned rows of Y */
    if (dec->stride > 0 && uv_off % (gsize)dec->stride == 0)
        return (int)(uv_off / (gsize)dec->stride);
    (void)stride_v;
    return dec->height;
}

Decoder *decoder_create(const char *url, int timeout_sec) {
    static int gst_inited = 0;
    if (!gst_inited) { gst_init(NULL, NULL); gst_inited = 1; }

    Decoder *dec = calloc(1, sizeof(Decoder));
    if (!dec) return NULL;
    for (int i = 0; i < DEC_NUM_BUF; i++) dec->buf_fd[i] = -1;
    pthread_mutex_init(&dec->buf_mutex, NULL);
    pthread_cond_init(&dec->buf_cond, NULL);
    for (int i = 0; i < DEC_NUM_BUF; i++) dec->buf_inuse[i] = 0;

    /* parsebin auto-detects whether the RTP stream is H.264, H.265, etc. from
     * the SDP encoding-name, so cameras need not advertise their codec via the
     * URL (a `.264`-less URL used to be forced to H.265, breaking H.264 cams
     * such as 192.168.1.120). mppvideodec does the hardware decode. */
    dec->codec_type = 1;

    /* --- Preferred path: copy into our own dma-heap buffers. The zero-copy
     * mppvideodec DMABuf path is skipped: on this board RGA cannot import the
     * MPP decoder's own dmabufs (importbuffer_fd fails -> tiles never appear).
     * Our dma-heap (system-uncached) buffers import cleanly into RGA. --- */
copy_path:

    /* --- Fallback: system-memory output + CPU copy into our own heap --- */
    dec->pipeline = build_pipeline(url, 0);
    if (!dec->pipeline) {
        pthread_mutex_destroy(&dec->buf_mutex);
        pthread_cond_destroy(&dec->buf_cond);
        free(dec);
        return NULL;
    }
    dec->appsink = gst_bin_get_by_name(GST_BIN(dec->pipeline), "sink");

    gst_element_set_state(dec->pipeline, GST_STATE_PLAYING);

    /* Wait for first frame with timeout */
    GstSample *sample = gst_app_sink_try_pull_sample(GST_APP_SINK(dec->appsink),
                                                      timeout_sec * GST_SECOND);
    if (!sample) {
        fprintf(stderr, "Timeout waiting for first frame from %s\n", url);
        destroy_pipeline(dec);
        pthread_mutex_destroy(&dec->buf_mutex);
        pthread_cond_destroy(&dec->buf_cond);
        free(dec);
        return NULL;
    }

    /* Get video info */
    if (probe_dims(dec, sample) < 0) {
        gst_sample_unref(sample);
        fprintf(stderr, "Failed to get video info\n");
        destroy_pipeline(dec);
        pthread_mutex_destroy(&dec->buf_mutex);
        pthread_cond_destroy(&dec->buf_cond);
        free(dec);
        return NULL;
    }

    /* Decode first frame to allocate buffer */
    GstBuffer *out_buf = gst_sample_get_buffer(sample);
    if (out_buf) {
        GstCaps *caps = gst_sample_get_caps(sample);
        GstVideoInfo vinfo;
        int stride_v = dec->height;
        if (caps && gst_video_info_from_caps(&vinfo, caps))
            stride_v = GST_VIDEO_INFO_PLANE_STRIDE(&vinfo, 1);
        size_t needed = dec->stride * dec->height + stride_v * dec->height / 2;

        for (int i = 0; i < DEC_NUM_BUF; i++) {
            if (dma_heap_alloc("cma", needed, &dec->buf_fd[i]) < 0) {
                if (dma_heap_alloc("system-uncached", needed, &dec->buf_fd[i]) < 0) {
                    dec->buf_fd[i] = -1;
                }
            }

            if (dec->buf_fd[i] >= 0) {
                dec->buf_map[i] = mmap(NULL, needed, PROT_READ | PROT_WRITE,
                                    MAP_SHARED, dec->buf_fd[i], 0);
                if (dec->buf_map[i] != MAP_FAILED) {
                    if (i == 0) {
                        dec->buf_size = needed;
                        dec->buf_width = dec->width;
                        dec->buf_height = dec->height;
                        dec->buf_stride = dec->stride;
                        dec->buf_stride_v = stride_v;
                        fprintf(stderr, "DBG: decoder buffer %dx%d stride=%d/%d size=%zu\n",
                                dec->width, dec->height, dec->stride, stride_v, needed);
                    }
                } else {
                    close(dec->buf_fd[i]);
                    dec->buf_fd[i] = -1;
                }
            }
        }
    }

    gst_sample_unref(sample);

    if (dec->buf_fd[0] < 0) {
        fprintf(stderr, "Failed to allocate decoder buffer\n");
        destroy_pipeline(dec);
        pthread_mutex_destroy(&dec->buf_mutex);
        pthread_cond_destroy(&dec->buf_cond);
        free(dec);
        return NULL;
    }

    return dec;
}

void decoder_destroy(Decoder *dec) {
    if (!dec) return;
    destroy_pipeline(dec);
    for (int i = 0; i < DEC_NUM_BUF; i++) {
        if (dec->slot_sample[i]) {
            gst_sample_unref(dec->slot_sample[i]);
            dec->slot_sample[i] = NULL;
        }
    }
    for (int i = 0; i < DEC_NUM_BUF; i++) {
        if (dec->buf_fd[i] >= 0) {
            if (dec->buf_map[i]) munmap(dec->buf_map[i], dec->buf_size);
            close(dec->buf_fd[i]);
        }
    }
    pthread_mutex_destroy(&dec->buf_mutex);
    pthread_cond_destroy(&dec->buf_cond);
    free(dec);
}

int decoder_get_codec(Decoder *dec) { return dec->codec_type; }
int decoder_get_width(Decoder *dec) { return dec->width; }
int decoder_get_height(Decoder *dec) { return dec->height; }
int decoder_get_stride(Decoder *dec) { return dec->stride; }

int decoder_get_frame(Decoder *dec, DecodedFrame *frame) {
    if (!dec || !dec->pipeline) return -1;

    /* --- zero-copy dmabuf path --- */
    if (dec->use_dmabuf) {
        GstSample *sample = gst_app_sink_try_pull_sample(GST_APP_SINK(dec->appsink), 0);
        if (!sample) return -3;
        GstBuffer *buf = gst_sample_get_buffer(sample);
        GstMemory *mem = buf ? gst_buffer_peek_memory(buf, 0) : NULL;
        if (!mem || !gst_is_dmabuf_memory(mem)) {
            gst_sample_unref(sample);
            return -1;
        }

        pthread_mutex_lock(&dec->buf_mutex);
        int idx = -1;
        for (int i = 0; i < DEC_NUM_BUF; i++) {
            if (!dec->slot_sample[i]) { idx = i; break; }
        }
        if (idx < 0) {
            /* all slots pinned (ring + held + in-flight): drop the new
             * frame, main catches up next iteration */
            pthread_mutex_unlock(&dec->buf_mutex);
            gst_sample_unref(sample);
            return -3;
        }
        dec->slot_sample[idx] = sample;   /* keep alive until release */
        pthread_mutex_unlock(&dec->buf_mutex);

        int fd = gst_dmabuf_memory_get_fd(mem);
        if (fd < 0) {
            pthread_mutex_lock(&dec->buf_mutex);
            dec->slot_sample[idx] = NULL;
            pthread_mutex_unlock(&dec->buf_mutex);
            gst_sample_unref(sample);
            return -1;
        }

        int ver = probe_dims(dec, sample);
        if (ver <= 0) ver = dec->height;

        frame->fd = fcntl(fd, F_DUPFD_CLOEXEC, 3);
        frame->width = dec->width;
        frame->height = dec->height;
        frame->hor_stride = dec->stride;
        frame->ver_stride = ver;
        frame->idx = idx;
        frame->pts = 0;
        if (frame->fd < 0) {
            pthread_mutex_lock(&dec->buf_mutex);
            dec->slot_sample[idx] = NULL;
            pthread_mutex_unlock(&dec->buf_mutex);
            gst_sample_unref(sample);
            return -1;
        }
        return 0;
    }

    /* --- fallback: system memory + CPU copy into our own heap --- */
    GstSample *sample = gst_app_sink_try_pull_sample(GST_APP_SINK(dec->appsink), 0);
    if (!sample) return -3;

    GstBuffer *out_buf = gst_sample_get_buffer(sample);
    GstCaps *caps = gst_sample_get_caps(sample);
    if (!caps || !out_buf) { gst_sample_unref(sample); return -1; }

    GstVideoInfo vinfo;
    if (!gst_video_info_from_caps(&vinfo, caps)) { gst_sample_unref(sample); return -1; }

    int w = GST_VIDEO_INFO_WIDTH(&vinfo);
    int h = GST_VIDEO_INFO_HEIGHT(&vinfo);
    int stride = GST_VIDEO_INFO_PLANE_STRIDE(&vinfo, 0);
    int stride_v = GST_VIDEO_INFO_PLANE_STRIDE(&vinfo, 1);
    size_t needed = stride * h + stride_v * h / 2;

    /* Reallocate if dimensions changed */
    if (w != dec->buf_width || h != dec->buf_height ||
        stride != dec->buf_stride || needed > dec->buf_size) {
        for (int i = 0; i < DEC_NUM_BUF; i++) {
            if (dec->buf_map[i]) munmap(dec->buf_map[i], dec->buf_size);
            if (dec->buf_fd[i] >= 0) close(dec->buf_fd[i]);
            dec->buf_fd[i] = -1;
            dec->buf_map[i] = NULL;

            if (dma_heap_alloc("cma", needed, &dec->buf_fd[i]) < 0) {
                if (dma_heap_alloc("system-uncached", needed, &dec->buf_fd[i]) < 0) {
                    dec->buf_fd[i] = -1;
                }
            }

            if (dec->buf_fd[i] < 0) { gst_sample_unref(sample); return -1; }

            dec->buf_map[i] = mmap(NULL, needed, PROT_READ | PROT_WRITE,
                                MAP_SHARED, dec->buf_fd[i], 0);
            if (dec->buf_map[i] == MAP_FAILED) {
                close(dec->buf_fd[i]);
                dec->buf_fd[i] = -1;
                gst_sample_unref(sample);
                return -1;
            }
        }
        dec->buf_size = needed;
        dec->buf_width = w;
        dec->buf_height = h;
        dec->buf_stride = stride;
        dec->buf_stride_v = stride_v;
    }

    pthread_mutex_lock(&dec->buf_mutex);
    while (dec->buf_inuse[dec->buf_index]) {
        struct timespec ts;
        clock_gettime(CLOCK_REALTIME, &ts);
        ts.tv_nsec += 200000000;
        if (ts.tv_nsec >= 1000000000) { ts.tv_sec++; ts.tv_nsec -= 1000000000; }
        pthread_cond_timedwait(&dec->buf_cond, &dec->buf_mutex, &ts);
    }
    int idx = dec->buf_index;
    dec->buf_inuse[idx] = 1;
    dec->buf_index = (idx + 1) % DEC_NUM_BUF;
    pthread_mutex_unlock(&dec->buf_mutex);

    GstVideoFrame vframe;
    if (!gst_video_frame_map(&vframe, &vinfo, out_buf, GST_MAP_READ)) {
        gst_sample_unref(sample);
        return -1;
    }

    uint8_t *src_y = GST_VIDEO_FRAME_PLANE_DATA(&vframe, 0);
    uint8_t *src_uv = GST_VIDEO_FRAME_PLANE_DATA(&vframe, 1);
    int src_stride = GST_VIDEO_FRAME_PLANE_STRIDE(&vframe, 0);
    int src_stride_uv = GST_VIDEO_FRAME_PLANE_STRIDE(&vframe, 1);
    uint8_t *dst = (uint8_t *)dec->buf_map[idx];

    /* CPU write start: invalidate/clean so cached dma-heap memory stays coherent */
    {
        struct dma_buf_sync sync = { .flags = DMA_BUF_SYNC_START | DMA_BUF_SYNC_WRITE };
        ioctl(dec->buf_fd[idx], DMA_BUF_IOCTL_SYNC, &sync);
    }

    if (src_stride == stride) {
        memcpy(dst, src_y, stride * h);
    } else {
        for (int y = 0; y < h; y++)
            memcpy(dst + y * stride, src_y + y * src_stride, stride < src_stride ? stride : src_stride);
    }

    dst += stride * h;
    int uv_h = h / 2;
    if (src_stride_uv == stride_v) {
        memcpy(dst, src_uv, stride_v * uv_h);
    } else {
        for (int y = 0; y < uv_h; y++)
            memcpy(dst + y * stride_v, src_uv + y * src_stride_uv, stride_v < src_stride_uv ? stride_v : src_stride_uv);
    }

    gst_video_frame_unmap(&vframe);
    gst_sample_unref(sample);

    /* CPU write done: flush cache so RGA DMA reads the latest frame data.
       Without this, fast motion shows a mix of old/new rows (tearing). */
    {
        struct dma_buf_sync sync = { .flags = DMA_BUF_SYNC_END | DMA_BUF_SYNC_WRITE };
        ioctl(dec->buf_fd[idx], DMA_BUF_IOCTL_SYNC, &sync);
    }

    frame->fd = fcntl(dec->buf_fd[idx], F_DUPFD_CLOEXEC, 3);
    frame->width = w;
    frame->height = h;
    frame->hor_stride = stride;
    frame->ver_stride = h;
    frame->idx = idx;
    frame->pts = 0;
    return 0;
}

int decoder_flush(Decoder *dec) {
    (void)dec;
    return 0;
}

int decoder_release_frame(Decoder *dec, int idx) {
    if (!dec || idx < 0 || idx >= DEC_NUM_BUF) return -1;
    pthread_mutex_lock(&dec->buf_mutex);
    if (dec->slot_sample[idx]) {
        gst_sample_unref(dec->slot_sample[idx]);
        dec->slot_sample[idx] = NULL;
        pthread_cond_signal(&dec->buf_cond);
    } else if (dec->buf_inuse[idx]) {
        dec->buf_inuse[idx] = 0;
        pthread_cond_signal(&dec->buf_cond);
    }
    pthread_mutex_unlock(&dec->buf_mutex);
    return 0;
}