/*
 * Copyright (c) 2018-2022, NVIDIA CORPORATION. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

#include <gst/gst.h>
#include <glib.h>
#include <stdio.h>
#include <cuda_runtime_api.h>
#include "gstnvdsmeta.h"
#include "nvds_yml_parser.h"

#define MAX_DISPLAY_LEN 64

#define PGIE_CLASS_ID_VEHICLE 0
#define PGIE_CLASS_ID_PERSON 2

/* The muxer output resolution must be set if the input streams will be of
 * different resolution. The muxer will scale all the input frames to this
 * resolution. */
#define MUXER_OUTPUT_WIDTH 1920
#define MUXER_OUTPUT_HEIGHT 1080

/* Muxer batch formation timeout, for e.g. 40 millisec. Should ideally be set
 * based on the fastest source's framerate. */
#define MUXER_BATCH_TIMEOUT_USEC 40000

gchar pgie_classes_str[4][32] = {"Vehicle", "TwoWheeler", "Person", "Roadsign"};

static GstPadProbeReturn
before_probe(GstPad *pad, GstPadProbeInfo *info,
                          gpointer u_data)
{
  GstBuffer *buf = (GstBuffer *)info->data;
  NvDsObjectMeta *obj_meta = NULL;
  NvDsMetaList *l_frame = NULL;
  NvDsMetaList *l_obj = NULL;

  NvDsBatchMeta *batch_meta = gst_buffer_get_nvds_batch_meta(buf);

  static guint batch_number = 0;
  g_message("----------");

  for (l_frame = batch_meta->frame_meta_list; l_frame != NULL; l_frame = l_frame->next)
  {
    NvDsFrameMeta *frame_meta = (NvDsFrameMeta *)(l_frame->data);

    int object_id = 0;
    for (l_obj = frame_meta->obj_meta_list; l_obj != NULL; l_obj = l_obj->next)
    {
      obj_meta = (NvDsObjectMeta *)(l_obj->data);

      g_message("before nvvideoconvert [batch = %d, batch_id = %d, src_id = %d, object_index = %d], x = %.2f, y = %.2f",
                batch_number, frame_meta->batch_id, frame_meta->source_id, object_id,
                obj_meta->detector_bbox_info.org_bbox_coords.top, obj_meta->detector_bbox_info.org_bbox_coords.left);

      object_id++;
    }
  }

  batch_number++;
  g_message("----------");

  return GST_PAD_PROBE_OK;
}

static GstPadProbeReturn
after_probe(GstPad *pad, GstPadProbeInfo *info,
                          gpointer u_data)
{
  GstBuffer *buf = (GstBuffer *)info->data;
  NvDsObjectMeta *obj_meta = NULL;
  NvDsMetaList *l_frame = NULL;
  NvDsMetaList *l_obj = NULL;

  NvDsBatchMeta *batch_meta = gst_buffer_get_nvds_batch_meta(buf);

  static guint batch_number = 0;
  g_message("----------");

  int object_id = 0;
  for (l_frame = batch_meta->frame_meta_list; l_frame != NULL; l_frame = l_frame->next)
  {
    NvDsFrameMeta *frame_meta = (NvDsFrameMeta *)(l_frame->data);
    for (l_obj = frame_meta->obj_meta_list; l_obj != NULL; l_obj = l_obj->next)
    {
      obj_meta = (NvDsObjectMeta *)(l_obj->data);

      g_message("after nvvideoconvert [batch = %d, batch_id = %d, src_id = %d, object_index = %d], x = %.2f, y = %.2f",
                batch_number, frame_meta->batch_id, frame_meta->source_id, object_id,
                obj_meta->detector_bbox_info.org_bbox_coords.top, obj_meta->detector_bbox_info.org_bbox_coords.left);

      object_id++;
    }
  }

  batch_number++;
  g_message("----------");

  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;
}

static GstElement *create_source_bin(guint index, const char* location)
{
  GstElement *bin = NULL,
             *source = NULL,
             *h264parser = NULL,
             *decoder = NULL;
  GstPad *real_pad, *ghost_pad;

  /* 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);

  /* Create elements in the bin */
  source = gst_element_factory_make ("filesrc", "file-source");
  h264parser = gst_element_factory_make ("h264parse", "h264-parser");
  decoder = gst_element_factory_make ("nvv4l2decoder", "nvv4l2-decoder");

  /* Set elements properties */
  g_object_set (G_OBJECT (source), "location", location, NULL);

  /* Add elements to the bin */
  gst_bin_add_many(GST_BIN(bin), source, h264parser, decoder, NULL);

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

  real_pad = gst_element_get_static_pad(decoder, "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;
}

int main(int argc, char *argv[])
{
  GMainLoop *loop = NULL;
  GstElement *pipeline = NULL,
             *streammux = NULL,
             *pgie = NULL,
             *tiler = NULL,
             *nvvidconv = NULL,
             *nvosd = NULL,
             *sink = NULL;

  GstElement *transform = NULL;
  GstBus *bus = NULL;
  guint bus_watch_id;

  if (argc < 2)
  {
    return -1;
  }

  int num_sources = argc - 1;

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

  /* Create gstreamer elements */
  pipeline = gst_pipeline_new("dstest1-pipeline");
  streammux = gst_element_factory_make("nvstreammux", "stream_muxer");
  pgie = gst_element_factory_make("nvinfer", "primary-nvinference-engine");
  tiler = gst_element_factory_make("nvmultistreamtiler", "nvtiler");
  nvvidconv = gst_element_factory_make("nvvideoconvert", "nvvideo_converter");
  nvosd = gst_element_factory_make("nvdsosd", "nv_onscreendisplay");
  sink = gst_element_factory_make("nveglglessink", "nvvideo-renderer");

  if (!pipeline || !streammux || !pgie || !tiler || !nvvidconv || !nvosd || !sink)
  {
    g_printerr("One element could not be created. Exiting.\n");
    return -1;
  }

  g_object_set(G_OBJECT(pgie), "config-file-path", "dstest1_pgie_config.txt", NULL);

  g_object_set(G_OBJECT(streammux),
               "batch-size", num_sources,
               "width", MUXER_OUTPUT_WIDTH,
               "height", MUXER_OUTPUT_HEIGHT,
               "batched-push-timeout", MUXER_BATCH_TIMEOUT_USEC, NULL);

  /* we add a 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);

  /* Set up the pipeline */
  gst_bin_add_many(GST_BIN(pipeline), streammux, pgie, tiler, nvvidconv, nvosd, sink, NULL);

  for (guint i = 0; i < num_sources; i++)
  {
    GstElement *source_bin = NULL;
    source_bin = create_source_bin(i, argv[i + 1]);
    gst_bin_add(GST_BIN(pipeline), 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);
    srcpad = gst_element_get_static_pad(source_bin, "src");
    gst_pad_link(srcpad, sinkpad);
    gst_object_unref(srcpad);
    gst_object_unref(sinkpad);
  }

  /* we link the elements together */
  if (!gst_element_link_many(streammux, pgie, tiler, nvvidconv, nvosd, sink, NULL))
  {
    g_printerr("Elements could not be linked: 2. Exiting.\n");
    return -1;
  }

  /* Lets add probes */
  GstPad *before_pad = gst_element_get_static_pad(nvvidconv, "sink");
  if (!before_pad)
  {
    g_print("Unable to get pad\n");
  }
  else
  {
    gst_pad_add_probe(before_pad, GST_PAD_PROBE_TYPE_BUFFER, before_probe, NULL, NULL);
  }
  gst_object_unref(before_pad);

  GstPad *after_pad = gst_element_get_static_pad(nvvidconv, "src");
  if (!after_pad)
  {
    g_print("Unable to get pad\n");
  }
  else
  {
    gst_pad_add_probe(after_pad, GST_PAD_PROBE_TYPE_BUFFER, after_probe, NULL, NULL);
  }
  gst_object_unref(after_pad);

  /* Set the pipeline to "playing" state */
  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);
  return 0;
}