#define NOMINMAX  
#include "hailo/hailort.hpp"
#include "hailo/vdevice.hpp"
#include "hailo/expected.hpp"
#include "hailo/hailort_common.hpp"
#include "hailo/runtime_statistics.hpp"
#include "VideoToAI.hpp"

#include <iostream>
#include <chrono>
#include <thread>
#include <fstream>

#include <cstring>
#if defined(__unix__)
#include <sys/mman.h>
#endif
//#include <include/common.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <mutex>
#include <condition_variable>
#include <queue>
#define BATCH_COUNT (1000)
#define BATCH_SIZE (4)
#define FPS (30)  // Set target FPS to 30

using namespace hailort;
using namespace VideoToAI;

std::mutex mtx;
std::condition_variable cv_var;
std::queue<cv::Mat> frameQueue;
bool doneReading = false;

cv::Mat readYUV420(const std::string& yuvFile, int width, int height, int frameIndex) {
    int frameSize = width * height + (width * height / 2);
    std::ifstream yuvStream(yuvFile, std::ios::binary);
    if (!yuvStream.is_open()) {
        throw std::runtime_error("Cannot open YUV file");
    }

    // Calculate file size for bounds checking
    yuvStream.seekg(0, std::ios::end);
    std::streamsize fileSize = yuvStream.tellg();
    if (frameIndex * frameSize >= fileSize) {
        throw std::runtime_error("Frame index out of range. Check the input file size.");
    }

    yuvStream.seekg(frameSize * frameIndex, std::ios::beg);
    std::vector<uint8_t> buffer(frameSize);
    yuvStream.read(reinterpret_cast<char*>(buffer.data()), frameSize);
    if (yuvStream.gcount() != frameSize) {
        throw std::runtime_error("Failed to read full YUV frame. Possibly corrupted file.");
    }

    return cv::Mat(height + height / 2, width, CV_8UC1, buffer.data()).clone();
}


void readYUVThread(const std::string& yuvFile, int width, int height, int numFrames,
    ConfiguredInferModel& configured_infer_model,
    std::shared_ptr<InferModel>& infer_model, VDevice* vdevice) {
    std::ifstream yuvStream(yuvFile, std::ios::binary);
    if (!yuvStream.is_open()) {
        throw std::runtime_error("Cannot open YUV file");
    }
    std::cout << "[INFO] Opened YUV file: " << yuvFile << std::endl;
    std::vector<VideoToAI::DetectionResult> detection_results;
    cv::Ptr<cv::BackgroundSubtractorMOG2> fgbg_mog2 = cv::createBackgroundSubtractorMOG2();

    for (int i = 0; i < numFrames; ++i) {
        try {
            std::cout << "[INFO] Reading frame " << i + 1 << " of " << numFrames << std::endl;

            // Read the YUV frame
            cv::Mat yuvFrame = readYUV420(yuvFile, width, height, i);
            if (yuvFrame.empty()) {
                throw std::runtime_error("YUV frame is empty.");
            }

            // Split into Y and UV planes
            cv::Mat yPlane = yuvFrame(cv::Rect(0, 0, width, height));
            cv::Mat uvPlane = yuvFrame(cv::Rect(0, height, width, height / 2));

            // Apply letterbox on Y and UV planes
            cv::Mat resizedYPlane, resizedUVPlane;
            letterbox(yPlane, resizedYPlane, cv::Size(640, 640));
            letterbox(uvPlane, resizedUVPlane, cv::Size(640, 640 / 2));
            if (resizedYPlane.cols != 640 || resizedYPlane.rows != 640 ||
                resizedUVPlane.cols != 640 || resizedUVPlane.rows != 320) {
                throw std::runtime_error("Unexpected size after letterbox. Ensure input dimensions and padding are correct.");
            }
            if (resizedYPlane.empty() || resizedUVPlane.empty()) {
                throw std::runtime_error("Letterbox failed to resize planes correctly.");
            }

            // Only Hailo infer
            hailoInference(resizedYPlane, resizedUVPlane, configured_infer_model, infer_model, vdevice, detection_results);

            // Catch detection results
            /*for (const auto& result : detection_results) {
                std::cout << "Class ID: " << result.class_id
                    << ", Confidence: " << result.confidence
                    << ", BBox: ["
                    << result.bbox.x_min << ", "
                    << result.bbox.y_min << ", "
                    << result.bbox.x_max << ", "
                    << result.bbox.y_max << "]" << std::endl;
            }*/

            std::cout << "[INFO] Processed YUV frame." << std::endl;

            // Combine Y and UV planes into a single frame
            std::vector<cv::Mat> channels = { resizedYPlane, resizedUVPlane };
            cv::Mat combinedFrame;
            cv::vconcat(channels, combinedFrame);

            // Push the processed frame to the queue
            std::unique_lock<std::mutex> lock(mtx);
            frameQueue.push(combinedFrame.clone());
            // lock.unlock(); // aviod didn't auto unlock
            cv_var.notify_one();
            // std::cout << "[INFO] Frame pushed to the queue." << std::endl;

        }
        catch (const std::exception& e) {
            std::cerr << "[ERROR] Read Error: " << e.what() << std::endl;
            break;
        }
    }

    std::unique_lock<std::mutex> lock(mtx);
    doneReading = true;
    cv_var.notify_one();
    // std::cout << "[INFO] Done reading YUV file." << std::endl;
}


void writeYUVThread(const std::string& outputFile) {
    std::ofstream outputYuvStream(outputFile, std::ios::binary);
    if (!outputYuvStream.is_open()) {
        std::cerr << "Cannot open output YUV file" << std::endl;
        return;
    }

    while (true) {
        std::unique_lock<std::mutex> lock(mtx);
        cv_var.wait(lock, [] { return !frameQueue.empty() || doneReading; });
        // lock.unlock(); // aviod didn't auto unlock
        while (!frameQueue.empty()) {
            cv::Mat frame = frameQueue.front();
            frameQueue.pop();
            outputYuvStream.write(reinterpret_cast<const char*>(frame.data), frame.total());
        }
        if (doneReading && frameQueue.empty()) {
            break;
        }
    }
    outputYuvStream.close();
    std::cout << "YUV video saved to " << outputFile << std::endl;
}

int main(int argc, char** argv) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <input_file> <output_file>" << std::endl;
        return -1;
    }
    // 
    std::string inputFilePath = argv[1];
    std::string outputFilePath = argv[2];

    try {
        auto vdevice = VDevice::create().expect("Failed to create vdevice");
        auto infer_model = vdevice->create_infer_model("yolov7_yuv21.hef").expect("Failed to create infer model");
        infer_model->set_batch_size(BATCH_SIZE);
        std::cout << "Set batch_size to " << BATCH_SIZE << std::endl;
        auto configured_infer_model = infer_model->configure().expect("Failed to configure infer model");
        //infer_model->set_batch_size(BATCH_SIZE);
        if (inputFilePath.size() >= 4 && inputFilePath.substr(inputFilePath.size() - 4) == ".yuv") {
            int width = 640;
            int height = 360;

            std::ifstream yuvStream(inputFilePath, std::ios::binary | std::ios::ate);
            if (!yuvStream.is_open()) {
                std::cerr << "Cannot open input YUV file" << std::endl;
                return -1;
            }

            std::streamsize fileSize = yuvStream.tellg();
            int frameSize = width * height + (width * height / 2);
            if (fileSize % frameSize != 0) {
                std::cerr << "Warning: File size is not an exact multiple of frame size. There might be incomplete frames." << std::endl;
            }
            int numFrames = static_cast<int>(fileSize / frameSize);
            std::thread reader(
                [&]() {
                    readYUVThread(inputFilePath, width, height, numFrames, configured_infer_model, infer_model, vdevice.get());
                });

            std::thread writer(writeYUVThread, outputFilePath);

            reader.join();
            writer.join();
        }
        else {
            std::cerr << "Unsupported file format. Please provide a .yuv file." << std::endl;
            return -1;
        }

        std::cout << "Video processing completed and saved to " << outputFilePath << std::endl;
    }
    catch (const hailort_error& e) {
        std::cerr << "Hailo initialization or inference error: " << e.what() << std::endl;
        return -1;
    }

    return 0;
}