/*
 * Contextra RF-DETR Custom Processor for gst-nvinferserver (DeepStream 9.1).
 *
 * CTXRF-01: C++은 RF-DETR tensor를 NvDsObjectMeta로 바꾸는 경계만 소유한다.
 * Tracking, Scenario, Event와 OSD 표현 정책은 Python Metadata Plane이 소유한다.
 * MODELCONTRACT-01: 고정 COCO slot mapping과 query/class 선택 뒤 Metadata를 생성한다.
 */

#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <string>
#include <unordered_map>
#include <vector>

#include "infer_custom_process.h"
#include "nvbufsurface.h"
#include "nvdsmeta.h"

#include "tensor_contract.hpp"
#include "detection_decode.hpp"
#include "frame_dimensions.hpp"
#include "frame_completion.hpp"
#include "build_identity.hpp"
#include "pose_decode.hpp"
#include "pose_metadata.hpp"
#include "crowd_decode.hpp"

namespace dsis = nvdsinferserver;

namespace {

__attribute__((used)) const char kBuildIdentity[] =
    "contextra-rfdetr-adapter;source_sha256=" CONTEXTRA_ADAPTER_SOURCE_SHA256
    ";base_digest=" CONTEXTRA_ADAPTER_BASE_DIGEST
    ";deepstream=" CONTEXTRA_ADAPTER_DEEPSTREAM_VERSION
    ";schema=dets:FP32[B,Q,4],labels:FP32[B,Q,91]";

#define LOG_INFO(fmt, ...) std::fprintf(stdout, "[rfdetr-adapter] " fmt "\n", ##__VA_ARGS__)
#define LOG_ERROR(fmt, ...) std::fprintf(stderr, "[rfdetr-adapter] " fmt "\n", ##__VA_ARGS__)

std::vector<uint32_t> dimensions(const dsis::InferBufferDescription& description)
{
    std::vector<uint32_t> result;
    result.reserve(description.dims.numDims);
    for (uint32_t index = 0; index < description.dims.numDims; ++index) {
        result.push_back(description.dims.d[index]);
    }
    return result;
}

class RfDetrProcessor : public dsis::IInferCustomProcessor {
    std::atomic<unsigned long long> entered_{0}, completed_{0};
    bool pose_;
    bool fire_;
    bool crowd_;
public:
    explicit RfDetrProcessor(bool pose=false, bool fire=false, bool crowd=false) : pose_(pose), fire_(fire), crowd_(crowd) {}
    void supportInputMemType(dsis::InferMemType& type) override
    {
        type = dsis::InferMemType::kCpu;
    }

    bool requireInferLoop() const override { return false; }

    NvDsInferStatus extraInputProcess(
        const std::vector<dsis::IBatchBuffer*>& primaryInputs,
        std::vector<dsis::IBatchBuffer*>& extraInputs,
        const dsis::IOptions* options) override
    {
        (void)primaryInputs;
        (void)extraInputs;
        (void)options;
        return NVDSINFER_SUCCESS;
    }

    NvDsInferStatus inferenceDone(
        const dsis::IBatchArray* outputs, const dsis::IOptions* options) override
    {
        if (++entered_ == 1) LOG_INFO("stage=native_enter count=1");
        if (outputs == nullptr || options == nullptr) {
            LOG_ERROR("missing output or frame options");
            return NVDSINFER_INVALID_PARAMS;
        }

        std::vector<uint64_t> streamIds;
        if (options->getValueArray(OPTION_NVDS_SREAM_IDS, streamIds) != NVDSINFER_SUCCESS
            || streamIds.empty()) {
            LOG_ERROR("stream identity metadata is unavailable");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        std::unordered_map<std::string, const dsis::IBatchBuffer*> tensors;
        for (uint32_t index = 0; index < outputs->getSize(); ++index) {
            const dsis::IBatchBuffer* buffer = outputs->getBuffer(index);
            if (buffer != nullptr) {
                tensors.emplace(buffer->getBufDesc().name, buffer);
            }
        }

        const auto detsEntry = tensors.find(contextra::kDetsTensorName);
        const auto labelsEntry = tensors.find(contextra::kLabelsTensorName);
        if (detsEntry == tensors.end() || labelsEntry == tensors.end()) {
            LOG_ERROR("required output tensor is unavailable");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        const auto detsDescription = detsEntry->second->getBufDesc();
        const auto labelsDescription = labelsEntry->second->getBufDesc();
        if (detsDescription.dims.numDims != 3 || labelsDescription.dims.numDims != 3) {
            LOG_ERROR("output tensor rank does not match RF-DETR contract");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }
        const contextra::TensorContract contract{
            detsDescription.dataType == dsis::InferDataType::kFp32,
            labelsDescription.dataType == dsis::InferDataType::kFp32,
            dimensions(detsDescription),
            dimensions(labelsDescription),
            static_cast<uint32_t>(streamIds.size()),
        };
        const auto keypointsEntry = tensors.find("keypoints");
        const bool poseValid = pose_ && contract.detsFp32 && contract.labelsFp32 && contract.runtimeBatchSize==1;
        if ((!pose_ && !((fire_ || crowd_) ? contextra::validateFireTensorContract(contract) : contextra::validateTensorContract(contract))) ||
            (pose_ && (!poseValid || contract.detsDims != std::vector<uint32_t>{1,100,4}
                || contract.labelsDims != std::vector<uint32_t>{1,100,2}
                || keypointsEntry == tensors.end()
                || keypointsEntry->second->getBufDesc().dataType != dsis::InferDataType::kFp32
                || dimensions(keypointsEntry->second->getBufDesc()) != std::vector<uint32_t>{1,100,34,8}))) {
            LOG_ERROR("output tensor dtype or shape does not match RF-DETR contract");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        const float* dets = static_cast<const float*>(detsEntry->second->getBufPtr(0));
        const float* labels = static_cast<const float*>(labelsEntry->second->getBufPtr(0));
        if (dets == nullptr || labels == nullptr) {
            LOG_ERROR("output tensor buffer is unavailable");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        std::vector<NvDsFrameMeta*> frameMetas;
        if (!options->hasValue(OPTION_NVDS_FRAME_META_LIST)
            || options->getValueArray(OPTION_NVDS_FRAME_META_LIST, frameMetas)
                != NVDSINFER_SUCCESS) {
            LOG_ERROR("frame metadata is unavailable");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        NvDsBatchMeta* batchMeta = nullptr;
        if (!options->hasValue(OPTION_NVDS_BATCH_META)
            || options->getObj(OPTION_NVDS_BATCH_META, batchMeta) != NVDSINFER_SUCCESS
            || batchMeta == nullptr) {
            LOG_ERROR("batch metadata is unavailable");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        const uint32_t batchSize = contract.detsDims[0];
        if (frameMetas.size() < batchSize) {
            LOG_ERROR("frame metadata count does not match tensor batch");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        std::vector<NvBufSurfaceParams*> surfaceParams;
        if (options->hasValue(OPTION_NVDS_BUF_SURFACE_PARAMS_LIST)
            && options->getValueArray(OPTION_NVDS_BUF_SURFACE_PARAMS_LIST, surfaceParams)
                != NVDSINFER_SUCCESS) {
            LOG_ERROR("surface metadata cannot be read");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        int64_t uniqueId = 0;
        if (options->hasValue(OPTION_NVDS_UNIQUE_ID)
            && options->getInt(OPTION_NVDS_UNIQUE_ID, uniqueId) != NVDSINFER_SUCCESS) {
            LOG_ERROR("inference component identity cannot be read");
            return NVDSINFER_CUSTOM_LIB_FAILED;
        }

        const uint32_t queries = contract.detsDims[1];
        const uint32_t classCount = contract.labelsDims[2];
        for (uint32_t batchIndex = 0; batchIndex < batchSize; ++batchIndex) {
            NvDsFrameMeta* frameMeta = frameMetas[batchIndex];
            if (frameMeta == nullptr) {
                LOG_ERROR("frame metadata entry is null");
                return NVDSINFER_CUSTOM_LIB_FAILED;
            }

            uint32_t frameWidth = frameMeta->source_frame_width;
            uint32_t frameHeight = frameMeta->source_frame_height;
            if (batchIndex < surfaceParams.size() && surfaceParams[batchIndex] != nullptr) {
                frameWidth = surfaceParams[batchIndex]->width;
                frameHeight = surfaceParams[batchIndex]->height;
            }
            if (frameWidth == 0 || frameHeight == 0) {
                LOG_ERROR("frame dimensions are invalid");
                return NVDSINFER_CUSTOM_LIB_FAILED;
            }

            const size_t firstRow = static_cast<size_t>(batchIndex) * queries;
            std::vector<contextra::PoseDetection> poses;
            if (pose_) {
                try {
                    poses=contextra::decodePose(dets,labels,static_cast<const float*>(keypointsEntry->second->getBufPtr(0)));
                } catch (const std::exception&) {
                    LOG_ERROR("invalid Pose output");
                    return NVDSINFER_CUSTOM_LIB_FAILED;
                }
            }
            auto detections = (pose_ || crowd_) ? std::vector<contextra::Detection>{} : contextra::decodeDetections(
                dets + firstRow * contextra::kBoxWidth,
                labels + firstRow * classCount, queries, frameWidth, frameHeight,
                {}, fire_);
            if (crowd_) {
                const contextra::CrowdRawContract c{300,3,0,1,{2},
                    "16092352a61d4bd0fe77be94db28db5fc19e9201540d6f39f3513f037fb03db1",
                    "98f5c29c60c4bb60120e26e0287bea456c1bb503b4c2d21f4b0d20d8e47868dc",
                    "awiros-person-head-r01", "CXCYWH_NORMALIZED", "RAW_SIGMOID_LOGITS", "GLOBAL_TOP300_THEN_MAP"};
                try {
                    const auto boxes=contextra::decodeCrowdFrame(c,
                        std::vector<float>(dets+firstRow*4,dets+(firstRow+queries)*4),
                        std::vector<float>(labels+firstRow*classCount,labels+(firstRow+queries)*classCount));
                    for (const auto& b:boxes) {
                        contextra::Detection d{}; d.classId=b.classId; d.score=b.score;
                        d.left=b.xyxy[0]*frameWidth;d.top=b.xyxy[1]*frameHeight;
                        d.width=(b.xyxy[2]-b.xyxy[0])*frameWidth;d.height=(b.xyxy[3]-b.xyxy[1])*frameHeight;
                        detections.push_back(d);
                    }
                } catch (const std::exception&) {return NVDSINFER_CUSTOM_LIB_FAILED;}
            }
            if (pose_) for (const auto& p:poses) {
                contextra::Detection d{};
                d.classId=0; d.score=p.score;
                d.left=p.box[0]*frameWidth; d.top=p.box[1]*frameHeight;
                d.width=(p.box[2]-p.box[0])*frameWidth; d.height=(p.box[3]-p.box[1])*frameHeight;
                detections.push_back(d);
            }

            nvds_acquire_meta_lock(batchMeta);
            if (!contextra::publishFrameDimensions(*frameMeta, frameWidth, frameHeight)) {
                nvds_release_meta_lock(batchMeta);
                LOG_ERROR("bbox and pipeline dimensions conflict");
                return NVDSINFER_CUSTOM_LIB_FAILED;
            }
            for (const contextra::Detection& detection : detections) {
                NvDsObjectMeta* objectMeta = nvds_acquire_obj_meta_from_pool(batchMeta);
                if (objectMeta == nullptr) {
                    nvds_release_meta_lock(batchMeta);
                    LOG_ERROR("object metadata pool is exhausted");
                    return NVDSINFER_CUSTOM_LIB_FAILED;
                }
                objectMeta->unique_component_id = static_cast<gint>(uniqueId);
                objectMeta->confidence = detection.score;
                objectMeta->object_id = UNTRACKED_OBJECT_ID;
                objectMeta->class_id = detection.classId;
                objectMeta->rect_params.left = detection.left;
                objectMeta->rect_params.top = detection.top;
                objectMeta->rect_params.width = detection.width;
                objectMeta->rect_params.height = detection.height;
                nvds_add_obj_meta_to_frame(frameMeta, objectMeta, nullptr);
            }
            frameMeta->bInferDone = TRUE;
            if (pose_ && !contextra::publishPose(batchMeta,frameMeta,static_cast<guint>(uniqueId),poses)) {
                nvds_release_meta_lock(batchMeta);
                return NVDSINFER_CUSTOM_LIB_FAILED;
            }
            if (!contextra::publishCompletion(batchMeta, frameMeta,
                    static_cast<guint>(uniqueId), static_cast<guint>(detections.size()), frameWidth, frameHeight)) {
                nvds_release_meta_lock(batchMeta);
                LOG_ERROR("completion metadata pool is exhausted");
                return NVDSINFER_CUSTOM_LIB_FAILED;
            }
            nvds_release_meta_lock(batchMeta);
            if (++completed_ == 1) LOG_INFO("stage=native_complete count=1");
        }
        return NVDSINFER_SUCCESS;
    }

    void notifyError(NvDsInferStatus status) override
    {
        LOG_ERROR("nvinferserver reported status=%d", static_cast<int>(status));
    }
};

}  // namespace

extern "C" __attribute__((visibility("default"))) dsis::IInferCustomProcessor*
CreateContextraRfDetrProcessor(const char* config, uint32_t configLength)
{
    (void)config;
    (void)configLength;
    LOG_INFO("processor loaded");
    return new RfDetrProcessor();
}

extern "C" __attribute__((visibility("default"))) dsis::IInferCustomProcessor*
CreateContextraPoseProcessor(const char*, uint32_t) {
    LOG_INFO("Pose processor loaded");
    return new RfDetrProcessor(true);
}

extern "C" __attribute__((visibility("default"))) dsis::IInferCustomProcessor*
CreateContextraFireProcessor(const char*, uint32_t) {
    LOG_INFO("Fire processor loaded");
    return new RfDetrProcessor(false, true);
}

extern "C" __attribute__((visibility("default"))) dsis::IInferCustomProcessor*
CreateContextraCrowdProcessor(const char*, uint32_t) {
    return new RfDetrProcessor(false, false, true);
}