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

void writeYUV420Frame(std::ofstream &file, cv::Mat &frame) {
    // Convert the frame from BGR to YUV
    cv::Mat yuv_frame;
    cv::cvtColor(frame, yuv_frame, cv::COLOR_BGR2YUV_I420);

    // Write Y plane
    file.write(reinterpret_cast<const char *>(yuv_frame.data), frame.rows * frame.cols);

    // Write U and V planes
    int uvHeight = frame.rows / 2;
    int uvWidth = frame.cols / 2;
    file.write(reinterpret_cast<const char *>(yuv_frame.data + frame.rows * frame.cols), uvHeight * uvWidth);
    file.write(reinterpret_cast<const char *>(yuv_frame.data + frame.rows * frame.cols + uvHeight * uvWidth),
               uvHeight * uvWidth);
}

int main(int argc, char **argv) {
    // Open the default camera
    cv::VideoCapture cap(0, cv::CAP_V4L2); // 改成0或合适的索引 /dev/video0
    if (!cap.isOpened()) {
        std::cerr << "Error: Could not open camera" << std::endl;
        return -1;
    }

    // Set frame width and height to 1280x720
    cap.set(cv::CAP_PROP_FRAME_WIDTH, 1280);
    cap.set(cv::CAP_PROP_FRAME_HEIGHT, 720);

    // Open a file to write the YUV420 video
    std::ofstream yuv_file("output.yuv", std::ios::out | std::ios::binary);
    if (!yuv_file.is_open()) {
        std::cerr << "Error: Could not open output file" << std::endl;
        return -1;
    }

    // Capture frames from the camera
    while (true) {
        cv::Mat frame;
        cap >> frame; // Capture a new frame
        if (frame.empty()) {
            std::cerr << "Error: Could not capture frame" << std::endl;
            break;
        }

        // Write the frame to the YUV file
        writeYUV420Frame(yuv_file, frame);

        // Display the frame (直接显示BGR帧)
        cv::imshow("Webcam", frame);

        // Break the loop on 'q' key press
        if (cv::waitKey(30) >= 0) {
            break;
        }
    }

    // Clean up
    yuv_file.close();
    cap.release();
    cv::destroyAllWindows();
    return 0;
}