#include <iostream>

#include <gst/gst.h>
#include <nvbufsurface.h>

const unsigned int WIDTH = 1280; /* or 848 */
const unsigned int HEIGHT = 720; /* 480 */
const unsigned int FRAMERATE = 5;
const char *COLOR_FORMAT = "RGB";

static GstPadProbeReturn streammux_src_pad_buffer_probe(GstPad *pad, GstPadProbeInfo *info,
                                                        gpointer u_data) {
    GstBuffer *buf = (GstBuffer *)info->data;

    GstMapInfo map;
    if (!gst_buffer_map(buf, &map, GST_MAP_READ)) {
        gst_buffer_unmap(buf, &map);
        return GST_PAD_PROBE_OK;
    }

    NvBufSurface *surf = (NvBufSurface *)map.data;
    std::cout << "colorFormat = " << surf->surfaceList[0].colorFormat << std::endl;
    gst_buffer_unmap(buf, &map);

    return GST_PAD_PROBE_OK;
}

static gboolean bus_call(GstBus *bus, GstMessage *msg, gpointer data) {
    GMainLoop *loop = (GMainLoop *)data;
    switch (GST_MESSAGE_TYPE(msg)) {
        case GST_MESSAGE_EOS:
            g_print("End of stream\n");
            g_main_loop_quit(loop);
            break;
        case GST_MESSAGE_ERROR: {
            gchar *debug;
            GError *error;
            gst_message_parse_error(msg, &error, &debug);
            g_printerr("ERROR from element %s: %s\n", GST_OBJECT_NAME(msg->src), error->message);
            if (debug) g_printerr("Error details: %s\n", debug);
            g_free(debug);
            g_error_free(error);
            g_main_loop_quit(loop);
            break;
        }
        default:
            break;
    }
    return TRUE;
}

int main(int argc, char *argv[]) {
    GMainLoop *loop = NULL;
    GstElement *pipeline = NULL;
    GstElement *src = NULL, *vidconv = NULL, *capsfilter = NULL, *nvvidconv = NULL,
               *nvcapsfilter = NULL;
    GstElement *streammux = NULL, *nvconv = NULL, *nvtiler = NULL, *nvosd = NULL, *transform = NULL,
               *sink = NULL;

    /* Standard GStreamer initialization */
    gst_init(&argc, &argv);
    loop = g_main_loop_new(NULL, FALSE);

    GstCaps *src_caps = gst_caps_new_simple(
        "video/x-raw", "format", G_TYPE_STRING, COLOR_FORMAT, "width", G_TYPE_INT, WIDTH, "height",
        G_TYPE_INT, HEIGHT, "framerate", GST_TYPE_FRACTION, FRAMERATE, 1, NULL);
    GstCaps *nv_caps = gst_caps_from_string("video/x-raw(memory:NVMM), format=NV12");

    /* Create gstreamer elements */
    pipeline = gst_pipeline_new(NULL);
    src = gst_element_factory_make("videotestsrc", NULL);
    capsfilter = gst_element_factory_make("capsfilter", NULL);
    vidconv = gst_element_factory_make("videoconvert", NULL);
    nvvidconv = gst_element_factory_make("nvvideoconvert", NULL);
    nvcapsfilter = gst_element_factory_make("capsfilter", NULL);

    streammux = gst_element_factory_make("nvstreammux", NULL);
    nvconv = gst_element_factory_make("nvvideoconvert", NULL);
    nvtiler = gst_element_factory_make("nvmultistreamtiler", NULL);
    nvosd = gst_element_factory_make("nvdsosd", NULL);
    transform = gst_element_factory_make("nvegltransform", NULL);
    sink = gst_element_factory_make("nveglglessink", NULL);

    /* Set properties */
    g_object_set(G_OBJECT(src), "is-live", TRUE, NULL);
    g_object_set(G_OBJECT(capsfilter), "caps", src_caps, NULL);
    g_object_set(G_OBJECT(nvcapsfilter), "caps", nv_caps, NULL);
    g_object_set(G_OBJECT(streammux), "batch-size", 1, "width", WIDTH, "height", HEIGHT,
                 "live-source", 1, NULL);
    g_object_set(G_OBJECT(nvtiler), "width", WIDTH, "height", HEIGHT, "columns", 2, "rows", 2,
                 NULL);

    /* We add a message handler */
    GstBus *bus = gst_pipeline_get_bus(GST_PIPELINE(pipeline));
    guint bus_watch_id = gst_bus_add_watch(bus, bus_call, loop);
    gst_object_unref(bus);

    /* Set up the pipeline */
    /* We add all elements into the pipeline */
    gst_bin_add_many(GST_BIN(pipeline), src, capsfilter, vidconv, nvvidconv, nvcapsfilter, NULL);
    gst_bin_add_many(GST_BIN(pipeline), streammux, nvconv, nvtiler, nvosd, transform, sink, NULL);

    GstPad *sinkpad, *srcpad;
    sinkpad = gst_element_get_request_pad(streammux, "sink_0");
    srcpad = gst_element_get_static_pad(nvcapsfilter, "src");
    if (gst_pad_link(srcpad, sinkpad) != GST_PAD_LINK_OK) {
        g_printerr("Failed to link decoder to stream muxer. Exiting.\n");
        return -1;
    }
    gst_object_unref(sinkpad);
    gst_object_unref(srcpad);

    /* We link the elements together */
    if (!gst_element_link_many(src, capsfilter, vidconv, nvvidconv, nvcapsfilter, NULL)) {
        g_printerr("Elements could not be linked: 1. Exiting.\n");
        return -1;
    }
    if (!gst_element_link_many(streammux, nvconv, nvtiler, nvosd, transform, sink, NULL)) {
        g_printerr("Elements could not be linked: 2. Exiting.\n");
        return -1;
    }

    /* We add probe */
    GstPad *streammux_src_pad = gst_element_get_static_pad(streammux, "src");
    gst_pad_add_probe(streammux_src_pad, GST_PAD_PROBE_TYPE_BUFFER, streammux_src_pad_buffer_probe,
                      NULL, NULL);

    /* Set the pipeline to "playing" state */
    g_print("Now playing: %s\n", argv[1]);
    gst_element_set_state(pipeline, GST_STATE_PLAYING);

    /* Wait till pipeline encounters an error or EOS */
    g_print("Running...\n");
    g_main_loop_run(loop);

    /* Out of the main loop, clean up nicely */
    g_print("Returned, stopping playback\n");
    gst_element_set_state(pipeline, GST_STATE_NULL);
    g_print("Deleting pipeline\n");
    gst_object_unref(GST_OBJECT(pipeline));
    g_source_remove(bus_watch_id);
    g_main_loop_unref(loop);

    gst_caps_unref(src_caps);
    gst_caps_unref(nv_caps);

    return 0;
}