JTJAI_Media/main.cpp

#include <iostream>
#include <signal.h>
#include <chrono>
#include <thread>
#include "config.h"
#include "scheduler.h"
#include "reporter.h"

using namespace jtjai_media;

// 全局变量用于信号处理
static std::unique_ptr<StreamScheduler> g_scheduler;
static std::atomic<bool> g_interrupted(false);

// 信号处理函数
void signal_handler(int signal) {
    std::cout << "\n接收到信号 " << signal << "，正在停止程序..." << std::endl;
    g_interrupted.store(true);
    if (g_scheduler) {
        g_scheduler->stop_execution();
    }
}

// 进度回调函数
void progress_callback(const StreamScheduler::SchedulerStats& stats) {
    std::cout << "\r进度: " << stats.completed_tasks << "/" << stats.total_tasks 
              << " 完成, " << stats.failed_tasks << " 失败, "
              << "当前并发: " << stats.max_concurrent_used << " 流" << std::flush;
}

// 报告生成回调函数
void report_callback(const ConfigManager& config_mgr, int cycle_number, 
                    const std::vector<RTSPClientStats>& client_stats,
                    const StreamScheduler::SchedulerStats& scheduler_stats) {
    std::cout << "生成第 " << cycle_number << " 轮执行报告..." << std::endl;
    
    // 为每个周期生成不同的报告文件
    ConfigManager cycle_config = config_mgr;
    auto cycle_global_config = cycle_config.get_global_config();
    
    // 修改报告文件名，包含周期编号
    std::string base_name = cycle_global_config.report_filename;
    size_t dot_pos = base_name.find_last_of('.');
    if (dot_pos != std::string::npos) {
        base_name = base_name.substr(0, dot_pos) + "_cycle" + std::to_string(cycle_number) + base_name.substr(dot_pos);
    } else {
        base_name = base_name + "_cycle" + std::to_string(cycle_number);
    }
    cycle_global_config.report_filename = base_name;
    cycle_config.set_global_config(cycle_global_config);
    
    ResultReporter cycle_reporter(cycle_config);
    if (cycle_reporter.generate_report(client_stats, scheduler_stats)) {
        std::cout << "第 " << cycle_number << " 轮报告生成成功" << std::endl;
    } else {
        std::cerr << "第 " << cycle_number << " 轮报告生成失败" << std::endl;
    }
}

int main(int argc, char* argv[]) {
    std::cout << "========================================" << std::endl;
    std::cout << "RTSP视频流并发拉取系统" << std::endl;
    std::cout << "========================================" << std::endl;
    
    // 注册信号处理器
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    // 解析命令行参数
    std::string config_file = "/Users/wenhongquan/CLionProjects/jtjai_media/config.json";
    if (argc > 1) {
        config_file = argv[1];
    }
    
    std::cout << "使用配置文件: " << config_file << std::endl;
    
    try {
        // 加载配置
        ConfigManager config_mgr;
        if (!config_mgr.load_from_file(config_file)) {
            std::cerr << "加载配置文件失败: " << config_file << std::endl;
            return 1;
        }
        
        std::cout << "配置加载成功" << std::endl;
        std::cout << config_mgr.to_string() << std::endl;
        
        // 创建调度器
        g_scheduler = std::make_unique<StreamScheduler>(config_mgr);
        
        // 设置进度回调
        g_scheduler->set_progress_callback(progress_callback);
        
        const auto& global_config = config_mgr.get_global_config();
        int current_cycle = 0;
        
        // 开始多轮询循环
        while ((global_config.poll_cycles == -1 || current_cycle < global_config.poll_cycles) && !g_interrupted.load()) {
            current_cycle++;
            std::cout << "\n========== 轮询周期 " << current_cycle;
            if (global_config.poll_cycles != -1) {
                std::cout << "/" << global_config.poll_cycles;
            }
            std::cout << " ==========" << std::endl;
            
            // 设置当前周期的报告回调
            g_scheduler->set_report_callback([&config_mgr, current_cycle](
                const std::vector<RTSPClientStats>& client_stats,
                const StreamScheduler::SchedulerStats& scheduler_stats) {
                report_callback(config_mgr, current_cycle, client_stats, scheduler_stats);
            });
            
            // 启动调度执行
            std::cout << "开始执行RTSP流拉取任务..." << std::endl;
            if (!g_scheduler->start_execution()) {
                std::cerr << "启动调度器失败" << std::endl;
                return 1;
            }
            
            // 等待执行完成或中断
            std::cout << "等待任务执行中..." << std::endl;
            while (g_scheduler->is_running() && !g_interrupted.load()) {
                std::this_thread::sleep_for(std::chrono::milliseconds(500));
            }
            
            std::cout << "任务执行完成，等待所有任务线程结束..." << std::endl;
            g_scheduler->wait_for_completion();
            
            // 检查是否需要继续下一轮
            if (g_interrupted.load()) {
                std::cout << "接收到中断信号，退出轮询" << std::endl;
                break;
            }
            
            if (global_config.poll_cycles == -1 || current_cycle < global_config.poll_cycles) {
                std::cout << "\n等待 " << global_config.cycle_interval_seconds << " 秒后开始下一轮询..." << std::endl;
                
                // 在等待期间检查中断信号
                for (int i = 0; i < global_config.cycle_interval_seconds && !g_interrupted.load(); ++i) {
                    std::this_thread::sleep_for(std::chrono::seconds(1));
                    if ((i + 1) % 10 == 0) {
                        std::cout << "\r剩余等待时间: " << (global_config.cycle_interval_seconds - i - 1) << " 秒" << std::flush;
                    }
                }
                std::cout << std::endl;
                
                if (!g_interrupted.load()) {
                    // 重置调度器状态，准备下一轮
                    if (!g_scheduler->reset_for_next_cycle()) {
                        std::cerr << "重置调度器失败" << std::endl;
                        break;
                    }
                }
            }
        }
        
        std::cout << "\n程序执行完成" << std::endl;
        
    } catch (const std::exception& e) {
        std::cerr << "程序执行异常: " << e.what() << std::endl;
        return 1;
    }
    
    return 0;
}