#include <cstdlib>
#include <cstring>
#include <sstream>
#include <gst/gst.h>

#include "nvbuf_utils.h"
#include <cuda.h>
#include <cuda_runtime.h>
#include <cudaEGL.h>
#include <opencv2/cudafilters.hpp>

using namespace std;

#define USE(x) ((void)(x))

static GstPipeline *gst_pipeline = nullptr;
static string launch_string;

GstClockTime usec = 1000000;
static int w = 1920;
static int h = 1080;
EGLDisplay egl_display;

static bool create_filter = true;
static cv::Ptr< cv::cuda::Filter > filter;

static GstPadProbeReturn
conv_src_pad_buffer_probe (GstPad * pad, GstPadProbeInfo * info,
    gpointer u_data)
{
    GstBuffer *buffer = (GstBuffer *) info->data;
    GstMapInfo map    = {0};
    int dmabuf_fd = 0;
    gst_buffer_map (buffer, &map, GST_MAP_READ);
    ExtractFdFromNvBuffer((void *)map.data, &dmabuf_fd);
    //CUDA postprocess
    {
        EGLImageKHR egl_image;
        egl_image = NvEGLImageFromFd(egl_display, dmabuf_fd);
        CUresult status;
        CUeglFrame eglFrame;
        CUgraphicsResource pResource = NULL;
        cudaFree(0);
        status = cuGraphicsEGLRegisterImage(&pResource,
                    egl_image,
                    CU_GRAPHICS_MAP_RESOURCE_FLAGS_NONE);
        if (status != CUDA_SUCCESS)
        {
            printf("cuGraphicsEGLRegisterImage failed: %d \n",status);
        }
        status = cuGraphicsResourceGetMappedEglFrame(&eglFrame, pResource, 0, 0);
        status = cuCtxSynchronize();
        if (create_filter) {
            filter = cv::cuda::createSobelFilter(CV_8UC4, CV_8UC4, 1, 0, 3, 1, cv::BORDER_DEFAULT);
            //filter = cv::cuda::createGaussianFilter(CV_8UC4, CV_8UC4, cv::Size(31,31), 0, 0, cv::BORDER_DEFAULT);
            create_filter = false;
            printf("filter created if needed\n");
        }
        printf("about to make GpuMat\n");
        cv::cuda::GpuMat d_mat(h, w, CV_8UC4, eglFrame.frame.pPitch[0]);
        printf("completed make GpuMat: %p\n", &d_mat);
        filter->apply (d_mat, d_mat);
        printf("filter applied\n");

        status = cuCtxSynchronize();
        printf("called cuCtxSynchronize\n");
        status = cuGraphicsUnregisterResource(pResource);
        printf("called cuGraphicsUnregisterResource\n");
        NvDestroyEGLImage(egl_display, egl_image);
    }
    gst_buffer_unmap(buffer, &map);

    return GST_PAD_PROBE_OK;
}

int main(int argc, char** argv) {
    USE(argc);
    USE(argv);

    gst_init (&argc, &argv);

    GMainLoop *main_loop;
    main_loop = g_main_loop_new (NULL, FALSE);
    ostringstream launch_stream;

    // egl_display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
    // eglInitialize(egl_display, NULL, NULL);

    launch_stream
    << "nvarguscamerasrc name=mysource ! "
    << "video/x-raw(memory:NVMM),width="<< w <<",height="<< h <<",framerate=30/1,format=NV12 ! "
    << "nvvidconv name=myconv ! "
    << "video/x-raw(memory:NVMM),format=RGBA ! "
    << "nvv4l2h264enc bitrate=8000000 ! h264parse ! qtmux ! filesink location=testnnnn.mp4 ";

    launch_string = launch_stream.str();

    g_print("Using launch string: %s\n", launch_string.c_str());

    GError *error = nullptr;
    gst_pipeline  = (GstPipeline*) gst_parse_launch(launch_string.c_str(), &error);

    if (gst_pipeline == nullptr) {
        g_print( "Failed to parse launch: %s\n", error->message);
        return -1;
    }
    if(error) g_error_free(error);
    GstElement* conv = gst_bin_get_by_name(GST_BIN(gst_pipeline), "myconv");
    GstPad *src_pad = gst_element_get_static_pad (conv, "src");
    gst_pad_add_probe (src_pad, GST_PAD_PROBE_TYPE_BUFFER,
        conv_src_pad_buffer_probe, NULL, NULL);

    gst_element_set_state((GstElement*)gst_pipeline, GST_STATE_PLAYING); 

    g_usleep(30*usec);

    GstElement* src = gst_bin_get_by_name(GST_BIN(gst_pipeline), "mysource");
    gst_element_send_event (src, gst_event_new_eos ());
    // Wait for EOS message
    GstBus *bus = gst_pipeline_get_bus(GST_PIPELINE(gst_pipeline));
    gst_bus_poll(bus, GST_MESSAGE_EOS, GST_CLOCK_TIME_NONE);

    gst_element_set_state((GstElement*)gst_pipeline, GST_STATE_NULL);
    gst_object_unref(GST_OBJECT(gst_pipeline));
    g_main_loop_unref(main_loop);

    // eglTerminate(egl_display);

    g_print("going to exit \n");
    return 0;
}