#include <fstream>
#include <opencv2/opencv.hpp>
#include <iostream>
#include <string>
#include <vector>
#include <sstream>
#include <cstdlib>

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

// Function to split a string by a delimiter and return a vector of tokens
std::vector<std::string> splitString(const std::string& str, char delimiter)
{
    std::vector<std::string> tokens;
    std::stringstream ss(str);
    std::string item;
    while (std::getline(ss, item, delimiter))
    {
        // Trim whitespace from each token
        size_t start = item.find_first_not_of(" \t");
        size_t end = item.find_last_not_of(" \t");
        if (start != std::string::npos && end != std::string::npos)
        {
            tokens.push_back(item.substr(start, end - start + 1));
        }
        else if (start != std::string::npos)
        {
            tokens.push_back(item.substr(start));
        }
        // Ignore tokens that are all whitespace
    }
    return tokens;
}

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

void printUsage(const char* programName)
{
    std::cout << "Usage: " << programName << " --videos <video_paths> [--model <model_path>] [--confidence <confidence_threshold>]\n";
    std::cout << "  --videos     : Comma-separated list of video file paths (required)\n";
    std::cout << "  --model      : Path to the detection model (optional, default: ../../yolo-crowd-ft-e60.trt)\n";
    std::cout << "  --confidence : Confidence threshold (optional, range: 0.0 - 1.0, default: 0.25)\n";
    std::cout << "Example:\n";
    std::cout << "  " << programName << " --videos video1.mp4,video2.mp4 --model ../../models/yolo-crowd-ft-e60.trt --confidence 0.3\n";
}

int main(int argc, char* argv[])
{
    // Initialize log
    INITIALIZE_LOGGER_NORMAL("test", "./test.log", 1, 100, 6, 1, 1);

    // Parse command-line arguments
    std::string videoPathsStr;
    std::string modelPath = "../../models/yolo-crowd-ft-e60.trt"; // Default model path
    float confidenceThreshold = 0.25f; // Default confidence threshold

    if (argc < 3)
    {
        std::cout << "Error: Insufficient arguments provided.\n";
        printUsage(argv[0]);
        return -1;
    }

    for (int i = 1; i < argc; ++i)
    {
        std::string arg = argv[i];
        if (arg == "--videos" && i + 1 < argc)
        {
            videoPathsStr = argv[++i];
        }
        else if (arg == "--model" && i + 1 < argc)
        {
            modelPath = argv[++i];
        }
        else if (arg == "--confidence" && i + 1 < argc)
        {
            try
            {
                confidenceThreshold = std::stof(argv[++i]);
                if (confidenceThreshold < 0.0f || confidenceThreshold > 1.0f)
                {
                    std::cout << "Error: Confidence threshold must be between 0.0 and 1.0.\n";
                    return -1;
                }
            }
            catch (const std::invalid_argument& e)
            {
                std::cout << "Error: Invalid confidence threshold value.\n";
                return -1;
            }
            catch (const std::out_of_range& e)
            {
                std::cout << "Error: Confidence threshold value out of range.\n";
                return -1;
            }
        }
        else
        {
            std::cout << "Error: Unknown or incomplete argument '" << arg << "'.\n";
            printUsage(argv[0]);
            return -1;
        }
    }

    // Check if --videos was provided
    if (videoPathsStr.empty())
    {
        std::cout << "Error: --videos argument is required.\n";
        printUsage(argv[0]);
        return -1;
    }

    // Split video paths
    std::vector<std::string> videoPaths = splitString(videoPathsStr, ',');

    if (videoPaths.empty())
    {
        std::cout << "Error: No video paths provided.\n";
        printUsage(argv[0]);
        return -1;
    }

    // Display the received parameters
    TZLogInfo("Received parameters:");
    TZLogInfo("  Video Paths:");
    for (const auto& path : videoPaths)
    {
        TZLogInfo("    %s", path.c_str());
    }
    TZLogInfo("  Model Path: %s", modelPath.c_str());
    TZLogInfo("  Confidence Threshold: %.2f", confidenceThreshold);

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

    // Iterate over each video path
    for (const auto& videoPath : videoPaths)
    {
        TZLogInfo("Processing video: %s", videoPath.c_str());

        // 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!!!");
            DestroyDetector(detector);
            Dispose();
            return -1;
        }
        TZLogInfo("Headcount strategy built successfully~~~");

        // Step 3:
        // Initialize the yolo-crowd detector with configuration parameters
        std::stringstream initParamStream;
        initParamStream << "{"
                        << "\"gpu_id\": 0, "
                        << "\"max_objects\": 1024, "
                        << "\"confidence_threshold\": " << confidenceThreshold << ", "
                        << "\"nms_threshold\": 0.5, "
                        << "\"model_path\": \"" << modelPath << "\""
                        << "}";

        std::string initParam = initParamStream.str();

        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 strategy!!!");
            DestroyDetector(detector);
            DestroyStrategy(strategy);
            Dispose();
            return -1;
        }
        TZLogInfo("Headcount strategy initialized successfully~~~");

        // Step 4:
        // Set yolo-crowd detection configuration (optional)
        std::stringstream detectConfigStream;
        detectConfigStream << "{"
                           << "\"freq\": 0, "
                           << "\"target\": {"
                           << "\"target_class\": 0, "
                           << "\"target_threshold\": " << confidenceThreshold
                           << "}, "
                           << "\"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 }"
                           << "]"
                           << "}";

        std::string detectConfig = detectConfigStream.str();

        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:
        // Open the video file
        cv::VideoCapture videoCapture(videoPath);
        if(!videoCapture.isOpened())
        {
            TZLogError("Failed to load video: %s!!!", videoPath.c_str());
            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));

        // Generate output video file name based on input video file name
        std::string outputVideoPath = "output_" + videoPath.substr(videoPath.find_last_of("/\\") + 1);
        cv::VideoWriter videoWriter(outputVideoPath,
                                    cv::VideoWriter::fourcc('X', '2', '6', '4'),
                                    videoFPS, cv::Size(videoWidth, videoHeight));
        if (!videoWriter.isOpened())
        {
            TZLogError("Failed to open video writer for %s!!!", outputVideoPath.c_str());
            DestroyDetector(detector);
            DestroyStrategy(strategy);
            Dispose();
            return -1;
        }

        TZLogInfo("Processing video: %s", videoPath.c_str());
        TZLogInfo("Output video will be saved as: %s", outputVideoPath.c_str());

        cv::Mat frame;
        while(videoCapture.read(frame))
        {
            if(frame.empty())
            {
                TZLogError("Failed to read frame from video: %s!!!", videoPath.c_str());
                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!!!");
            DestroyDetector(detector);
            DestroyStrategy(strategy);
            Dispose();
            return -1;
        }
        SaveDetectorInfo(detectorInfo);

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

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