rtsp_tensorrt/pipeline/output/output_manager.cpp

#include "output_manager.hpp"
#include "video_writer.hpp"
#include "rtsp_writer.hpp"
#include "../common/logger.hpp"

namespace pipeline {

OutputManager::OutputManager() = default;

OutputManager::~OutputManager() {
    cleanup();
}

bool OutputManager::addTarget(const OutputTargetConfig& config) {
    std::lock_guard<std::mutex> lock(mutex_);

    // 检查名称是否已存在
    if (writers_.find(config.name) != writers_.end()) {
        Logger::warning("Output target already exists: " + config.name);
        return false;
    }

    // 创建写入器实例
    auto writer = createWriter(config.type);
    if (!writer) {
        Logger::error("Failed to create writer for type: " + config.type);
        return false;
    }

    // 初始化写入器
    if (!writer->init(config)) {
        Logger::error("Failed to initialize writer: " + config.name);
        return false;
    }

    // 添加到映射
    writers_[config.name] = std::move(writer);
    Logger::info("Added output target: " + config.name);
    return true;
}

bool OutputManager::writeFrames(const cv::Mat& frame) {
    std::lock_guard<std::mutex> lock(mutex_);

    if (frame.empty()) {
        Logger::error("Input frame is empty");
        return false;
    }

    bool all_success = true;
    for (auto& [name, writer] : writers_) {
        if (!writer->write(frame)) {
            Logger::error("Failed to write frame to target: " + name);
            all_success = false;
        }
    }

    return all_success;
}

bool OutputManager::writeFrames(const std::string& source_name, const cv::Mat& frame) {
    std::lock_guard<std::mutex> lock(mutex_);

    if (frame.empty()) {
        Logger::error("Input frame is empty");
        return false;
    }

    // 检查输入源是否有映射的输出目标
    auto mapping_it = source_target_mapping_.find(source_name);
    if (mapping_it == source_target_mapping_.end()) {
        Logger::warning("No output targets mapped for source: " + source_name);
        return false;
    }

    bool all_success = true;
    for (const auto& target_name : mapping_it->second) {
        auto writer_it = writers_.find(target_name);
        if (writer_it == writers_.end()) {
            Logger::error("Writer not found for target: " + target_name);
            all_success = false;
            continue;
        }

        if (!writer_it->second->write(frame)) {
            Logger::error("Failed to write frame to target: " + target_name);
            all_success = false;
        }
    }

    return all_success;
}

bool OutputManager::getTargetStatus(const std::string& name, std::string& error_msg) const {
    std::lock_guard<std::mutex> lock(mutex_);

    auto it = writers_.find(name);
    if (it == writers_.end()) {
        error_msg = "Target not found: " + name;
        return false;
    }

    return it->second->getStatus(error_msg);
}

size_t OutputManager::getTargetCount() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return writers_.size();
}

std::vector<std::string> OutputManager::getTargetNames() const {
    std::lock_guard<std::mutex> lock(mutex_);
    std::vector<std::string> names;
    names.reserve(writers_.size());
    for (const auto& [name, _] : writers_) {
        names.push_back(name);
    }
    return names;
}

bool OutputManager::removeTarget(const std::string& name) {
    std::lock_guard<std::mutex> lock(mutex_);

    auto it = writers_.find(name);
    if (it == writers_.end()) {
        Logger::warning("Target not found: " + name);
        return false;
    }

    it->second->release();
    writers_.erase(it);
    Logger::info("Removed output target: " + name);
    return true;
}

void OutputManager::cleanup() {
    std::lock_guard<std::mutex> lock(mutex_);
    for (auto& [name, writer] : writers_) {
        writer->release();
    }
    writers_.clear();
    initialized_ = false;
}

std::unique_ptr<WriterInterface> OutputManager::createWriter(const std::string& type) {
    if (type == "video") {
        return std::make_unique<VideoWriter>();
    } else if (type == "rtsp") {
        return std::make_unique<RtspWriter>();
    } else {
        Logger::error("Unsupported writer type: " + type);
        return nullptr;
    }
}

bool OutputManager::addSourceTargetMapping(const std::string& source_name, 
                                         const std::vector<std::string>& target_names) {
    std::lock_guard<std::mutex> lock(mutex_);

    // 验证所有目标是否存在
    for (const auto& target : target_names) {
        if (writers_.find(target) == writers_.end()) {
            Logger::error("Target does not exist: " + target);
            return false;
        }
    }

    // 添加或更新映射
    auto& targets = source_target_mapping_[source_name];
    targets.clear();
    targets.insert(target_names.begin(), target_names.end());
    
    Logger::info("Added mapping for source: " + source_name + 
                " with " + std::to_string(target_names.size()) + " targets");
    return true;
}

bool OutputManager::removeSourceMapping(const std::string& source_name) {
    std::lock_guard<std::mutex> lock(mutex_);

    auto it = source_target_mapping_.find(source_name);
    if (it == source_target_mapping_.end()) {
        Logger::warning("No mapping found for source: " + source_name);
        return false;
    }

    source_target_mapping_.erase(it);
    Logger::info("Removed mapping for source: " + source_name);
    return true;
}

} // namespace pipeline