rtsp_tensorrt/pipeline/output/video_writer.cpp

#include "video_writer.hpp"
#include "../common/logger.hpp"
#include <filesystem>
#include <fstream>

namespace pipeline {

VideoWriter::VideoWriter() = default;

VideoWriter::~VideoWriter() {
    release();
}

bool VideoWriter::validateConfig(const OutputTargetConfig& config) {
    if (config.name.empty()) {
        setLastError("Empty name in config");
        return false;
    }

    if (config.path.empty()) {
        setLastError("Empty path in config");
        return false;
    }

    if (config.fps <= 0) {
        setLastError("Invalid fps in config: " + std::to_string(config.fps));
        return false;
    }

    // 检查输出目录是否可写
    try {
        auto path = std::filesystem::path(config.path);
        auto dir = path.parent_path();
        Logger::info("Checking directory: " + dir.string());

        // 检查路径安全性
        if (dir.string().find("/tmp") != 0 && 
            dir.string().find("/var/tmp") != 0 && 
            dir.string().find("/home") != 0) {
            Logger::info("Invalid directory location");
            setLastError("Output directory must be under /tmp, /var/tmp or /home: " + dir.string());
            return false;
        }

        // 检查路径深度
        int depth = 0;
        for (const auto& part : dir) {
            depth++;
            if (depth > 8) {  // 限制路径深度
                Logger::info("Directory path too deep");
                setLastError("Directory path exceeds maximum depth: " + dir.string());
                return false;
            }
        }
        
        // 检查目录是否存在且可写
        std::error_code ec;
        if (std::filesystem::exists(dir, ec)) {
            if (!std::filesystem::is_directory(dir)) {
                Logger::info("Path exists but is not a directory");
                setLastError("Output path exists but is not a directory: " + dir.string());
                return false;
            }
            
            // 检查目录权限
            auto perms = std::filesystem::status(dir, ec).permissions();
            if (ec) {
                Logger::info("Error getting directory status: " + ec.message());
                setLastError("Error checking directory: " + ec.message());
                return false;
            }
            
            // 在Unix系统中,检查用户写权限
            if ((perms & std::filesystem::perms::owner_write) == std::filesystem::perms::none) {
                Logger::info("No write permission for directory");
                setLastError("No write permission for directory: " + dir.string());
                return false;
            }
            
            // 如果目录存在且有写权限,尝试创建一个临时文件
            auto test_file = dir / "test_write_permission";
            {
                std::ofstream test(test_file);
                if (!test) {
                    Logger::info("Cannot create test file in directory");
                    setLastError("Directory is not writable: " + dir.string());
                    return false;
                }
                test << "test";
            }
            std::filesystem::remove(test_file);
            
            // 检查实际输出文件
            if (std::filesystem::exists(path)) {
                // 如果文件已存在,检查是否可写
                auto file_perms = std::filesystem::status(path, ec).permissions();
                if (ec) {
                    Logger::info("Error getting file status: " + ec.message());
                    setLastError("Error checking output file: " + ec.message());
                    return false;
                }
                
                if ((file_perms & std::filesystem::perms::owner_write) == std::filesystem::perms::none) {
                    Logger::info("No write permission for existing file");
                    setLastError("No write permission for existing file: " + path.string());
                    return false;
                }
                
                // 尝试打开文件进行写入
                std::ofstream test_write(path);
                if (!test_write) {
                    Logger::info("Cannot open existing file for writing");
                    setLastError("Cannot open existing file for writing: " + path.string());
                    return false;
                }
            } else {
                // 如果文件不存在,尝试创建
                std::ofstream test_write(path);
                if (!test_write) {
                    Logger::info("Cannot create output file");
                    setLastError("Cannot create output file: " + path.string());
                    return false;
                }
                // 创建成功后删除测试文件
                test_write.close();
                std::filesystem::remove(path);
            }
            
        } else {
            // 目录不存在,检查父目录是否可写
            auto parent = dir;
            while (!parent.empty() && !std::filesystem::exists(parent)) {
                parent = parent.parent_path();
            }
            
            if (parent.empty()) {
                Logger::info("No valid parent directory found");
                setLastError("No valid parent directory found for: " + dir.string());
                return false;
            }
            
            // 检查父目录权限
            auto perms = std::filesystem::status(parent, ec).permissions();
            if (ec) {
                Logger::info("Error getting parent directory status: " + ec.message());
                setLastError("Error checking parent directory: " + ec.message());
                return false;
            }
            
            if ((perms & std::filesystem::perms::owner_write) == std::filesystem::perms::none) {
                Logger::info("No write permission for parent directory");
                setLastError("No write permission for parent directory: " + parent.string());
                return false;
            }
            
            // 尝试创建目录
            if (!std::filesystem::create_directories(dir)) {
                Logger::info("Failed to create directory");
                setLastError("Failed to create output directory: " + dir.string());
                return false;
            }
            
            // 尝试创建测试文件
            std::ofstream test_write(path);
            if (!test_write) {
                Logger::info("Cannot create output file in new directory");
                setLastError("Cannot create output file in new directory: " + path.string());
                return false;
            }
            // 创建成功后删除测试文件
            test_write.close();
            std::filesystem::remove(path);
        }
        
    } catch (const std::exception& e) {
        Logger::info("Exception while checking directory: " + std::string(e.what()));
        setLastError("Error checking output directory: " + std::string(e.what()));
        return false;
    }

    return true;
}

bool VideoWriter::init(const OutputTargetConfig& config) {
    std::lock_guard<std::mutex> lock(mutex_);
    
    if (is_initialized_) {
        Logger::warning("VideoWriter already initialized");
        return true;
    }

    if (!validateConfig(config)) {
        return false;
    }

    name_ = config.name;
    output_path_ = config.path;
    fps_ = config.fps;
    codec_ = config.codec;
    bitrate_ = config.bitrate;

    // 立即尝试打开文件
    cv::Size initial_size(640, 480); // 使用一个初始大小,后续会根据实际帧大小调整
    if (!openWriter(initial_size)) {
        return false;
    }
    
    is_initialized_ = true;
    Logger::info("VideoWriter initialized successfully: " + name_);
    return true;
}

bool VideoWriter::openWriter(const cv::Size& size) {
    if (writer_.isOpened()) {
        if (frame_size_ == size) {
            return true;
        }
        writer_.release();
    }

    frame_size_ = size;
    int fourcc = getFourcc(codec_);
    
    try {
        // 设置临时文件路径
        std::string temp_path = output_path_ + ".tmp";
        
        // 打开临时文件
        writer_.open(temp_path, fourcc, fps_, size, true);  // 添加isColor参数
        if (!writer_.isOpened()) {
            setLastError("Failed to open video writer: " + temp_path);
            is_opened_ = false;
            return false;
        }
        
        // 设置编码器参数
        writer_.set(cv::VIDEOWRITER_PROP_QUALITY, 100);  // 最高质量
        
        is_opened_ = true;
        return true;
    } catch (const cv::Exception& e) {
        setLastError("OpenCV error: " + std::string(e.what()));
        is_opened_ = false;
        return false;
    } catch (const std::exception& e) {
        setLastError("Error opening video writer: " + std::string(e.what()));
        is_opened_ = false;
        return false;
    }
}

bool VideoWriter::write(const cv::Mat& frame) {
    std::lock_guard<std::mutex> lock(mutex_);
    
    if (!is_initialized_) {
        setLastError("VideoWriter not initialized");
        return false;
    }

    if (frame.empty()) {
        setLastError("Input frame is empty");
        return false;
    }

    try {
        if (!openWriter(frame.size())) {
            return false;
        }

        // 确保帧格式正确
        cv::Mat write_frame;
        if (frame.channels() != 3) {
            cv::cvtColor(frame, write_frame, cv::COLOR_GRAY2BGR);
        } else {
            write_frame = frame;
        }

        // 写入帧
        writer_.write(write_frame);
        return true;
    } catch (const cv::Exception& e) {
        setLastError("OpenCV error: " + std::string(e.what()));
        return false;
    } catch (const std::exception& e) {
        setLastError("Error writing frame: " + std::string(e.what()));
        return false;
    }
}

bool VideoWriter::isOpened() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return is_initialized_ && is_opened_;
}

void VideoWriter::release() {
    std::lock_guard<std::mutex> lock(mutex_);
    if (writer_.isOpened()) {
        writer_.release();
        
        // 如果使用了临时文件，在这里重命名
        std::string temp_path = output_path_ + ".tmp";
        if (std::filesystem::exists(temp_path)) {
            try {
                std::filesystem::rename(temp_path, output_path_);
            } catch (const std::exception& e) {
                Logger::error("Failed to rename temporary file: " + std::string(e.what()));
            }
        }
    }
    is_initialized_ = false;
    is_opened_ = false;
    frame_size_ = cv::Size(0, 0);
}

std::string VideoWriter::getName() const {
    return name_;
}

bool VideoWriter::getStatus(std::string& error_msg) const {
    std::lock_guard<std::mutex> lock(mutex_);
    error_msg = last_error_;
    return is_initialized_ && is_opened_;
}

int VideoWriter::getFourcc(const std::string& codec) const {
    if (codec == "h264") {
        return cv::VideoWriter::fourcc('a','v','c','1');  // 使用avc1代替H264
    } else if (codec == "mp4v") {
        return cv::VideoWriter::fourcc('m','p','4','v');
    } else if (codec == "mjpg") {
        return cv::VideoWriter::fourcc('M','J','P','G');
    } else if (codec == "xvid") {
        return cv::VideoWriter::fourcc('X','V','I','D');
    } else {
        Logger::warning("Unknown codec: " + codec + ", using default avc1");
        return cv::VideoWriter::fourcc('a','v','c','1');
    }
}

void VideoWriter::setLastError(const std::string& error) {
    last_error_ = error;
    Logger::error("VideoWriter error: " + error);
}

} // namespace pipeline