#include <fstream>
#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"
#include "ilogger.hpp"
#include "trt_infer.hpp"

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

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

#define RUN_TEXT_DRAWING 0

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

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 || !callbackContext->img || callbackContext->strategy == nullptr)
    {
        TZLogError("Error: Invalid context or image pointer!!!");
        return -1;
    }

    cv::Mat* img = callbackContext->img;
    masd::Strategy* strategy = callbackContext->strategy;

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

    // Process the strategy inside the callback
    TZ_INT strategyResult = strategy->DoStrategy(media);
    if (strategyResult != masd::MEC_OK)
    {
        TZLogError("Headcount strategy failed to process the stream info!!!");
        return -1;
    }
    TZLogInfo("Headcount strategy processed the stream info successfully~~~");

    SPtr<masd::StraProducing> straProducing = media->GetStraProducing();
    if (straProducing)
    {
        std::string headcountResult = straProducing->Result.RstName;
        TZLogInfo("Detected headcount: %s", headcountResult.c_str());

        int fontFace = cv::FONT_HERSHEY_SIMPLEX;
        double fontScale = 0.8;
        int thickness = 2;
        cv::Scalar textColor(0, 255, 0);
        cv::Point textPosition(10, 30);

        cv::putText(*img, headcountResult, textPosition, fontFace, fontScale, textColor, thickness);
    }
    else
    {
        TZLogWarn("No headcount result found in the stream info.");
    }

    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 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~~~");

    // Build headcount strategy
    masd::Strategy* strategy = BuildStrategy();
    if(strategy == nullptr)
    {
        TZLogError("Failed to build headcount strategy!!!");
        Dispose();
        return -1;
    }
    TZLogInfo("Headcount strategy 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);
        DestroyStrategy(strategy);
        Dispose();
        return -1;
    }
    TZLogInfo("Yolo-crowd detector initialized successfully~~~");

    // Initialize the headcount strategy
    TZ_INT initStraResult = strategy->Initialize();
    if (initStraResult != masd::MEC_OK)
    {
        TZLogError("Failed to initialize the headcount detector!!!");
        DestroyDetector(detector);
        DestroyStrategy(strategy);
        Dispose();
        return -1;
    }
    TZLogInfo("Headcount strategy initialized successfully~~~");

    // Step 4:
    // Set yolo-crowd detection configuration (optional)
    std::string detectConfig = R"({
        "freq": 0,
        "target": {
            "target_class": 0,
            "target_threshold": 0.7
        },
        "focusArea": [
            { "LTX": 0.0, "LTY": 0.0, "RBX": 1.0, "RBY": 1.0 }
        ],
        "ignoreArea": [
            { "LTX": 0.0, "LTY": 0.0, "RBX": 0.0, "RBY": 0.0 }
        ]
    })";

    TZ_INT setDetCfgResult = detector->SetDetectCfg(detectConfig);
    if (setDetCfgResult != masd::MEC_OK)
    {
        TZLogError("Failed to set yolo-crowd detection configuration!!!");
        DestroyDetector(detector);
        DestroyStrategy(strategy);
        Dispose();
        return -1;
    }
    TZLogInfo("Yolo-crowd detection configuration set successfully~~~");

    // Set headcount strategy configuration (optional)
    std::string headcountConfig = R"({
        "TimeThreshold": 5
    })";

    TZ_INT setStraCfgResult = strategy->SetStrategyCfg(headcountConfig);
    if (setStraCfgResult != masd::MEC_OK)
    {
        TZLogError("Failed to set headcount strategy configuration!!!");
        DestroyDetector(detector);
        DestroyStrategy(strategy);
        Dispose();
        return -1;
    }
    TZLogInfo("Headcount strategy configuration set 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);
        DestroyStrategy(strategy);
        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);
        DestroyStrategy(strategy);
        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();
        SPtr<masd::Media> mediaResource = std::make_shared<masd::Media>(length);
        mediaResource->Width = frame.cols;
        mediaResource->Height = frame.rows;
        mediaResource->DataType = frame.type();
        mediaResource->Mem = frame.data;

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

        CallbackContext callbackContext{strategy, &frame, &videoWriter};
        detector->DoDetect(streamInfo, DetectionCallback, &callbackContext);

        SEMA.Wait();
    }

    // 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~~~");

    // Destroy the strategy
    DestroyStrategy(strategy);
    TZLogInfo("Detector strategy 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;
}