import asyncio
import weakref
from typing import Dict, Set, Optional
from utils.logger import session_logger


class TaskManager:
    """管理异步任务，防止内存泄露"""

    def __init__(self):
        self.device_tasks: Dict[str, Set[asyncio.Task]] = {}
        self.session_tasks: Dict[tuple, Set[asyncio.Task]] = {}
        self.cleanup_tasks: Set[asyncio.Task] = set()
        self.lock = asyncio.Lock()

    async def create_task(self, coro, device_id: Optional[str] = None, session_key: Optional[tuple] = None, task_type: str = "general"):
        """创建并跟踪异步任务"""
        task = asyncio.create_task(coro)

        async with self.lock:
            # 按设备ID跟踪
            if device_id:
                if device_id not in self.device_tasks:
                    self.device_tasks[device_id] = set()
                self.device_tasks[device_id].add(task)

            # 按会话跟踪
            if session_key:
                if session_key not in self.session_tasks:
                    self.session_tasks[session_key] = set()
                self.session_tasks[session_key].add(task)

            # 跟踪清理任务
            if task_type == "cleanup":
                self.cleanup_tasks.add(task)

        # 任务完成后自动清理
        task.add_done_callback(lambda t: asyncio.create_task(self._cleanup_completed_task(t, device_id, session_key, task_type)))

        session_logger.info(
            device_id or "system",
            session_key[1] if session_key else "task_manager",
            f"创建任务: {task_type}, 总任务数: {len(self.cleanup_tasks) if task_type == 'cleanup' else 'tracked'}"
        )

        return task

    async def _cleanup_completed_task(self, task: asyncio.Task, device_id: Optional[str], session_key: Optional[tuple], task_type: str):
        """清理已完成的任务"""
        async with self.lock:
            if device_id and device_id in self.device_tasks:
                self.device_tasks[device_id].discard(task)
                if not self.device_tasks[device_id]:
                    del self.device_tasks[device_id]

            if session_key and session_key in self.session_tasks:
                self.session_tasks[session_key].discard(task)
                if not self.session_tasks[session_key]:
                    del self.session_tasks[session_key]

            if task_type == "cleanup":
                self.cleanup_tasks.discard(task)

    async def cancel_device_tasks(self, device_id: str):
        """取消设备的所有任务"""
        async with self.lock:
            if device_id in self.device_tasks:
                tasks = list(self.device_tasks[device_id])
                session_logger.info(device_id, "task_manager", f"取消设备任务数量: {len(tasks)}")

                for task in tasks:
                    if not task.done():
                        task.cancel()

                # 等待任务取消完成
                if tasks:
                    try:
                        await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=5.0)
                    except asyncio.TimeoutError:
                        session_logger.warning(device_id, "task_manager", "部分任务取消超时")

                del self.device_tasks[device_id]

    async def cancel_session_tasks(self, session_key: tuple):
        """取消会话的所有任务"""
        async with self.lock:
            if session_key in self.session_tasks:
                tasks = list(self.session_tasks[session_key])
                device_id, session_id = session_key
                session_logger.info(device_id, session_id, f"取消会话任务数量: {len(tasks)}")

                for task in tasks:
                    if not task.done():
                        task.cancel()

                # 等待任务取消完成
                if tasks:
                    try:
                        await asyncio.wait_for(asyncio.gather(*tasks, return_exceptions=True), timeout=3.0)
                    except asyncio.TimeoutError:
                        session_logger.warning(device_id, session_id, "部分会话任务取消超时")

                del self.session_tasks[session_key]

    async def get_task_stats(self) -> dict:
        """获取任务统计信息"""
        async with self.lock:
            device_count = sum(len(tasks) for tasks in self.device_tasks.values())
            session_count = sum(len(tasks) for tasks in self.session_tasks.values())

            return {
                "device_tasks": device_count,
                "session_tasks": session_count,
                "cleanup_tasks": len(self.cleanup_tasks),
                "total_tracked_devices": len(self.device_tasks),
                "total_tracked_sessions": len(self.session_tasks)
            }

    async def cleanup_all_tasks(self):
        """清理所有任务"""
        async with self.lock:
            all_tasks = []

            # 收集所有任务
            for tasks in self.device_tasks.values():
                all_tasks.extend(tasks)
            for tasks in self.session_tasks.values():
                all_tasks.extend(tasks)
            all_tasks.extend(self.cleanup_tasks)

            session_logger.info("system", "task_manager", f"开始清理所有任务，总数: {len(all_tasks)}")

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

            # 等待任务取消完成
            if all_tasks:
                try:
                    await asyncio.wait_for(asyncio.gather(*all_tasks, return_exceptions=True), timeout=10.0)
                except asyncio.TimeoutError:
                    session_logger.warning("system", "task_manager", "部分任务清理超时")

            # 清空所有跟踪
            self.device_tasks.clear()
            self.session_tasks.clear()
            self.cleanup_tasks.clear()

            session_logger.info("system", "task_manager", "任务管理器清理完成")


# 全局任务管理器实例
task_manager = TaskManager()