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

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

// libmascommon
#include "MemPool.h"

// deps
#include "DetUtils.h"

// libabndobjdetector
#include "DetectorAPI.h"
#include "AbandObjDetector.h"

// libheadcountstrategy
//#include "StrategyAPI.h"
//#include "HeadCountStrategy.h"

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

CallbackContext callbackContext;

tzc::Semaphore SEMA;

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) || 
    (callbackContext->videoWriter == nullptr) || 
    (!callbackContext->videoWriter->isOpened()))
    {
        TZLogError("Error: Invalid context, image pointer, or uninitialized video writer!!!");
        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());
        TZLogInfo("Detection Result: %s~~~", detProducing->Result.c_str());

        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);
        }
    }

    callbackContext->videoWriter->write(*img);

    SEMA.Signal();

    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 abandobj detector
    masd::Detector* detector = BuildDetector();
    if (detector == nullptr)
    {
        TZLogError("Failed to build abndobj detector!!!");
        Dispose();
        return -1;
    }
    TZLogInfo("AbandObj detector built successfully~~~");

    // Step 3:
    // Initialize the yolo-crowd detector with configuration parameters
    const std::string initParam = R"(
    {
        "algo": "MOG2",
        "video_path": "../../media/video_cutter.mp4",
        "factor": 4,
        "short_term_rate": 0.01,
        "short_term_history": 200,
        "long_term_rate": 0.0005,
        "long_term_history": 5000,
        "var_threshold": 16.0,
        "detect_shadows": false,
        "iou_threshold": 0.6,
        "area_threshold": 625.0,
        "perimeter_threshold": 100.0
    }
    )";
    TZ_INT initDetResult = detector->Initialize(initParam);
    if (initDetResult != masd::MEC_OK)
    {
        TZLogError("Failed to initialize the abandobj detector!!!");
        DestroyDetector(detector);
        Dispose();
        return -1;
    }
    TZLogInfo("AbandObj detector initialized successfully~~~");


    // Step 5:
    // Simulate frame processing with using test MP4
    cv::VideoCapture videoCapture("../../media/video.mp4");
    if(!videoCapture.isOpened())
    {
        TZLogError("Failed to load test video!!!");
        DestroyDetector(detector);
        Dispose();
        return -1;
    }
    
    TZ_INT videoWidth = static_cast<TZ_INT>(videoCapture.get(cv::CAP_PROP_FRAME_WIDTH));
    TZ_INT videoHeight = static_cast<TZ_INT>(videoCapture.get(cv::CAP_PROP_FRAME_HEIGHT));
    TZ_INT videoFPS = static_cast<TZ_INT>(videoCapture.get(cv::CAP_PROP_FPS));

    cv::VideoWriter videoWriter("output_video.mp4",
                                cv::VideoWriter::fourcc('X', '2', '6', '4'),
                                videoFPS, cv::Size(videoWidth, videoHeight));
    if (!videoWriter.isOpened())
    {
        TZLogError("Failed to open video writer!!!");
        DestroyDetector(detector);
        Dispose();
        return -1;
    }

    cv::Mat frame;
    while(videoCapture.read(frame))
    {
        if(frame.empty())
        {
            TZLogError("Failed to read frame from video!!!");
            break;
        }

        TZ_INT length = frame.total() * frame.elemSize();

        void* allocatedMem = pool->AllocAvailMem(length);
        if (!allocatedMem)
        {
            TZLogError("Failed to allocate memory from the pool!");
            break;
        }

        SPtr<masd::Media> mediaResource = std::make_shared<masd::Media>(length);
        mediaResource->Width = frame.cols;
        mediaResource->Height = frame.rows;
        mediaResource->DataType = frame.type();
        mediaResource->Mem = allocatedMem;
        std::memcpy(mediaResource->Mem, frame.data, length);

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

        callbackContext.img = new cv::Mat(frame.clone());
        callbackContext.videoWriter = &videoWriter; 

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

        SEMA.Wait();

        pool->ReturnFreeMem(allocatedMem, length);
        TZ_delete(callbackContext.img);
    }

    // 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;
}