#include "gstnvdsmeta.h"
#include "nvbufsurface.h"
#include "nvds_obj_encode.h"
#include <gst/gst.h>
#include <gst/app/gstappsrc.h>
#include <stdio.h>
#include <unistd.h>
#include <stdbool.h>
#include <stdlib.h>

#define MUXER_OUTPUT_WIDTH 960
#define MUXER_OUTPUT_HEIGHT 544
#define MUXER_BATCH_SIZE 2
#define MUXER_BATCH_TIMEOUT_USEC 35000
#define NUM_SOURCES 2
#define MODEL_CONFIG_FILE "/opt/nvidia/deepstream/deepstream-5.0/samples/configs/deepstream-app/config_infer_primary_nano.txt"

// bus callback
static gboolean
bus_call(GstBus *bus, GstMessage *msg, gpointer data)
{
    GMainLoop *loop = (GMainLoop *)data;
    // busData *bus_data = (busData *)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;
}

// create live source bin
static GstElement *
create_live_source_bin(guint index)
{
    GstElement *bin = NULL,
               *source = NULL,
               *vidconv = NULL,
               *nvvidconv = NULL,
               *cap_filter1 = NULL,
               *cap_filter2 = NULL;
    GstPad *real_pad, *ghost_pad;
    GstCaps *caps1, *caps2 = NULL;

    /* Create a source GstBin to abstract this bin's content from the rest of the pipeline */
    gchar bin_name[16] = {};
    g_snprintf(bin_name, 15, "source-bin-%02d", index);
    bin = gst_bin_new(bin_name);

    source = gst_element_factory_make("v4l2src", "camera");
    cap_filter1 = gst_element_factory_make("capsfilter", "cap_filter1");

    if (!bin || !source || !cap_filter1)
    {
        g_printerr("One element in source bin could not be created.\n");
        return NULL;
    }

    gchar device_name[16] = {};
    g_snprintf(device_name, 15, "/dev/video%d", index);
    g_object_set(G_OBJECT(source), "device", device_name, NULL);

    // we need this cap for v4l2src so that it capture video at our specified config
    caps1 = gst_caps_from_string("video/x-raw, width=640, height=480, framerate=25/1");
    g_object_set(G_OBJECT(cap_filter1), "caps", caps1, NULL);
    gst_caps_unref(caps1);
    
    vidconv = gst_element_factory_make("videoconvert", "vidconv");
    nvvidconv = gst_element_factory_make("nvvideoconvert", "nvvidconv");
    cap_filter2 = gst_element_factory_make("capsfilter", "cap_filter2");

    if (!vidconv || !nvvidconv || !cap_filter2)
    {
        g_printerr("One element in source bin could not be created.\n");
        return NULL;
    }

    /* Set elements properties */
    g_object_set(G_OBJECT(nvvidconv), "nvbuf-memory-type", 0, NULL);

    caps2 = gst_caps_from_string("video/x-raw(memory:NVMM), format=NV12");
    g_object_set(G_OBJECT(cap_filter2), "caps", caps2, NULL);
    gst_caps_unref(caps2);

    /* Add elements to the bin */
    gst_bin_add_many(GST_BIN(bin), source, cap_filter1, vidconv, nvvidconv, cap_filter2, NULL);

    /* Link elements in the bin*/
    if (!gst_element_link_many(source, cap_filter1, vidconv, nvvidconv, cap_filter2, NULL))
    {
        g_printerr("One element in source bin could not be linked.\n");
        return NULL;
    }

    /* Create a ghost pad for the bin */
    real_pad = gst_element_get_static_pad(cap_filter2, "src");
    ghost_pad = gst_ghost_pad_new("src", real_pad);
    gst_pad_set_active(ghost_pad, TRUE);
    if (!gst_element_add_pad(bin, ghost_pad))
    {
        g_printerr("Failed to add ghost pad in source bin\n");
        return NULL;
    }
    gst_object_unref(real_pad);

    return bin;
}

// create display sink bin
static GstElement *
create_display_sink_bin(guint index)
{
    GstElement *bin = NULL,
               *transform = NULL,
               *video_sink = NULL,
               *queue = NULL;
    GstPad *real_pad, *ghost_pad;

    gchar bin_name[16] = {};
    gchar ele_name[32] = {};
    g_snprintf(bin_name, 15, "sink-bin-%02d", index);
    bin = gst_bin_new(bin_name);

    g_snprintf(ele_name, 18, "nvegl-transform%02d", index);
    transform = gst_element_factory_make("nvegltransform", ele_name);
    g_snprintf(ele_name, 19, "nvvideo-renderer%02d", index);
    video_sink = gst_element_factory_make("nveglglessink", ele_name);
    g_snprintf(ele_name, 8, "queue%02d", index);
    queue = gst_element_factory_make("queue", ele_name);

    g_object_set(G_OBJECT(video_sink), "sync", FALSE, NULL);

    if (!bin || !transform || !video_sink || !queue)
    {
        g_printerr("One element in sink bin could not be created.\n");
        return NULL;
    }

    /* Add elements to the bin */
    gst_bin_add_many(GST_BIN(bin), queue, transform, video_sink, NULL);

    /* Link elements in the bin*/
    if (!gst_element_link_many(queue, transform, video_sink, NULL))
    {
        g_printerr("One element in sink bin could not be linked.\n");
        return NULL;
    }

    /* Create a ghost pad for the bin */
    real_pad = gst_element_get_static_pad(queue, "sink");
    ghost_pad = gst_ghost_pad_new("sink", real_pad);
    gst_pad_set_active(ghost_pad, TRUE);
    if (!gst_element_add_pad(bin, ghost_pad))
    {
        g_printerr("Failed to add ghost pad in sink bin\n");
        return NULL;
    }
    gst_object_unref(real_pad);

    return bin;
}

/* { Tracker config parsing */

#define CHECK_ERROR(error)                                                   \
    if (error)                                                               \
    {                                                                        \
        g_printerr("Error while parsing config file: %s\n", error->message); \
        goto done;                                                           \
    }

#define TRACKER_CONFIG_FILE "/opt/nvidia/deepstream/deepstream-5.0/sources/apps/sample_apps/deepstream-test2/dstest2_tracker_config.txt"
#define MAX_TRACKING_ID_LEN 16
#define CONFIG_GROUP_TRACKER "tracker"
#define CONFIG_GROUP_TRACKER_WIDTH "tracker-width"
#define CONFIG_GROUP_TRACKER_HEIGHT "tracker-height"
#define CONFIG_GROUP_TRACKER_LL_CONFIG_FILE "ll-config-file"
#define CONFIG_GROUP_TRACKER_LL_LIB_FILE "ll-lib-file"
#define CONFIG_GROUP_TRACKER_ENABLE_BATCH_PROCESS "enable-batch-process"
#define CONFIG_GPU_ID "gpu-id"

static gchar *
get_absolute_file_path(gchar *cfg_file_path, gchar *file_path)
{
    gchar abs_cfg_path[PATH_MAX + 1];
    gchar *abs_file_path;
    gchar *delim;

    if (file_path && file_path[0] == '/')
    {
        return file_path;
    }

    if (!realpath(cfg_file_path, abs_cfg_path))
    {
        g_free(file_path);
        return NULL;
    }

    // Return absolute path of config file if file_path is NULL.
    if (!file_path)
    {
        abs_file_path = g_strdup(abs_cfg_path);
        return abs_file_path;
    }

    delim = g_strrstr(abs_cfg_path, "/");
    *(delim + 1) = '\0';

    abs_file_path = g_strconcat(abs_cfg_path, file_path, NULL);
    g_free(file_path);

    return abs_file_path;
}

static gboolean
set_tracker_properties(GstElement *nvtracker)
{
    gboolean ret = FALSE;
    GError *error = NULL;
    gchar **keys = NULL;
    gchar **key = NULL;
    GKeyFile *key_file = g_key_file_new();

    if (!g_key_file_load_from_file(key_file, TRACKER_CONFIG_FILE, G_KEY_FILE_NONE,
                                   &error))
    {
        g_printerr("Failed to load config file: %s\n", error->message);
        return FALSE;
    }

    keys = g_key_file_get_keys(key_file, CONFIG_GROUP_TRACKER, NULL, &error);
    CHECK_ERROR(error);

    for (key = keys; *key; key++)
    {
        if (!g_strcmp0(*key, CONFIG_GROUP_TRACKER_WIDTH))
        {
            gint width =
                g_key_file_get_integer(key_file, CONFIG_GROUP_TRACKER,
                                       CONFIG_GROUP_TRACKER_WIDTH, &error);
            CHECK_ERROR(error);
            g_object_set(G_OBJECT(nvtracker), "tracker-width", width, NULL);
        }
        else if (!g_strcmp0(*key, CONFIG_GROUP_TRACKER_HEIGHT))
        {
            gint height =
                g_key_file_get_integer(key_file, CONFIG_GROUP_TRACKER,
                                       CONFIG_GROUP_TRACKER_HEIGHT, &error);
            CHECK_ERROR(error);
            g_object_set(G_OBJECT(nvtracker), "tracker-height", height, NULL);
        }
        else if (!g_strcmp0(*key, CONFIG_GPU_ID))
        {
            guint gpu_id =
                g_key_file_get_integer(key_file, CONFIG_GROUP_TRACKER,
                                       CONFIG_GPU_ID, &error);
            CHECK_ERROR(error);
            g_object_set(G_OBJECT(nvtracker), "gpu_id", gpu_id, NULL);
        }
        else if (!g_strcmp0(*key, CONFIG_GROUP_TRACKER_LL_CONFIG_FILE))
        {
            char *ll_config_file = get_absolute_file_path(TRACKER_CONFIG_FILE,
                                                          g_key_file_get_string(key_file,
                                                                                CONFIG_GROUP_TRACKER,
                                                                                CONFIG_GROUP_TRACKER_LL_CONFIG_FILE, &error));
            CHECK_ERROR(error);
            g_object_set(G_OBJECT(nvtracker), "ll-config-file", ll_config_file, NULL);
        }
        else if (!g_strcmp0(*key, CONFIG_GROUP_TRACKER_LL_LIB_FILE))
        {
            char *ll_lib_file = get_absolute_file_path(TRACKER_CONFIG_FILE,
                                                       g_key_file_get_string(key_file,
                                                                             CONFIG_GROUP_TRACKER,
                                                                             CONFIG_GROUP_TRACKER_LL_LIB_FILE, &error));
            CHECK_ERROR(error);
            g_object_set(G_OBJECT(nvtracker), "ll-lib-file", ll_lib_file, NULL);
        }
        else if (!g_strcmp0(*key, CONFIG_GROUP_TRACKER_ENABLE_BATCH_PROCESS))
        {
            gboolean enable_batch_process =
                g_key_file_get_integer(key_file, CONFIG_GROUP_TRACKER,
                                       CONFIG_GROUP_TRACKER_ENABLE_BATCH_PROCESS, &error);
            CHECK_ERROR(error);
            g_object_set(G_OBJECT(nvtracker), "enable_batch_process",
                         enable_batch_process, NULL);
        }
        else
        {
            g_printerr("Unknown key '%s' for group [%s]", *key,
                       CONFIG_GROUP_TRACKER);
        }
    }

    ret = TRUE;
done:
    if (error)
    {
        g_error_free(error);
    }
    if (keys)
    {
        g_strfreev(keys);
    }
    if (!ret)
    {
        g_printerr("%s failed", __func__);
    }
    return ret;
}

/* Tracker config parsing }*/

int main(int argc, char *argv[])
{
    
    /*Gstreamer Variables declaration*/
    GMainLoop *loop = NULL;
    GstElement *pipeline = NULL,
               *streammux = NULL,
               *pgie = NULL,
               *nvtracker = NULL,
               *demux = NULL;

    GstBus *bus = NULL;
    guint bus_watch_id;

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

    /* Create gstreamer elements */
    pipeline = gst_pipeline_new("ds-custom-pipeline");
    streammux = gst_element_factory_make("nvstreammux", "stream_muxer");
    pgie = gst_element_factory_make("nvinfer", "primary-nvinference-engine");
    nvtracker = gst_element_factory_make("nvtracker", "tracker");
    demux = gst_element_factory_make("nvstreamdemux", "demuxer");

    // exit if fail to create any element
    if (!pipeline)
    {
        g_printerr("One element could not be created 1. Exiting.\n");
        return EXIT_FAILURE;
    }
    if (!streammux || !pgie || !nvtracker || !demux)
    {
        g_printerr("One element could not be created 2. Exiting.\n");
        return EXIT_FAILURE;
    }

    /* Set gstreamer elements properties */
    g_object_set(G_OBJECT(streammux),
                 "width", MUXER_OUTPUT_WIDTH,
                 "height", MUXER_OUTPUT_HEIGHT,
                 "batch-size", MUXER_BATCH_SIZE,
                 "batched-push-timeout", MUXER_BATCH_TIMEOUT_USEC,
                 "live-source", TRUE,
                 NULL);

    g_object_set(G_OBJECT(pgie), "config-file-path", MODEL_CONFIG_FILE, NULL);

    /* Set necessary properties of the tracker element. */
    if (!set_tracker_properties(nvtracker))
    {
        g_printerr("Failed to set tracker properties. Exiting.\n");
        return -1;
    }

    /* Add a bus message handler */
    bus = gst_pipeline_get_bus(GST_PIPELINE(pipeline));
    bus_watch_id = gst_bus_add_watch(bus, bus_call, loop);
    gst_object_unref(bus);

    /* Add common elements into the pipeline */
    gst_bin_add_many(GST_BIN(pipeline),
                     streammux, pgie, nvtracker,
                     demux,
                     NULL);

    /* Link elements */
    // Link all elements that can be automatically linked because they have "Always" pads
    if (!gst_element_link_many(streammux, pgie, nvtracker,
                               demux, NULL))
    {
        g_printerr("Elements could not be linked: 1. Exiting.\n");
        return EXIT_FAILURE;
    }

    /* Set up source bins (must done this after streammux is added to the pipeline) */
    for (guint i = 0; i < NUM_SOURCES; i++)
    {
        GstElement *source_bin = NULL;
        source_bin = create_live_source_bin(i);
        
        if (!source_bin)
        {
            g_printerr("Failed to create source bin. Exiting.\n");
            return EXIT_FAILURE;
        }

        // add each source bin to the pipeline
        gst_bin_add(GST_BIN(pipeline), source_bin);

        // manually link streammux to each source bin
        GstPad *sinkpad, *srcpad;
        gchar pad_name[16] = {};
        g_snprintf(pad_name, 15, "sink_%u", i);
        sinkpad = gst_element_get_request_pad(streammux, pad_name);
        if (!sinkpad)
        {
            g_printerr("Streammux request sink pad failed. Exiting.\n");
            return EXIT_FAILURE;
        }

        srcpad = gst_element_get_static_pad(source_bin, "src");
        if (!srcpad)
        {
            g_printerr("Failed to get src pad of source bin. Exiting.\n");
            return EXIT_FAILURE;
        }

        if (gst_pad_link(srcpad, sinkpad) != GST_PAD_LINK_OK)
        {
            g_printerr("Failed to link source bin to stream muxer. Exiting.\n");
            return EXIT_FAILURE;
        }

        gst_object_unref(srcpad);
        gst_object_unref(sinkpad);
    }

    for (guint i = 0; i < NUM_SOURCES;i++)
    {
        GstElement *sink_bin = NULL;
        
        sink_bin = create_display_sink_bin(i);

        if (!sink_bin)
        {
            g_printerr("Failed to create sink bin. Exiting.\n");
            return EXIT_FAILURE;
        }

        // add each sink bin to the pipeline
        gst_bin_add(GST_BIN(pipeline), sink_bin);

        // manually link demux to each sink bin
        GstPad *sinkpad, *srcpad;
        gchar pad_name[16] = {};
        g_snprintf(pad_name, 15, "src_%u", i);
        srcpad = gst_element_get_request_pad(demux, pad_name);
        if (!srcpad)
        {
            g_printerr("demux request source pad failed. Exiting.\n");
            return EXIT_FAILURE;
        }

        sinkpad = gst_element_get_static_pad(sink_bin, "sink");
        if (!sinkpad)
        {
            g_printerr("Failed to get sink pad of sink bin. Exiting.\n");
            return EXIT_FAILURE;
        }

        if (gst_pad_link(srcpad, sinkpad) != GST_PAD_LINK_OK)
        {
            g_printerr("Failed to link source bin to stream muxer. Exiting.\n");
            return EXIT_FAILURE;
        }

        gst_object_unref(srcpad);
        gst_object_unref(sinkpad);
    }

    /* 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");
    GST_DEBUG_BIN_TO_DOT_FILE_WITH_TS(GST_BIN(pipeline), GST_DEBUG_GRAPH_SHOW_ALL, "ds-app-playing");
    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");
    g_source_remove(bus_watch_id);
    g_main_loop_unref(loop);
    gst_object_unref(GST_OBJECT(pipeline));

    printf("Exit program\n");
    return EXIT_SUCCESS;
}