#include <gst/gst.h>
#include <gstnvdsmeta.h>
#include <nvdsmeta.h>
#include "nvdspreprocess_meta.h"
#include <atomic>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <vector>
#include <cmath>
#include "/source/native/nvinferserver_rfdetr/frame_completion.hpp"
static std::atomic<int> inferred{0},completed{0},objects{0},invalid{0};
static std::vector<std::string> samples;
static GstPadProbeReturn consumer(GstPad*,GstPadProbeInfo* info,gpointer){
 ++inferred;
 auto* batch=gst_buffer_get_nvds_batch_meta(GST_PAD_PROBE_INFO_BUFFER(info));
 if(batch)for(auto* f=batch->frame_meta_list;f;f=f->next){
  auto* frame=static_cast<NvDsFrameMeta*>(f->data);
  for(auto* m=frame->frame_user_meta_list;m;m=m->next){
   auto* u=static_cast<NvDsUserMeta*>(m->data);
   if(u->base_meta.meta_type==contextra::completionMetaType()){
    ++completed;if(samples.size()<2)samples.emplace_back(static_cast<const char*>(u->user_meta_data));
   }
  }
  for(auto* o=frame->obj_meta_list;o;o=o->next){
   ++objects;auto* ob=static_cast<NvDsObjectMeta*>(o->data);const auto& r=ob->rect_params;
   if(!std::isfinite(ob->confidence)||ob->confidence<0||ob->confidence>1||ob->class_id<0||ob->class_id>=80||
      !std::isfinite(r.left)||!std::isfinite(r.top)||!std::isfinite(r.width)||!std::isfinite(r.height)||
      r.left<0||r.top<0||r.width<=0||r.height<=0||r.left+r.width>1920.1||r.top+r.height>1080.1)++invalid;
  }
 }
 return GST_PAD_PROBE_OK;
}
static std::atomic<int> before{0},after{0},valid{0};
static guint64 target;
static GstPadProbeReturn pre(GstPad*,GstPadProbeInfo*,gpointer){++before;return GST_PAD_PROBE_OK;}
static GstPadProbeReturn post(GstPad*,GstPadProbeInfo* info,gpointer){
 ++after;
 auto* batch=gst_buffer_get_nvds_batch_meta(GST_PAD_PROBE_INFO_BUFFER(info));
 if(batch) for(auto* node=batch->batch_user_meta_list;node;node=node->next){
  auto* user=static_cast<NvDsUserMeta*>(node->data);
  if(user->base_meta.meta_type!=NVDS_PREPROCESS_BATCH_META) continue;
  auto* meta=static_cast<GstNvDsPreProcessBatchMeta*>(user->user_meta_data);
  auto* tensor=meta->tensor_meta;
  if(tensor && tensor->tensor_shape==std::vector<int>({1,3,576,576}) &&
     tensor->tensor_name=="input" && tensor->data_type==NvDsDataType_FP32 &&
     tensor->buffer_size==3*576*576*sizeof(float) && tensor->raw_tensor_buffer &&
     meta->target_unique_ids==std::vector<guint64>({target})) ++valid;
 }
 return GST_PAD_PROBE_OK;
}
int main(int argc,char** argv){
 setvbuf(stdout,nullptr,_IONBF,0);if(argc!=4)return 2;
 target=std::strtoull(argv[2],nullptr,10);gst_init(nullptr,nullptr);
 std::string desc="videotestsrc num-buffers=12 pattern=black ! video/x-raw,format=NV12,width=320,height=180,framerate=10/1 ! nvvideoconvert gpu-id=0 ! video/x-raw(memory:NVMM),format=NV12 ! mux.sink_0 nvstreammux name=mux batch-size=1 width=1920 height=1080 batched-push-timeout=33000 buffer-pool-size=4 live-source=true compute-hw=1 ! nvdspreprocess name=prep config-file=";
 desc+=argv[1];desc+=" ! queue ! nvinferserver name=infer config-file-path=";desc+=argv[3];desc+=" batch-size=1 ! fakesink sync=false async=false";
 GError* err=nullptr;auto* pipeline=gst_parse_launch(desc.c_str(),&err);
 if(!pipeline||err){std::puts("{\"stage\":\"parse_failed\"}");return 3;}
 auto* mux=gst_bin_get_by_name(GST_BIN(pipeline),"mux");auto* prep=gst_bin_get_by_name(GST_BIN(pipeline),"prep");
 std::string uid=std::to_string(target);g_object_set(prep,"target-unique-ids",uid.c_str(),nullptr);
 auto* a=gst_element_get_static_pad(mux,"src");auto* b=gst_element_get_static_pad(prep,"src");
 gst_pad_add_probe(a,GST_PAD_PROBE_TYPE_BUFFER,pre,nullptr,nullptr);
 gst_pad_add_probe(b,GST_PAD_PROBE_TYPE_BUFFER,post,nullptr,nullptr);
 auto* infer=gst_bin_get_by_name(GST_BIN(pipeline),"infer");auto* ip=gst_element_get_static_pad(infer,"src");
 gst_pad_add_probe(ip,GST_PAD_PROBE_TYPE_BUFFER,consumer,nullptr,nullptr);
 auto* bus=gst_element_get_bus(pipeline);
 if(gst_element_set_state(pipeline,GST_STATE_PLAYING)==GST_STATE_CHANGE_FAILURE)return 4;
 auto* msg=gst_bus_timed_pop_filtered(bus,30*GST_SECOND,static_cast<GstMessageType>(GST_MESSAGE_EOS|GST_MESSAGE_ERROR));
 const bool eos=msg && GST_MESSAGE_TYPE(msg)==GST_MESSAGE_EOS;
 std::printf("{\"stage\":\"before_stop\",\"before\":%d,\"after\":%d,\"valid_meta\":%d,\"eos\":%s}\n",before.load(),after.load(),valid.load(),eos?"true":"false");
 std::printf("{\"inferred\":%d,\"completed\":%d,\"objects\":%d,\"invalid_objects\":%d,\"samples\":[",inferred.load(),completed.load(),objects.load(),invalid.load());
 for(size_t i=0;i<samples.size();++i)std::printf("%s%s",i?",":"",samples[i].c_str());std::puts("]}");
 if(msg)gst_message_unref(msg);
 gst_element_set_state(pipeline,GST_STATE_NULL);
 gst_object_unref(ip);gst_object_unref(infer);gst_object_unref(a);gst_object_unref(b);gst_object_unref(mux);gst_object_unref(prep);gst_object_unref(bus);gst_object_unref(pipeline);
 const bool ok=eos && before==12 && after==12 && valid==12 && inferred==12 && completed==12 && invalid==0;
 std::printf("{\"stage\":\"stopped\",\"pass\":%s}\n",ok?"true":"false");return ok?0:5;
}