#include <iostream>
#include <fstream>
#include <thread>
#include <mutex>
#include <atomic>
#include <vector>
#include <filesystem>  // 確保資料夾存在
#include "libobsensor/ObSensor.hpp"
#include "opencv2/opencv.hpp"
#include "utils.hpp"

std::mutex fileMutex;
std::mutex bufferMutex;
std::atomic<bool> isRunning(true);
std::vector<cv::Mat> imageBuffer;
std::vector<std::pair<uint64_t, uint64_t>> timestampBuffer;
std::string outputDir = "C:\\Users\\ShawnPeng\\Desktop\\orbbec_output";

void createOutputDir(const std::string& dir) {
    if (!std::filesystem::exists(dir)) {
        std::filesystem::create_directories(dir);
    }
}

// **儲存 IMU 資料**
void saveIMUData(std::shared_ptr<ob::Sensor> sensor, const std::string& outputDir, std::string type) {
    std::string filename = outputDir + "/" + type + "_data.txt";
    std::ofstream imuFile(filename, std::ios::app);
    if (!imuFile) {
        std::cerr << "Error: Cannot open " << filename << std::endl;
        return;
    }

    sensor->start(sensor->getStreamProfileList()->getProfile(OB_PROFILE_DEFAULT), 
        [filename, type](std::shared_ptr<ob::Frame> frame) {
            std::lock_guard<std::mutex> lock(fileMutex);
            auto timeStamp = frame->timeStampUs();

            if (type == "Gyro") {
                auto imuFrame = frame->as<ob::GyroFrame>();
                if (imuFrame) {
                    auto value = imuFrame->value();
                    std::ofstream imuFile(filename, std::ios::app);
                    imuFile << timeStamp << "\t" << "\t"
                            << value.x << "\t" << value.y << "\t" << value.z << "\n";
                }
            } else {
                auto imuFrame = frame->as<ob::AccelFrame>();
                if (imuFrame) {
                    auto value = imuFrame->value();
                    std::ofstream imuFile(filename, std::ios::app);
                    imuFile << timeStamp << "\t" << "\t"
                            << value.x << "\t" << value.y << "\t" << value.z << "\n";
                }
            }
        });

    while (isRunning) {
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }

    sensor->stop();
}

void writeToDisk()
{
    while (isRunning)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
        std::lock_guard<std::mutex> lock(bufferMutex);
        while (!imageBuffer.empty())
        {
            std::string filename = outputDir + "/color_" + std::to_string(timestampBuffer[0].first) + "_" + std::to_string(timestampBuffer[0].second) + ".png";
            cv::imwrite(filename, imageBuffer[0]);
            imageBuffer.erase(imageBuffer.begin());
            timestampBuffer.erase(timestampBuffer.begin());
        }
    }
}

void saveColorImage(ob::Pipeline& pipeline) 
{
    ob::FormatConvertFilter formatConvertFilter;
    while (isRunning) 
    {
        auto frameset = pipeline.waitForFrames(3000000);
        if (!frameset) continue;
        auto colorFrame = frameset->colorFrame();
        if (!colorFrame) continue;
		
		/*
        if (colorFrame->format() != OB_FORMAT_RGB) {
            formatConvertFilter.setFormatConvertType(FORMAT_RGB_TO_BGR);
            colorFrame = formatConvertFilter.process(colorFrame)->as<ob::ColorFrame>();
        }*/

        auto frameTimestamp = colorFrame->getMetadataValue((OBFrameMetadataType)0);
        auto sensorTimestamp = colorFrame->getMetadataValue((OBFrameMetadataType)1);
        auto framenumber = colorFrame->getMetadataValue((OBFrameMetadataType)2);

        std::cout << "framenumber: " << framenumber << std::endl;
        cv::Mat colorRawMat(colorFrame->height(), colorFrame->width(), CV_8UC3, colorFrame->data());
        {
            std::lock_guard<std::mutex> lock(bufferMutex);
            imageBuffer.push_back(colorRawMat.clone());
            timestampBuffer.push_back({ frameTimestamp, sensorTimestamp });
        }
    }
}

int main() {
    createOutputDir(outputDir);
    ob::Pipeline pipeline;
    std::shared_ptr<ob::Config> config = std::make_shared<ob::Config>();
    auto colorProfiles = pipeline.getStreamProfileList(OB_SENSOR_COLOR);
    if (colorProfiles) {
        //auto colorProfile = colorProfiles->getVideoStreamProfile(1280, 720, OB_FORMAT_RGB, 30);
        auto colorProfile = colorProfiles->getVideoStreamProfile(640, 360, OB_FORMAT_RGB, 30);

        config->enableStream(colorProfile);
    }
    auto devList = ob::Context().queryDeviceList();
    if (devList->deviceCount() == 0) {
        std::cerr << "No device found!" << std::endl;
        return -1;
    }
    auto dev = devList->getDevice(0);
    auto sensorList = dev->getSensorList();

    std::shared_ptr<ob::Sensor> gyroSensor = sensorList->getSensor(OB_SENSOR_GYRO);
    std::shared_ptr<ob::Sensor> accelSensor = sensorList->getSensor(OB_SENSOR_ACCEL);
    
    
    pipeline.start(config);


// **啟動 IMU 與影像儲存執行緒**
    std::thread gyroThread(saveIMUData, gyroSensor, outputDir, "Gyro");
    std::thread accelThread(saveIMUData, accelSensor, outputDir, "Accel");
    std::thread colorThread(saveColorImage, std::ref(pipeline));
    std::thread diskThread(writeToDisk);

    std::cout << "Press ESC to stop recording..." << std::endl;
    while (true) {
        int key = _getch();
        if (key == 27) break;
    }

    isRunning = false;
    gyroThread.join();
    accelThread.join();
    colorThread.join();
    diskThread.join();
    pipeline.stop();
    return 0;
}