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

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

// libmascommon
#include "MemPool.h"

// deps
#include "DetUtils.h"

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

#define RUN_TEXT_DRAWING 0

struct CallbackContext
{
    cv::Mat* img;
};

void SaveDetectorInfo(const char* info)
{
    std::ofstream outFile("detector_info.json", std::ios::out | std::ios::app);
    
    if (outFile.is_open())
    {
        outFile << "{\n\t\"Detector Information\": \"" << info << "\"\n}" << std::endl;
        outFile.close();
        TZLogInfo("Detector information saved to detector_info.json~~~");
    }
    else
    {
        TZLogError("Failed to open file for writing!!!");
    }
}

TZ_INT DetectionCallback(SPtr<masd::StreamInfo>& media, void* ctx)
{
    TZLogInfo("Detection callback triggered!~~~");

    // Retrieve the context
    CallbackContext* callbackContext = reinterpret_cast<CallbackContext*>(ctx);
    if ((callbackContext == nullptr) || 
    (callbackContext->img == nullptr))
    {
        TZLogError("Error: Invalid context or image pointer!!!");
        return -1;
    }

    cv::Mat* img = callbackContext->img;
    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);

                #if RUN_TEXT_DRAWING
                if (!rect.Text.Text.empty())
                {
                    cv::putText(*img, rect.Text.Text, 
                                cv::Point(topLeft.x, topLeft.y - 10), 
                                cv::FONT_HERSHEY_SIMPLEX, 0.8, color, 2);
                }
                #endif
            }
        }
        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 Clip detector
    masd::Detector* detector = BuildDetector();
    if (detector == nullptr)
    {
        TZLogError("Failed to build Clip detector!!!");
        Dispose();
        return -1;
    }
    TZLogInfo("Clip detector built successfully~~~");

    // Step 3:
    // Initialize the Clip detector with configuration parameters
    std::string initParam = 
    "{\"gpu_id\": 0, \"model_path\": \"/home/flechazo/workspace/abandoned-object-detection-algo/clip/models/visual.onnx\"}";
    TZ_INT initDetResult = detector->Initialize(initParam);
    if (initDetResult != masd::MEC_OK)
    {
        TZLogError("Failed to initialize the Clip detector!!!");
        DestroyDetector(detector);
        Dispose();
        return -1;
    }
    TZLogInfo("Clip detector initialized successfully~~~");

    // Step 5:
    // Simulate frame processing with using test image
    cv::Mat testImage = cv::imread("/home/flechazo/workspace/abandoned-object-detection-algo/proposals/proposal_7.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);

    CallbackContext callbackContext{&testImage};
    detector->DoDetect(streamInfo, DetectionCallback, &callbackContext);

    // Step 6:
    // Print DetGetInformation
    char detectorInfo[4096];
    TZ_INT infoResult = GetInformation(detectorInfo);
    if (infoResult != masd::MEC_OK)
    {
        TZLogError("Failed to get detector information!!!");
        Dispose();
        return -1;
    }
    SaveDetectorInfo(detectorInfo);

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

    // Step 8:
    // 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;
}