#include <fstream>
#include <opencv2/opencv.hpp>

// libcommon
#include "Logger.h"
#include "OSTime.h"
#include "SysUtils.h"

// libmascommon
#include "MemPool.h"

// deps
#include "DetUtils.h"
#include "ilogger.hpp"
#include "trt_infer.hpp"

// libyolocrowddetector
#include "DetectorAPI.h"
#include "YoloCrowdDetector.h"

TZ_INT DetectionCallback(SPtr<masd::StreamInfo>& media, void* ctx)
{
    TZLogInfo("Detection callback triggered!~~~");
	
    if (ctx == nullptr)
    {
        TZLogError("Error: Invalid image pointer!!!");
        return -1;
    }

    cv::Mat* img = reinterpret_cast<cv::Mat*>(ctx);
    if (img->empty())
    {
        TZLogError("Error: Invalid image pointer!!!");
        return -1;
    }

    auto allDetRst = media->GetAllDetRst();
    int imgWidth = img->cols, imgHeight = img->rows;

    for (auto it = allDetRst.begin(); it != allDetRst.end(); ++it)
    {
        const std::string& detKey = it->first;
        const SPtr<masd::DetProducing>& detProducing = it->second;

        TZLogInfo("Detection Key: %s~~~", detKey.c_str());
        std::cout << "Detection Result: " << detProducing->Result << std::endl;

        if (!detProducing->Draw.Rects.empty())
        {
            TZLogInfo("Processing Draw Info...~~~");
            for (const auto& rect : detProducing->Draw.Rects)
            {
                TZLogInfo("Rect: LTX: %.2f, LTY: %.2f, "
                          "RBX: %.2f, RBY: %.2f, Color: %s, Thickness: %d~~~",
                          rect.LTX, rect.LTY, rect.RBX, rect.RBY,
                          rect.Color.c_str(), rect.Thickness);

                if (!rect.Text.Text.empty())
                {
                    TZLogInfo("Text: %s~~~", rect.Text.Text.c_str());
                }

                cv::Scalar color;
                {
                    std::stringstream colorStream(rect.Color);
                    int r, g, b;
                    char comma;
                    colorStream >> r >> comma >> g >> comma >> b;
                    color = cv::Scalar(b, g, r);
                }

                cv::Point topLeft(rect.LTX * imgWidth, rect.LTY * imgHeight);
                cv::Point bottomRight(rect.RBX * imgWidth, rect.RBY * imgHeight);

                cv::rectangle(*img, topLeft, bottomRight, color, rect.Thickness);
            }
        }
        else
        {
            TZLogInfo("No Draw Info available.~~~");
        }

        if (detProducing->DetMedia)
        {
            const auto& media = detProducing->DetMedia;
            TZLogInfo("Media Length: %d~~~", media->Length);
            TZLogInfo("Media DataType: %d~~~", media->DataType);
            TZLogInfo("Media Height: %d, Width: %d~~~", media->Height, media->Width);
        }
    }

    if (cv::imwrite("DetRst.jpg", *img))
    {
        TZLogInfo("Image saved as DetRst.jpg~~~");
    }
    else
    {
        TZLogError("Failed to save the image!!!");
    }

    return 0;
}

int main()
{
    // Initialize log
    INITIALIZE_LOGGER_NORMAL("test", "./test.log", 1, 100, 6, 1, 1);

    // Initialize memory pool
    masd::MemPool *pool = masd::MEMPOOL;
    if (pool->Initialize() != masd::MEC_OK)
    {
        TZLogError("Memory pool initialization failed!!!");
        return -1;
    }

    /* Calling libyolocrowddetector and libheadcountstrategy */
    // Step 1:
    // Initialize the SDK
    TZ_INT initResult = Initialize();
    if (initResult != masd::MEC_OK)
    {
        TZLogError("Failed to initialize the SDK!!!");
        return -1;
    }
    TZLogInfo("SDK Initialized Successfully~~~");

    // Step 2:
    // Build yolo-crowd detector
    masd::Detector* detector = BuildDetector();
    if (detector == nullptr)
    {
        TZLogError("Failed to build yolo-crowd detector!!!");
        Dispose();
        return -1;
    }
    TZLogInfo("Yolo-crowd detector built successfully~~~");

    // Step 3:
    // Initialize the yolo-crowd detector with configuration parameters
    std::string initParam = 
    "{\"gpu_id\": 0, \"max_objects\": 1024, "
    "\"confidence_threshold\": 0.2, \"nms_threshold\": 0.5, "
    "\"model_path\": \"../../models/yolo-crowd-output0.trt\"}";
    TZ_INT initDetResult = detector->Initialize(initParam);
    if (initDetResult != masd::MEC_OK)
    {
        TZLogError("Failed to initialize the yolo-crowd detector!!!");
        DestroyDetector(detector);
        Dispose();
        return -1;
    }
    TZLogInfo("Yolo-crowd detector initialized successfully~~~");

    // Step 4:
    // Simulate frame processing with using test image
    cv::Mat testImage = cv::imread("../../media/image.jpg");
    if (testImage.empty())
    {
        TZLogError("Failed to load test image!!!");
        DestroyDetector(detector);
        Dispose();
        return -1;
    }

    TZ_INT length = testImage.total() * testImage.elemSize();
    SPtr<masd::Media> mediaResource = std::make_shared<masd::Media>(length);
    mediaResource->Width = testImage.cols;
    mediaResource->Height = testImage.rows;
    mediaResource->DataType = testImage.type();
    mediaResource->Mem = testImage.data;

    SPtr<masd::StreamInfo> streamInfo = std::make_shared<masd::StreamInfo>();
    streamInfo->SetMediaRsc(mediaResource);

    detector->DoDetect(streamInfo, DetectionCallback, &testImage);

    // Step 5:
    // Destroy the detector
    DestroyDetector(detector);
    TZLogInfo("Detector destroyed successfully~~~");

    // Step 6:
    // Dispose the SDK
    TZ_INT disposeResult = Dispose();
    if (disposeResult != masd::MEC_OK)
    {
        TZLogError("Failed to dispose the SDK!!!");
        return -1;
	}
	TZLogInfo("SDK disposed successfully~~~");
	
	return 0;
}