/*
 ********************************************************************************
 * NOTICE
 * This software is the property of HiRain Technologies. Any information
 * contained in this doc should not be reproduced, or used, or disclosed
 * without the written authorization from HiRain Technologies.
 *
 ********************************************************************************
 *   Description:
 *   - Pull NV12 frames from camera pipeline
 *   - Run YOLOX model (416x416) inference
 *   - Draw boxes on the full frame (w x h)
 *   - Push annotated video in real-time via UDP using GStreamer appsrc -> H264 RTP
 ********************************************************************************
 */

#include <chrono>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <vector>

#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <cstring>
#include <fstream>
#include <mutex>
#include <filesystem>

#include <opencv2/opencv.hpp>

// GStreamer
#include <gst/gst.h>
#include <gst/app/gstappsrc.h>

#include "camera_handler_type.h"
#include "camera_handler_interface.h"
#include "inference_interface.h"
#include "utils.h"

using namespace inference;
using namespace camera_handler;

static bool stop_flag = false;
static uint8_t* image_data = nullptr;

static CameraHandlerStatus camera_status;
static InferenceStatus infer_status;
static AppImageInfo image_info;

static int w = 640;
static int h = 640;

// Model directory
static std::string directory = "/opt/model_zoo/ONR-OD-8200-yolox-nano-lite-mmdet-coco-416x416/";

// UDP target (receiver)
static std::string ip   = "192.168.178.77";
static std::string port = "5001";

// Interrupt function, used to respond to Ctrl+C input
void handle_sigint(int) {
  std::cout << "User inputs ctrl+c" << std::endl;
  stop_flag = true;
}

int main(int argc, char** argv) {
  // Register Interrupt Function
  signal(SIGINT, handle_sigint);

  // Initialize GStreamer (required)
  gst_init(&argc, &argv);

  // ---------------- CAMERA PIPELINE (PULL) ----------------
  auto pipe_ptr = CreatePipeline(CAMERA_HANDLER_NV12, 640, 640);
  camera_status = GetImageInfo(pipe_ptr, image_info);
  if (camera_status == CAMERA_HANDLER_NULLPTR) {
    std::cout << "CreatePipeline/GetImageInfo returned nullptr" << std::endl;
    return 0;
  }

  w = image_info.width_;
  h = image_info.height_;

  std::cout << "Camera output: " << w << "x" << h << " (NV12)" << std::endl;

  // ---------------- INFERER ----------------
  auto* inferer = CreateDdawInferer(directory);
  if (!inferer) {
    std::cout << "CreateDdawInferer returned nullptr!!!" << std::endl;
    return 0;
  }

  infer_status = DumpInputInfo(inferer);
  infer_status = DumpOutputInfo(inferer);
  if (infer_status != INFERENCE_SUCCESS) {
    std::cout << "Dump info failed!!!" << std::endl;
  }

  // ---------------- PUSH PIPELINE (UDP) ----------------
  // We will push annotated frames as BGR (OpenCV default) and let videoconvert convert to NV12 for encoder.
  std::string push_pipe =
      "appsrc name=appsrc is-live=true block=true do-timestamp=true format=GST_FORMAT_TIME ! "
      "video/x-raw,format=BGR,width=" + std::to_string(w) + ",height=" + std::to_string(h) + ",framerate=30/1 ! "
      "queue leaky=2 max-size-buffers=1 ! videoconvert ! video/x-raw,format=NV12 ! "
      "queue leaky=2 max-size-buffers=1 ! v4l2h264enc ! rtph264pay pt=96 config-interval=1 ! "
      "udpsink host=" + ip + " port=" + port + " sync=false";

  GstElement* push_pipeline = gst_parse_launch(push_pipe.c_str(), NULL);
  if (!push_pipeline) {
    std::cout << "Failed to create push pipeline." << std::endl;
    DestroyDdawInferer(inferer);
    DestroyPipeline(pipe_ptr);
    return 0;
  }

  GstElement* appsrc = gst_bin_get_by_name(GST_BIN(push_pipeline), "appsrc");
  if (!appsrc) {
    std::cout << "Failed to get appsrc from push pipeline." << std::endl;
    gst_object_unref(push_pipeline);
    DestroyDdawInferer(inferer);
    DestroyPipeline(pipe_ptr);
    return 0;
  }

  // Optional: set appsrc properties (safer for live streaming)
  g_object_set(G_OBJECT(appsrc),
               "stream-type", 0, // GST_APP_STREAM_TYPE_STREAM
               "format", GST_FORMAT_TIME,
               "is-live", TRUE,
               "block", TRUE,
               NULL);

  gst_element_set_state(push_pipeline, GST_STATE_PLAYING);

  const int bgr_size = w * h * 3;

  std::cout << "Streaming to udp://" << ip << ":" << port << std::endl;
  std::cout << "Receiver (PC) example:\n"
            << "gst-launch-1.0 -v udpsrc port=" << port
            << " caps=\"application/x-rtp,media=video,encoding-name=H264,payload=96\" "
            << "! rtph264depay ! h264parse ! avdec_h264 ! videoconvert ! autovideosink sync=false\n";

  // ---------------- MAIN LOOP ----------------
  while (!stop_flag) {
    // Pull image (NV12)
    camera_status = PullImageFromPipeline(pipe_ptr, image_data);
    if (camera_status != CAMERA_HANDLER_SUCCESS || !image_data) {
      std::cout << "cannot get image from pipeline" << std::endl;
      break;
    }

    // NV12 raw buffer -> OpenCV Mat
    cv::Mat raw_nv12 = cv::Mat(int(1.5 * h), w, CV_8UC1, image_data);

    // Convert to BGR (full frame)
    cv::Mat bgr_full;
    cv::cvtColor(raw_nv12, bgr_full, cv::COLOR_YUV2BGR_NV12);

    // Preprocess for model (416x416)
    cv::Mat resize_image;
    cv::resize(bgr_full, resize_image, cv::Size(416, 416));

    auto* inputbuff = GetInputBufferData(inferer, 0);
    Preprocess(resize_image, inputbuff);

    // Inference
    infer_status = RunModel(inferer);
    if (infer_status != INFERENCE_SUCCESS) {
      std::cout << "infer failed!!!" << std::endl;
      // continue; // you can continue or break; keeping continue avoids stopping stream
    }

    // Postprocess (boxes in 416x416 coordinates)
    auto* detdata = GetOutputBufferData(inferer, 0);
    std::vector<std::array<int, 4>> infer_result;
    Postprocess(infer_result, detdata);

    // Draw on full frame (scale 416 -> w/h)
    if (!infer_result.empty()) {
      float sx = static_cast<float>(w) / 416.0f;
      float sy = static_cast<float>(h) / 416.0f;

      for (auto& box : infer_result) {
        int x1 = static_cast<int>(box[0] * sx);
        int y1 = static_cast<int>(box[1] * sy);
        int x2 = static_cast<int>(box[2] * sx);
        int y2 = static_cast<int>(box[3] * sy);

        x1 = std::max(0, std::min(x1, w - 1));
        y1 = std::max(0, std::min(y1, h - 1));
        x2 = std::max(0, std::min(x2, w - 1));
        y2 = std::max(0, std::min(y2, h - 1));

        cv::rectangle(bgr_full,
                      cv::Point(x1, y1),
                      cv::Point(x2, y2),
                      cv::Scalar(0, 0, 255),
                      2);
      }
    }

    // Push annotated frame via appsrc (BGR)
    GstBuffer* buffer = gst_buffer_new_allocate(NULL, bgr_size, NULL);
    if (!buffer) {
      std::cout << "Failed to allocate GstBuffer" << std::endl;
      break;
    }

    GstMapInfo map;
    if (!gst_buffer_map(buffer, &map, GST_MAP_WRITE)) {
      std::cout << "Failed to map GstBuffer" << std::endl;
      gst_buffer_unref(buffer);
      break;
    }

    memcpy(map.data, bgr_full.data, bgr_size);
    gst_buffer_unmap(buffer, &map);

    // (Optional) set PTS/DTS for smoother playback (not strictly required with do-timestamp=true)
    // GST_BUFFER_PTS(buffer) = gst_util_uint64_scale(frame_count, GST_SECOND, 30);
    // GST_BUFFER_DURATION(buffer) = gst_util_uint64_scale(1, GST_SECOND, 30);

    GstFlowReturn ret = gst_app_src_push_buffer(GST_APP_SRC(appsrc), buffer);
    if (ret != GST_FLOW_OK) {
      std::cout << "gst_app_src_push_buffer failed: " << ret << std::endl;
      break;
    }

    // Optional debug output (careful: writing every frame is slow)
    // cv::imwrite("/home/ddaw_image.jpg", bgr_full);
  }

  // Send EOS (optional, helps receiver end cleanly)
  gst_app_src_end_of_stream(GST_APP_SRC(appsrc));

  // ---------------- CLEANUP ----------------
  infer_status = DestroyDdawInferer(inferer);

  gst_element_set_state(push_pipeline, GST_STATE_NULL);
  gst_object_unref(appsrc);
  gst_object_unref(push_pipeline);

  camera_status = DestroyPipeline(pipe_ptr);

  return 0;
}