"""
任务队列管理模块

提供基于优先级的异步任务队列管理。
"""

import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime
from enum import IntEnum
from typing import Optional, Callable, Dict, Any

from prometheus_client import Gauge, Counter, Histogram

logger = logging.getLogger(__name__)


class TaskPriority(IntEnum):
    """任务优先级"""
    LOW = 0
    NORMAL = 1
    HIGH = 2
    URGENT = 3


@dataclass(order=True)
class TaskQueueItem:
    """
    任务队列项

    支持优先级排序，优先级高的在前，相同优先级按创建时间排序。
    """
    priority: int = field(compare=True)
    created_at: float = field(compare=True)
    task_id: str = field(compare=False)
    task_type: str = field(compare=False)
    user_id: Optional[str] = field(default=None, compare=False)
    project_id: Optional[str] = field(default=None, compare=False)

    def __init__(
        self,
        task_id: str,
        task_type: str,
        priority: TaskPriority = TaskPriority.NORMAL,
        user_id: Optional[str] = None,
        project_id: Optional[str] = None
    ):
        self.priority = priority
        self.created_at = datetime.now().timestamp()
        self.task_id = task_id
        self.task_type = task_type
        self.user_id = user_id
        self.project_id = project_id

    def __lt__(self, other):
        """优先级比较（用于优先队列）"""
        if self.priority != other.priority:
            return self.priority > other.priority  # 高优先级在前
        return self.created_at < other.created_at  # 相同优先级，先创建的在前


# Prometheus 指标
QUEUE_SIZE_GAUGE = Gauge(
    'task_queue_size',
    'Current number of tasks in queue',
    ['task_type']
)

TASKS_ENQUEUED_COUNTER = Counter(
    'tasks_enqueued_total',
    'Total number of tasks enqueued',
    ['task_type', 'priority']
)

TASKS_DEQUEUED_COUNTER = Counter(
    'tasks_dequeued_total',
    'Total number of tasks dequeued',
    ['task_type']
)

WAIT_TIME_HISTOGRAM = Histogram(
    'task_wait_time_seconds',
    'Time spent waiting in queue',
    ['task_type'],
    buckets=[0.1, 0.5, 1, 5, 10, 30, 60, 120, 300]
)

QUEUE_FULL_COUNTER = Counter(
    'task_queue_full_total',
    'Number of times queue was full'
)


class TaskQueue:
    """
    任务队列管理器

    管理基于优先级的任务队列，支持多种任务类型。
    """

    def __init__(self, maxsize: int = 1000):
        """
        初始化任务队列

        Args:
            maxsize: 队列最大容量
        """
        self._queue: asyncio.PriorityQueue[TaskQueueItem] = asyncio.PriorityQueue(maxsize=maxsize)
        self._maxsize = maxsize
        self._item_count = 0

        # 追踪队列中的任务（用于快速查找）
        self._task_registry: Dict[str, TaskQueueItem] = {}
        self._registry_lock = asyncio.Lock()

        # 队列统计
        self._stats = {
            "total_enqueued": 0,
            "total_dequeued": 0,
            "rejected": 0
        }

        logger.info(f"TaskQueue initialized with maxsize={maxsize}")

    @property
    def size(self) -> int:
        """当前队列大小"""
        return self._queue.qsize()

    @property
    def is_full(self) -> bool:
        """队列是否已满"""
        return self._queue.full()

    @property
    def is_empty(self) -> bool:
        """队列是否为空"""
        return self._queue.empty()

    @property
    def maxsize(self) -> int:
        """队列最大容量"""
        return self._maxsize

    async def enqueue(
        self,
        task_id: str,
        task_type: str,
        priority: TaskPriority = TaskPriority.NORMAL,
        user_id: Optional[str] = None,
        project_id: Optional[str] = None,
        timeout: Optional[float] = None
    ) -> bool:
        """
        将任务加入队列

        Args:
            task_id: 任务 ID
            task_type: 任务类型
            priority: 任务优先级
            user_id: 用户 ID
            project_id: 项目 ID
            timeout: 等待队列有空位的超时时间

        Returns:
            是否成功加入队列

        Raises:
            asyncio.TimeoutError: 如果超时
        """
        item = TaskQueueItem(
            task_id=task_id,
            task_type=task_type,
            priority=priority,
            user_id=user_id,
            project_id=project_id
        )

        # 检查队列是否已满
        if self._queue.full():
            if timeout is None:
                QUEUE_FULL_COUNTER.inc()
                self._stats["rejected"] += 1
                logger.warning(f"Queue full, task {task_id} rejected")
                return False

            # 等待队列有空位
            try:
                await asyncio.wait_for(self._wait_for_space(), timeout=timeout)
            except asyncio.TimeoutError:
                QUEUE_FULL_COUNTER.inc()
                self._stats["rejected"] += 1
                logger.warning(f"Queue full, task {task_id} timed out waiting")
                raise

        # 加入队列
        await self._queue.put(item)

        # 更新注册表
        async with self._registry_lock:
            self._task_registry[task_id] = item

        # 更新统计
        self._item_count += 1
        self._stats["total_enqueued"] += 1

        # 更新指标
        TASKS_ENQUEUED_COUNTER.labels(
            task_type=task_type,
            priority=priority.name
        ).inc()
        QUEUE_SIZE_GAUGE.labels(task_type=task_type).set(self._queue.qsize())

        logger.debug(
            f"Task enqueued: {task_id} (type={task_type}, priority={priority.name}, "
            f"queue_size={self._queue.qsize()})"
        )

        return True

    async def dequeue(self, timeout: Optional[float] = None) -> Optional[TaskQueueItem]:
        """
        从队列取出任务

        Args:
            timeout: 等待的超时时间

        Returns:
            TaskQueueItem 或 None（如果超时）
        """
        try:
            if timeout:
                item = await asyncio.wait_for(self._queue.get(), timeout=timeout)
            else:
                item = await self._queue.get()
        except asyncio.TimeoutError:
            return None

        # 计算等待时间
        wait_time = datetime.now().timestamp() - item.created_at

        # 更新注册表
        async with self._registry_lock:
            self._task_registry.pop(item.task_id, None)

        # 更新统计
        self._stats["total_dequeued"] += 1

        # 更新指标
        TASKS_DEQUEUED_COUNTER.labels(task_type=item.task_type).inc()
        WAIT_TIME_HISTOGRAM.labels(task_type=item.task_type).observe(wait_time)
        QUEUE_SIZE_GAUGE.labels(task_type=item.task_type).set(self._queue.qsize())

        logger.debug(
            f"Task dequeued: {item.task_id} (type={item.task_type}, "
            f"waited={wait_time:.2f}s, queue_size={self._queue.qsize()})"
        )

        return item

    async def get_position(self, task_id: str) -> Optional[int]:
        """
        获取任务在队列中的位置

        Args:
            task_id: 任务 ID

        Returns:
            队列位置（1-based），如果不在队列中返回 None
        """
        async with self._registry_lock:
            if task_id not in self._task_registry:
                return None

        # 需要遍历队列查找位置（性能注意：O(n)）
        # 在生产环境中，可能需要更高效的数据结构
        position = 0
        temp_items = []

        try:
            while not self._queue.empty():
                item = await asyncio.wait_for(self._queue.get(), timeout=0.1)
                position += 1
                temp_items.append(item)

                if item.task_id == task_id:
                    # 找到任务，把取出的任务放回
                    for temp_item in temp_items:
                        await self._queue.put(temp_item)
                    return position
        except asyncio.TimeoutError:
            pass

        # 没找到，放回所有任务
        for temp_item in temp_items:
            await self._queue.put(temp_item)

        return None

    def get_stats(self) -> Dict[str, Any]:
        """
        获取队列统计信息

        Returns:
            统计信息字典
        """
        return {
            **self._stats,
            "current_size": self._queue.qsize(),
            "max_size": self._maxsize,
            "utilization": self._queue.qsize() / self._maxsize if self._maxsize > 0 else 0
        }

    async def wait_for_empty(self, check_interval: float = 0.5):
        """
        等待队列变为空

        Args:
            check_interval: 检查间隔（秒）
        """
        while not self._queue.empty():
            await asyncio.sleep(check_interval)

    async def clear(self) -> int:
        """
        清空队列

        Returns:
            清空的任务数量
        """
        count = 0
        async with self._registry_lock:
            while not self._queue.empty():
                try:
                    await self._queue.get()
                    count += 1
                except asyncio.QueueEmpty:
                    break
            self._task_registry.clear()

        logger.info(f"Queue cleared, removed {count} tasks")
        return count

    async def _wait_for_space(self):
        """等待队列有空位"""
        while self._queue.full():
            await asyncio.sleep(0.1)

    async def get_tasks_by_type(self, task_type: str) -> list:
        """
        获取特定类型的所有任务

        Args:
            task_type: 任务类型

        Returns:
            任务 ID 列表
        """
        async with self._registry_lock:
            return [
                task_id for task_id, item in self._task_registry.items()
                if item.task_type == task_type
            ]

    async def get_tasks_by_user(self, user_id: str) -> list:
        """
        获取特定用户的所有任务

        Args:
            user_id: 用户 ID

        Returns:
            任务 ID 列表
        """
        async with self._registry_lock:
            return [
                task_id for task_id, item in self._task_registry.items()
                if item.user_id == user_id
            ]

    async def get_tasks_by_project(self, project_id: str) -> list:
        """
        获取特定项目的所有任务

        Args:
            project_id: 项目 ID

        Returns:
            任务 ID 列表
        """
        async with self._registry_lock:
            return [
                task_id for task_id, item in self._task_registry.items()
                if item.project_id == project_id
            ]
