import json

import redis
import subprocess
import time
import os
import psutil  # 用于通过PID检查进程是否存活

def main():
    # Redis连接配置（根据实际情况修改）
    redis_host = '172.16.0.24'
    redis_port = 6379
    redis_db = 8
    redis_password = 'aiyingli@@123'
    redis_list_key = 'dso_image_gen'  # 存储任务的list键名
    redis_list_key_task = 'dso_image_gen_task'  # 存储任务的list键名

    # 最大并发数
    max_concurrent = 10

    try:
        # 连接Redis
        r = redis.Redis(
            host=redis_host,
            port=redis_port,
            db=redis_db,
            password=redis_password,
            decode_responses=True  # 自动将返回值解码为字符串
        )
        r.ping()
        print("成功连接到Redis")
    except Exception as e:
        print(f"Redis连接失败: {e}")
        return

    # 存储当前运行的实际任务进程PID（而非shell进程）
    task_pids = []

    try:
        print("开始监控任务队列... (按Ctrl+C停止)")
        while True:
            # 清理已结束的任务进程（检查PID是否存活）
            active_pids = []
            for pid in task_pids:
                if psutil.pid_exists(pid):
                    print(f"任务进程PID {pid} 正在执行")
                    # 进程仍在运行
                    active_pids.append(pid)
                else:
                    print(f"任务进程PID {pid} 已结束")
            task_pids = active_pids  # 更新存活的PID列表

            # 检查是否可以启动新任务
            if len(task_pids) < max_concurrent:
                # 从Redis列表左侧获取一个任务
                task = r.rpop(redis_list_key)

                if task:
                    # print(f"获取到新任务: {json.loads(task)}")
                    task = r.lpush(redis_list_key_task,task)
                    try:
                        # 启动任务并捕获实际进程的PID
                        # 命令说明：
                        # 1. 用nohup启动任务，将PID写入临时文件
                        # 2. 从临时文件读取PID并记录
                        pid_file = f"/tmp/video_task_{int(time.time())}.pid"
                        cmd = (
                            f'nohup python3.9 ./image_gen_process.py'
                            f'> /dev/null 2>&1 & echo $! > {pid_file}'
                        )
                        # 执行命令（启动任务并写入PID）
                        process = subprocess.Popen(
                            cmd,
                            shell=True,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE,
                            text=True
                        )
                        # 等待命令执行完成（确保PID文件已生成）
                        stdout, stderr = process.communicate()
                        if stderr:
                            raise Exception(f"命令执行错误: {stderr}")

                        # 读取PID文件获取实际任务的PID
                        if os.path.exists(pid_file):
                            with open(pid_file, 'r') as f:
                                pid = int(f.read().strip())
                            os.remove(pid_file)  # 清理临时文件
                            task_pids.append(pid)
                            print(f"启动任务进程（PID: {pid}），当前并发数: {len(task_pids)}/{max_concurrent}")
                        else:
                            raise Exception("未生成PID文件，无法获取任务进程ID")

                    except Exception as e:
                        print(f"启动任务失败: {e}")
                        r.rpush(redis_list_key, task)  # 失败时放回队列
                        time.sleep(5)
                else:
                    # 没有任务时短暂休息
                    time.sleep(5)
            else:
                # 达到最大并发数，等待10秒
                print(f"已达到最大并发数 {max_concurrent}，等待10秒...")
                time.sleep(10)

    except KeyboardInterrupt:
        print("\n收到停止信号，正在等待所有任务结束...")
        # 等待所有任务进程结束（可选：如果需要强制终止，可调用psutil.Process(pid).terminate()）
        for pid in task_pids:
            if psutil.pid_exists(pid):
                psutil.Process(pid).wait()  # 等待进程自然结束
                print(f"任务进程PID {pid} 已终止")
        print("所有任务已处理完毕，程序退出")


if __name__ == "__main__":
    main()