import json
import logging
import threading
from abc import ABC, abstractmethod
from concurrent.futures import ThreadPoolExecutor, wait as futures_wait
from typing import Any, Generic, Optional, TypeVar
from redis import Redis
from .types import RedisClient
logger = logging.getLogger(__name__)
TaskT = TypeVar("TaskT")
class RedisWork(Generic[TaskT], ABC):
    """
    generated by ai(kimi): https://www.kimi.com/chat/19ab06e4-1232-8061-8000-0905a428a309
    监听 Redis list，线程池并发处理任务。

    用法：
    >>> class MyWork(RedisWork[dict]):
    ...     def handle(self, task: dict) -> None:
    ...         print(task)
    >>> w = MyWork("myqueue", max_workers=4)
    >>> w.start()
    """

    def __init__(
        self,
        list_key: str,
        redis_client: RedisClient,
        max_workers: int = 4,
        *,
        timeout: int = 5,
        decode_responses: bool = True,
    ) -> None:
        self.list_key = list_key
        self.timeout = timeout

        # Redis
        self.rdb = redis_client

        # 线程池
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self._max_workers = max_workers

        # 生命周期
        self._stop_evt = threading.Event()
        self._first_exc: Optional[BaseException] = None

        # 条件变量：线程池满载时阻塞
        self._lock = threading.Lock()
        self._cond = threading.Condition(self._lock)
        self._futures: set[Any] = set()

    # ---------------- 子类必须实现 ----------------
    @abstractmethod
    def handle(self, task: TaskT) -> Any:
        raise NotImplementedError

    # ---------------- 公共 API ----------------
    def start(self) -> None:
        """阻塞运行，直到手动 stop 或任务抛错。"""
        logger.info("RedisWork<%s> started, listening key='%s'", self.__class__.__name__, self.list_key)
        try:
            while not self._stop_evt.is_set():
                if self._first_exc is not None:
                    break

                # 1. 线程池满就阻塞
                with self._lock:
                    while len(self._futures) >= self._max_workers:
                        logger.warning("并发达到上限，暂停监听 Redis")
                        self._cond.wait()

                # 2. 取任务
                raw_task = self.rdb.blpop(self.list_key, timeout=self.timeout)
                if raw_task is None:
                    continue

                _, raw = raw_task
                try:
                    task = json.loads(raw)
                except Exception as exc:
                    logger.error("Bad task format: %s", raw, exc_info=exc)
                    continue

                # 3. 提交
                fut = self.executor.submit(self._safe_handle, task)
                with self._lock:
                    self._futures.add(fut)
                fut.add_done_callback(self._on_future_done)

        except KeyboardInterrupt:
            logger.info("Caught Ctrl-C, shutting down…")
        finally:
            self.shutdown()
            if self._first_exc is not None:
                raise self._first_exc

    def shutdown(self, wait: bool = True) -> None:
        self._stop_evt.set()
        logger.info("等待线程池清场…")
        if wait:
            futures_wait(self._futures, timeout=None)
        self.executor.shutdown(wait=wait)
        logger.info("RedisWork<%s> stopped.", self.__class__.__name__)

    # ---------------- 内部 ----------------
    def _safe_handle(self, task: TaskT) -> None:
        try:
            self.handle(task)
        except Exception as exc:
            with self._lock:
                if self._first_exc is None:
                    self._first_exc = exc
                    logger.exception("Task handle failed，将停止监听")
            raise

    def _on_future_done(self, fut: Any) -> None:
        with self._lock:
            self._futures.discard(fut)
            if len(self._futures) < self._max_workers:
                self._cond.notify()
