#include "nv_osd_manager.hpp"

#define SET_GPU_ID(object, gpu_id) \
    g_object_set(G_OBJECT(object), "gpu-id", gpu_id, NULL);
#define GPU_ID 0
#define OSD_PROCESS_MODE \
    1  // use GPU to draw rectangles, keypoints and text on frame if
       // OSD_PROCESS_MODE set to 1
#define OSD_DISPLAY_TEXT 1
#define MAX_DISPLAY_LEN 64
#define PGIE_CLASS_ID_PERSON 0

gint NvOsdManager::frame_number = 0;
bool NvOsdManager::write_full_frame_to_disk = false;
bool NvOsdManager::write_cropped_objects_to_disk = false;

NvOsdManager::NvOsdManager() {
    const auto &config = ConfigManager::get_instance().get_config();
    write_full_frame_to_disk =
        config.at("write_full_frame_to_disk").get<bool>();
    write_cropped_objects_to_disk =
        config.at("write_cropped_objects_to_disk").get<bool>();
}

bool NvOsdManager::create_nv_osd() {
    /* Create OSD to draw on the converted RGBA buffer */
    nvosd = gst_element_factory_make("nvdsosd", "nv-onscreendisplay");
    /* Finally render the osd output */
    g_object_set(G_OBJECT(nvosd), "process-mode", OSD_PROCESS_MODE,
                 "display-text", OSD_DISPLAY_TEXT, NULL);
    SET_GPU_ID(nvosd, GPU_ID);

    if (!nvosd) {
        g_printerr("Unable to create NVOSD. Exiting.\n");
        return false;
    }
    return true;
}

// Attach probe to a pad in the pipeline
void NvOsdManager::attach_probe_to_src_nvosd(
    NvDsObjEncCtxHandle obj_ctx_handle) {
    GstPad *src_pad = gst_element_get_static_pad(nvosd, "src");
    if (!src_pad) {
        std::cerr << "Unable to get nvosd src pad\n";
        return;
    }

    gst_pad_add_probe(src_pad, GST_PAD_PROBE_TYPE_BUFFER,
                      osd_src_pad_buffer_probe, (gpointer)obj_ctx_handle, NULL);
    gst_object_unref(src_pad);
}

void NvOsdManager::save_full_frame(NvDsFrameMeta *frame_meta) {
    char fileFrameNameString[FILE_NAME_SIZE];
    const char *osd_string = "OSD";

    /* For Demonstration Purposes we are writing metadata to jpeg images of
     * the first 10 frames only.
     * The files generated have an 'OSD' prefix. */
    NvDsUserMetaList *usrMetaList = frame_meta->frame_user_meta_list;
    FILE *file;
    int stream_num = 0;
    while (usrMetaList != NULL) {
        NvDsUserMeta *usrMetaData = (NvDsUserMeta *)usrMetaList->data;
        // std::cout<<"usrMetaData->base_meta.meta_type frame =
        // "<<usrMetaData->base_meta.meta_type<<std::endl;
        if (usrMetaData->base_meta.meta_type == NVDS_CROP_IMAGE_META) {
            snprintf(fileFrameNameString, FILE_NAME_SIZE, "%s_frame_%d_%d.jpg",
                     osd_string, frame_number, stream_num++);
            NvDsObjEncOutParams *enc_jpeg_image =
                (NvDsObjEncOutParams *)usrMetaData->user_meta_data;
            /* Write to File */
            file = fopen(fileFrameNameString, "wb");
            fwrite(enc_jpeg_image->outBuffer, sizeof(uint8_t),
                   enc_jpeg_image->outLen, file);
            fclose(file);
            // std::cout<<"fileFrameNameString =
            // "<<fileFrameNameString<<std::endl;
        }
        usrMetaList = usrMetaList->next;
    }
}

void NvOsdManager::save_cropped_objects(NvDsFrameMeta *frame_meta,
                                        NvDsObjectMeta *obj_meta,
                                        guint num_rects) {
    const char *osd_string = "OSD";
    // std::cout<<"obj_meta->unique_component_id =
    // "<<obj_meta->unique_component_id<<std::endl;
    char fileObjNameString[FILE_NAME_SIZE];

    /* For Demonstration Purposes we are writing metadata to jpeg images of
     * vehicles or persons for the first 100 frames only.
     * The files generated have a 'OSD' prefix. */
    //   std::cout<<"in obj_meta_list NvOsdManager"<<std::endl;
    //   std::quick_exit(0);
    NvDsUserMetaList *usrMetaList = obj_meta->obj_user_meta_list;
    FILE *file;
    while (usrMetaList != NULL) {
        NvDsUserMeta *usrMetaData = (NvDsUserMeta *)usrMetaList->data;
        // std::cout<<"usrMetaData->base_meta.meta_type object =
        // "<<usrMetaData->base_meta.meta_type<<std::endl;
        if (usrMetaData->base_meta.meta_type == NVDS_CROP_IMAGE_META) {
            NvDsObjEncOutParams *enc_jpeg_image =
                (NvDsObjEncOutParams *)usrMetaData->user_meta_data;

            snprintf(fileObjNameString, FILE_NAME_SIZE, "%s_%d_%d_%d_%s.jpg",
                     osd_string, frame_number, frame_meta->batch_id, num_rects,
                     obj_meta->obj_label);
            /* Write to File */
            file = fopen(fileObjNameString, "wb");
            fwrite(enc_jpeg_image->outBuffer, sizeof(uint8_t),
                   enc_jpeg_image->outLen, file);
            fclose(file);
            std::cout << "fileObjNameString = " << fileObjNameString
                      << " obj_meta->class_id =" << obj_meta->class_id
                      << std::endl;
            usrMetaList = NULL;
        } else {
            usrMetaList = usrMetaList->next;
        }
    }
}

/* This is the buffer probe function that we have registered on the sink pad
 * of the OSD element. All the infer elements in the pipeline shall attach
 * their metadata to the GstBuffer, here we will iterate & process the metadata
 * forex: class ids to strings, counting of class_id objects etc. */
GstPadProbeReturn NvOsdManager::osd_src_pad_buffer_probe(GstPad *pad,
                                                         GstPadProbeInfo *info,
                                                         gpointer u_data) {
    (void)pad;
    (void)u_data;
    GstBuffer *buf = (GstBuffer *)info->data;
    guint num_rects = 0;
    guint person_count = 0;
    NvDsObjectMeta *obj_meta = NULL;
    NvDsMetaList *l_frame = NULL;
    NvDsMetaList *l_obj = NULL;
    NvDsDisplayMeta *display_meta = NULL;

    NvDsBatchMeta *batch_meta = gst_buffer_get_nvds_batch_meta(buf);

    for (l_frame = batch_meta->frame_meta_list; l_frame != NULL;
         l_frame = l_frame->next) {
        NvDsFrameMeta *frame_meta = (NvDsFrameMeta *)(l_frame->data);
        int offset = 0;
        if (write_full_frame_to_disk == true) save_full_frame(frame_meta);
        for (l_obj = frame_meta->obj_meta_list; l_obj != NULL;
             l_obj = l_obj->next) {
            obj_meta = (NvDsObjectMeta *)(l_obj->data);
            if (obj_meta->class_id == PGIE_CLASS_ID_PERSON) {
                person_count++;
                num_rects++;
                // std::cout << "In OSD sink "
                //           << "x = " << obj_meta->rect_params.left
                //           << " y = " << obj_meta->rect_params.top
                //           << " w = " << obj_meta->rect_params.width
                //           << " h = " << obj_meta->rect_params.height
                //           << " score = " << obj_meta->confidence
                //           << "  Object ID: " << obj_meta->object_id
                //           << std::endl;
            }
            if (write_cropped_objects_to_disk == true)
                save_cropped_objects(frame_meta, obj_meta, num_rects);
        }
        display_meta = nvds_acquire_display_meta_from_pool(batch_meta);
        NvOSD_TextParams *txt_params = &display_meta->text_params[0];
        display_meta->num_labels = 1;
        txt_params->display_text = (gchar *)g_malloc0(MAX_DISPLAY_LEN);
        offset = snprintf(txt_params->display_text, MAX_DISPLAY_LEN,
                          "Person = %d ", person_count);
        (void)offset;

        /* Now set the offsets where the string should appear */
        txt_params->x_offset = 10;
        txt_params->y_offset = 12;

        /* Font , font-color and font-size */
        txt_params->font_params.font_name = (gchar *)"Serif";
        txt_params->font_params.font_size = 10;
        txt_params->font_params.font_color.red = 1.0;
        txt_params->font_params.font_color.green = 1.0;
        txt_params->font_params.font_color.blue = 1.0;
        txt_params->font_params.font_color.alpha = 1.0;

        /* Text background color */
        txt_params->set_bg_clr = 1;
        txt_params->text_bg_clr.red = 0.0;
        txt_params->text_bg_clr.green = 0.0;
        txt_params->text_bg_clr.blue = 0.0;
        txt_params->text_bg_clr.alpha = 1.0;

        nvds_add_display_meta_to_frame(frame_meta, display_meta);
    }
    // g_print(
    //     "In OSD sink "
    //     "Frame Number = %d "
    //     "Person Count = %d\n",
    //     frame_number, person_count);

    frame_number++;
    return GST_PAD_PROBE_OK;
}

// Attach probe to a pad in the pipeline
void NvOsdManager::attach_probe_to_sink_nvosd() {
    GstPad *sink_pad = gst_element_get_static_pad(nvosd, "sink");
    if (!sink_pad) {
        std::cerr << "Unable to get nvosd src pad\n";
        return;
    }

    gst_pad_add_probe(sink_pad, GST_PAD_PROBE_TYPE_BUFFER,
                      osd_sink_pad_buffer_probe, NULL, NULL);
    gst_object_unref(sink_pad);
}

GstPadProbeReturn NvOsdManager::osd_sink_pad_buffer_probe(GstPad *,
                                                          GstPadProbeInfo *,
                                                          gpointer) {
    return GST_PAD_PROBE_OK;
}