# main.py
# 主程序入口
import os
import signal
import threading
import time
import schedule
import torch
from log import logger
from global_data import gd
from server import socketio, app
from task_manager import task_manager

# 初始化全局数据
gd._init()

# 设置全局变量
gd.set_value('task_manager', task_manager)  # 直接在这里初始化
gd.set_value('latest_drone_data', None)
gd.set_value('mqtt_data_lock', threading.Lock())

# 初始化任务管理器
logger.info("任务管理器初始化...")
logger.info(f"任务管理器实例: {task_manager}")


def cleanup_stopped_tasks():
    """定期清理已停止的任务"""
    try:
        if not task_manager:
            return

        all_tasks = task_manager.get_all_tasks()
        cleaned_count = 0

        for task in all_tasks:
            task_id = task.get('task_id')
            status = task.get('status', 'unknown')

            # 清理已停止、失败或错误的任务
            if status in ['stopped', 'failed', 'error']:
                # 检查任务线程是否真的已经停止
                thread = task_manager.tasks.get(task_id, {}).get('thread')
                if thread and hasattr(thread, 'is_alive'):
                    if not thread.is_alive():
                        if task_manager.cleanup_task(task_id):
                            cleaned_count += 1
                else:
                    if task_manager.cleanup_task(task_id):
                        cleaned_count += 1

        if cleaned_count > 0:
            logger.info(f"定时清理已完成，清理了 {cleaned_count} 个任务")

    except Exception as e:
        logger.error(f"定时清理任务失败: {str(e)}")


def start_cleanup_scheduler():
    """启动定时清理调度器"""
    # 每5分钟清理一次
    schedule.every(5).minutes.do(cleanup_stopped_tasks)

    # 每分钟检查一次待清理任务
    schedule.every(1).minutes.do(check_pending_cleanup)

    def run_scheduler():
        while True:
            try:
                schedule.run_pending()
            except Exception as e:
                logger.error(f"定时任务执行失败: {str(e)}")
            time.sleep(60)  # 每分钟检查一次

    scheduler_thread = threading.Thread(target=run_scheduler, daemon=True)
    scheduler_thread.start()
    logger.info("定时清理调度器已启动")


def check_pending_cleanup():
    """检查待清理的任务"""
    try:
        if not task_manager:
            return

        # 检查是否有长时间处于停止状态但未清理的任务
        for task_id, task_info in task_manager.tasks.items():
            status = task_info.get('status', 'unknown')
            last_updated = task_info.get('last_updated', 0)

            if status == 'stopped' and time.time() - last_updated > 60:  # 停止超过60秒
                logger.info(f"发现长时间停止未清理的任务: {task_id}")
                task_manager.cleanup_task(task_id)

    except Exception as e:
        logger.error(f"检查待清理任务失败: {str(e)}")


# main.py 修改部分

if __name__ == '__main__':
    logger.info("启动多任务版YOLOv8服务")
    logger.info(f"PyTorch版本: {torch.__version__}, CUDA可用: {torch.cuda.is_available()}")

    # 启动定时清理
    start_cleanup_scheduler()

    # 初始化任务推流管理器
    try:
        from task_stream_manager import task_stream_manager

        task_stream_manager.start_health_monitor()
        logger.info("任务推流管理器初始化完成")
    except Exception as e:
        logger.error(f"初始化推流管理器失败: {str(e)}")


    # 退出服务
    def graceful_exit(signum, frame):
        logger.info("收到退出信号，停止所有服务...")

        # 停止所有任务推流
        try:
            from task_stream_manager import task_stream_manager
            task_stream_manager.cleanup_all()
        except:
            pass

        # 停止所有任务
        if task_manager:
            cleaned_count = task_manager.cleanup_all_tasks()
            logger.info(f"清理了 {cleaned_count} 个任务")

        logger.info("服务已退出")
        os._exit(0)


    signal.signal(signal.SIGINT, graceful_exit)
    signal.signal(signal.SIGTERM, graceful_exit)

    # 启动服务
    socketio.run(app,
                 host='0.0.0.0',
                 port=9309,
                 debug=True,
                 use_reloader=False,
                 allow_unsafe_werkzeug=True)