buff/backend/app/background_tasks.py

"""
后台定时任务管理
"""
import asyncio
from typing import List, Optional
from contextlib import asynccontextmanager
from app.database import get_db
from app.core.logger import logger
from app.services.key_rotation_service import KeyRotationService
from app.services.session_service import SessionService
from app.services.heartbeat_monitor import HeartbeatMonitor
from app.config import settings


class BackgroundTaskManager:
    """后台任务管理器"""

    def __init__(self):
        self.tasks: List[asyncio.Task] = []
        self.running = False

    async def start_all_tasks(self):
        """启动所有后台任务"""
        if self.running:
            logger.warning("后台任务已在运行")
            return

        self.running = True
        logger.info("启动后台定时任务...")

        # 启动各个任务
        self.tasks = [
            asyncio.create_task(self._key_rotation_task()),
            asyncio.create_task(self._session_cleanup_task()),
            asyncio.create_task(self._heartbeat_cleanup_task()),
        ]

        logger.info(f"已启动 {len(self.tasks)} 个后台任务")

    async def stop_all_tasks(self):
        """停止所有后台任务"""
        if not self.running:
            logger.warning("后台任务未运行")
            return

        logger.info("停止后台定时任务...")
        self.running = False

        # 取消所有任务
        for task in self.tasks:
            if not task.done():
                task.cancel()

        # 等待任务完成
        try:
            await asyncio.gather(*self.tasks, return_exceptions=True)
        except asyncio.CancelledError:
            pass

        self.tasks.clear()
        logger.info("后台定时任务已停止")

    async def _key_rotation_task(self):
        """密钥轮换任务 - 每5分钟执行一次"""
        logger.info("启动密钥轮换定时任务 (间隔: 5分钟)")

        while self.running:
            try:
                logger.debug("执行密钥轮换检查...")

                db = next(get_db())
                try:
                    # 检查是否需要轮换密钥
                    if KeyRotationService.should_rotate_key(db):
                        logger.info("检测到需要轮换密钥，开始执行...")
                        success, public_key, error_msg = KeyRotationService.rotate_key(db)
                        if success:
                            logger.info("密钥轮换成功")
                        else:
                            logger.error(f"密钥轮换失败: {error_msg}")
                    else:
                        logger.debug("当前无需轮换密钥")

                    # 清理过期密钥
                    cleaned_count = KeyRotationService.cleanup_expired_keys(db)
                    if cleaned_count > 0:
                        logger.info(f"清理了 {cleaned_count} 个过期密钥")

                finally:
                    db.close()

            except Exception as e:
                logger.exception("密钥轮换任务异常")

            # 等待5分钟
            await asyncio.sleep(300)  # 5分钟 = 300秒

    async def _session_cleanup_task(self):
        """Session清理任务 - 每2小时执行一次"""
        logger.info("启动Session清理定时任务 (间隔: 2小时)")

        while self.running:
            try:
                logger.debug("执行Session清理...")

                db = next(get_db())
                try:
                    cleaned_count = SessionService.cleanup_expired_sessions(db)
                    if cleaned_count > 0:
                        logger.info(f"清理了 {cleaned_count} 个30天前的非活跃Session记录")
                    else:
                        logger.debug("没有需要清理的旧Session记录")

                finally:
                    db.close()

            except Exception as e:
                logger.exception("Session清理任务异常")

            # 等待2小时
            await asyncio.sleep(7200)  # 2小时 = 7200秒

    async def _heartbeat_cleanup_task(self):
        """心跳监控缓存清理任务 - 每1小时执行一次"""
        logger.info("启动心跳缓存清理定时任务 (间隔: 1小时)")

        while self.running:
            try:
                logger.debug("执行心跳缓存清理...")

                cleaned_count = HeartbeatMonitor.cleanup_cache()
                if cleaned_count > 0:
                    logger.info(f"清理了 {cleaned_count} 个过期的缓存条目")
                else:
                    logger.debug("心跳缓存状态正常，无需清理")

                # 输出监控统计信息
                stats = HeartbeatMonitor.get_monitoring_stats()
                logger.info(f"心跳监控统计: 监控用户={stats.get('total_monitored_users', 0)}, "
                           f"受保护用户={stats.get('protected_users', 0)}, "
                           f"失败用户={stats.get('failing_users', 0)}, "
                           f"缓存大小={stats.get('cache_size', 0)}")

            except Exception as e:
                logger.exception("心跳缓存清理任务异常")

            # 等待1小时
            await asyncio.sleep(3600)  # 1小时 = 3600秒


# 全局后台任务管理器实例
background_task_manager = BackgroundTaskManager()


@asynccontextmanager
async def lifespan_manager(app):
    """FastAPI lifespan 事件管理器"""
    # 启动时执行
    logger.info("应用启动，开始初始化后台任务...")
    await background_task_manager.start_all_tasks()

    yield

    # 关闭时执行
    logger.info("应用关闭，开始清理后台任务...")
    await background_task_manager.stop_all_tasks()


# 便捷函数
async def start_background_tasks():
    """启动后台任务（手动调用）"""
    await background_task_manager.start_all_tasks()


async def stop_background_tasks():
    """停止后台任务（手动调用）"""
    await background_task_manager.stop_all_tasks()


def get_task_manager() -> BackgroundTaskManager:
    """获取后台任务管理器实例"""
    return background_task_manager