"""
任务队列管理器
支持优先级和并发控制
"""
import asyncio
import uuid
from typing import Optional, Dict, Any, Callable, Awaitable, TypeVar, List
from datetime import datetime
from loguru import logger

from .types import TaskStatus, TaskQueueConfig, QueueTask

T = TypeVar('T')


class TaskQueue:
    """任务队列管理器，支持优先级和并发控制"""

    def __init__(self, config: Optional[TaskQueueConfig] = None):
        self.config = config or TaskQueueConfig()
        self._queue: asyncio.PriorityQueue = asyncio.PriorityQueue()
        self._tasks: Dict[str, QueueTask] = {}
        self._running_tasks: Dict[str, asyncio.Task] = {}
        self._results: Dict[str, Any] = {}
        self._is_running = False
        self._workers: List[asyncio.Task] = []
        self._semaphore: Optional[asyncio.Semaphore] = None

    async def start(self):
        if self._is_running:
            return
        self._is_running = True
        self._semaphore = asyncio.Semaphore(self.config.max_concurrent)
        for i in range(self.config.max_concurrent):
            worker = asyncio.create_task(self._worker(i))
            self._workers.append(worker)

    async def stop(self, wait: bool = True):
        self._is_running = False
        if wait:
            while self._running_tasks:
                await asyncio.sleep(0.1)
        for worker in self._workers:
            worker.cancel()
        self._workers.clear()

    async def __aenter__(self):
        await self.start()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.stop()

    def submit(
        self,
        handler: Callable[..., Awaitable[T]],
        name: Optional[str] = None,
        priority: int = 5,
        args: Optional[tuple] = None,
        kwargs: Optional[Dict] = None,
        dependencies: Optional[List[str]] = None,
        max_retries: int = 3
    ) -> str:
        task_id = str(uuid.uuid4())[:8]
        task = QueueTask(
            id=task_id, name=name or handler.__name__, handler=handler,
            priority=priority, args=args or (), kwargs=kwargs or {},
            dependencies=dependencies or [], max_retries=max_retries,
        )
        self._tasks[task_id] = task
        self._queue.put_nowait((-priority, task.created_at.timestamp(), task_id))
        return task_id

    async def wait_for(self, task_id: str, timeout: Optional[float] = None) -> Any:
        timeout = timeout or self.config.timeout
        start_time = datetime.now()
        while True:
            task = self._tasks.get(task_id)
            if not task:
                raise ValueError(f"任务不存在: {task_id}")
            if task.status == TaskStatus.COMPLETED:
                return self._results.get(task_id)
            if task.status == TaskStatus.FAILED:
                raise RuntimeError(f"任务失败: {task.error}")
            if task.status == TaskStatus.CANCELLED:
                raise asyncio.CancelledError("任务已取消")
            if (datetime.now() - start_time).total_seconds() > timeout:
                raise asyncio.TimeoutError("等待任务超时")
            await asyncio.sleep(0.1)

    def get_task(self, task_id: str) -> Optional[QueueTask]:
        return self._tasks.get(task_id)

    def get_stats(self) -> Dict[str, Any]:
        stats = {"total": len(self._tasks)}
        for status in TaskStatus:
            stats[status.value] = sum(1 for t in self._tasks.values() if t.status == status)
        return stats

    async def _worker(self, worker_id: int):
        while self._is_running:
            try:
                try:
                    _, _, task_id = await asyncio.wait_for(self._queue.get(), timeout=1.0)
                except asyncio.TimeoutError:
                    continue
                task = self._tasks.get(task_id)
                if not task or task.status == TaskStatus.CANCELLED:
                    continue
                async with self._semaphore:
                    await self._execute_task(task)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"工作协程 {worker_id} 错误: {e}")

    async def _execute_task(self, task: QueueTask):
        task.status = TaskStatus.RUNNING
        task.started_at = datetime.now()
        try:
            result = await asyncio.wait_for(
                task.handler(*task.args, **task.kwargs),
                timeout=self.config.timeout
            )
            task.result = result
            task.status = TaskStatus.COMPLETED
            self._results[task.id] = result
        except asyncio.TimeoutError:
            task.error = "任务超时"
            task.status = TaskStatus.FAILED
        except Exception as e:
            task.error = str(e)
            if task.retry_count < task.max_retries:
                task.retry_count += 1
                task.status = TaskStatus.PENDING
                await asyncio.sleep(self.config.retry_delay)
                self._queue.put_nowait((-task.priority, task.created_at.timestamp(), task.id))
            else:
                task.status = TaskStatus.FAILED
        finally:
            task.completed_at = datetime.now()

    def cancel(self, task_id: str) -> bool:
        """取消任务"""
        task = self._tasks.get(task_id)
        if task and task.status == TaskStatus.PENDING:
            task.status = TaskStatus.CANCELLED
            return True
        return False

    def get_pending_count(self) -> int:
        """获取待处理任务数"""
        return sum(1 for t in self._tasks.values() if t.status == TaskStatus.PENDING)

    def get_running_count(self) -> int:
        """获取运行中任务数"""
        return sum(1 for t in self._tasks.values() if t.status == TaskStatus.RUNNING)
