import asyncio
import logging
from abc import ABC, abstractmethod
from typing import Generic, TypeVar

from pydantic import BaseModel

logger = logging.getLogger(__name__)

_V = TypeVar("_V", bound=BaseModel)


class TaskQueue(Generic[_V], ABC):
    """
    一个基于 asyncio.Queue 实现的内存任务队列。
    """

    def __init__(self, max_workers: int) -> None:
        self.max_workers = max_workers
        self.queue: asyncio.Queue[_V] = asyncio.Queue()
        self.workers: list[asyncio.Task] = []
        self._shutdown = False

    def start(self) -> None:
        """
        启动 worker 协程
        """
        for i in range(0, self.max_workers):
            task = asyncio.create_task(self._worker(i), name=f"worker {i}")
            self.workers.append(task)

    async def put(self, item: _V) -> None:
        """
        向队列提交任务。

        Args:
            item: 任务

        Raises:
            RuntimeError: 在队列关闭的过程中提交新任务。
        """
        if self._shutdown:
            raise RuntimeError("队列正在关闭中，无法提交新任务。")
        await self.queue.put(item)

    async def _worker(self, worker_id: int) -> None:
        """
        消费逻辑。

        Args:
            worker_id: 消费者 ID
        """
        while True:
            try:
                # 使用 timeout 以便优雅地检查 shutdown
                item = await asyncio.wait_for(self.queue.get(), timeout=1.0)
            except asyncio.TimeoutError:
                if self._shutdown:
                    break
                continue

            try:
                logger.info(f"[Worker {worker_id}] Processing: {item}")
                await self._process(item)
            except Exception as e:
                logger.error(f"[Worker {worker_id}] Error processing {item}: {e}")
            finally:
                self.queue.task_done()

    async def join(self) -> None:
        """等待队列中所有任务完成"""
        await self.queue.join()

    async def shutdown(self) -> None:
        """优雅关闭"""
        self._shutdown = True
        await self.join()
        for w in self.workers:
            w.cancel()
        await asyncio.gather(*self.workers, return_exceptions=True)
        logger.info("队列成功关闭。")

    @abstractmethod
    async def _process(self, item: _V) -> None:
        """实际执行的工作。接收 item，返回 None。请 overload 此方法。"""
        ...
