"""
限流器管理器

统一管理所有任务的限流器实例

职责（简化后）：
1. 管理 TaskRateLimiter 实例
2. 每次 acquire 时从 DynamicConfig 获取最新配置
3. 配置变化时自动重建限流器

不再负责：
- 配置存储和读取（由 DynamicConfig 统一管理）
- 配置变更监听（DynamicConfig 已处理 Pub/Sub）
"""

import asyncio
import logging
from redis.asyncio import Redis
from typing import Dict, Optional, Any, TYPE_CHECKING

from .task_limiter import TaskRateLimiter
from .config import RateLimitConfig, QPSLimit, ConcurrencyLimit

if TYPE_CHECKING:
    from jettask.config.dynamic_config import DynamicConfig

logger = logging.getLogger('app')


class RateLimiterManager:
    """限流器管理器

    管理多个任务的限流器。
    每次 acquire 时从 DynamicConfig 获取最新配置，自动检测配置变化并重建限流器。
    """

    def __init__(
        self,
        redis_client: Redis,
        worker_id: str,
        redis_prefix: str = "jettask",
        dynamic_config: 'DynamicConfig' = None,
        worker_state_manager=None
    ):
        self._redis = redis_client
        self._worker_id = worker_id
        self._prefix = redis_prefix
        self._dynamic_config = dynamic_config
        self._worker_state_manager = worker_state_manager

        self._limiters: Dict[str, TaskRateLimiter] = {}

        self._config_snapshots: Dict[str, Optional[Dict]] = {}

        self._rebuild_locks: Dict[str, asyncio.Lock] = {}

        logger.debug(f"RateLimiterManager initialized for worker {worker_id}")

    def set_dynamic_config(self, dynamic_config: 'DynamicConfig'):
        self._dynamic_config = dynamic_config
        logger.debug("DynamicConfig injected into RateLimiterManager")

    async def preload_all(self):
        if not self._dynamic_config:
            logger.warning("[RateLimiter] DynamicConfig 未设置，跳过预加载")
            return

        task_names = self._dynamic_config.get_all_task_names()
        if not task_names:
            logger.debug("[RateLimiter] 没有已配置的任务，跳过预加载")
            return

        loaded_count = 0
        for task_name in task_names:
            config = self._dynamic_config.get_rate_limit(task_name)
            if config:
                try:
                    await self._rebuild_limiter(task_name, config)
                    loaded_count += 1
                except Exception as e:
                    logger.error(f"[RateLimiter] 预加载限流器失败: {task_name}, error: {e}")

        logger.info(f"[RateLimiter] 预加载完成，已加载 {loaded_count} 个限流器")

    def _get_rebuild_lock(self, task_name: str) -> asyncio.Lock:
        if task_name not in self._rebuild_locks:
            self._rebuild_locks[task_name] = asyncio.Lock()
        return self._rebuild_locks[task_name]

    async def acquire(self, task_name: str, timeout: float = None) -> Optional[str]:
        config = self._dynamic_config.get_rate_limit(task_name) if self._dynamic_config else None

        if self._need_rebuild(task_name, config):
            lock = self._get_rebuild_lock(task_name)
            async with lock:
                if self._need_rebuild(task_name, config):
                    await self._rebuild_limiter(task_name, config)

        limiter = self._limiters.get(task_name)
        if limiter:
            return await limiter.acquire(timeout=timeout)

        return True

    async def release(self, task_name: str, task_id: Optional[str] = None):
        limiter = self._limiters.get(task_name)
        if limiter:
            await limiter.release(task_id)

    def _need_rebuild(self, task_name: str, new_config: Optional[Dict]) -> bool:
        old_config = self._config_snapshots.get(task_name)

        if old_config is None and new_config is not None:
            return True

        if old_config is not None and new_config is None:
            return True

        if old_config != new_config:
            return True

        if new_config is not None and task_name not in self._limiters:
            return True

        return False

    async def _rebuild_limiter(self, task_name: str, config: Optional[Dict]):
        old_limiter = self._limiters.pop(task_name, None)
        if old_limiter:
            try:
                await old_limiter.stop()
            except Exception as e:
                logger.warning(f"[RateLimiter] 停止旧限流器失败: {task_name}, error: {e}")

        self._config_snapshots[task_name] = config

        if config is None:
            logger.info(f"[RateLimiter] 移除限流器: {task_name}")
            return

        rate_limit_config = self._parse_config(config)
        if rate_limit_config:
            limiter = TaskRateLimiter(
                redis_client=self._redis,
                task_name=task_name,
                worker_id=self._worker_id,
                config=rate_limit_config,
                redis_prefix=self._prefix,
                worker_state_manager=self._worker_state_manager
            )
            await limiter.start()
            self._limiters[task_name] = limiter
            logger.info(f"[RateLimiter] 重建限流器: {task_name} = {config}")

    def _parse_config(self, config: Dict) -> Optional[RateLimitConfig]:
        rate_type = config.get("type")
        if rate_type == "qps":
            return QPSLimit(
                qps=config.get("qps", 10),
                window_size=config.get("window_size", 1.0)
            )
        elif rate_type == "concurrency":
            return ConcurrencyLimit(
                max_concurrency=config.get("max_concurrency", 10)
            )
        return None

    async def stop_all(self):
        for task_name, limiter in list(self._limiters.items()):
            try:
                await limiter.stop()
            except Exception as e:
                logger.warning(f"[RateLimiter] 停止限流器失败: {task_name}, error: {e}")
        self._limiters.clear()
        self._config_snapshots.clear()
        logger.debug("Stopped all rate limiters")

    def get_all_stats(self) -> Dict[str, dict]:
        stats = {}
        for task_name, limiter in self._limiters.items():
            if hasattr(limiter, 'get_stats'):
                stats[task_name] = limiter.get_stats()
        return stats


    @staticmethod
    def register_rate_limit_config(
        redis_client,
        task_name: str,
        config: RateLimitConfig,
        redis_prefix: str = "jettask",
        force: bool = False
    ):
        try:
            config_key = f"{redis_prefix}:CONFIG:TASK:{task_name}"
            index_key = f"{redis_prefix}:CONFIG:TASK:INDEX"
            version_key = f"{redis_prefix}:CONFIG:VERSION"

            if not force:
                existing_type = redis_client.hget(config_key, "rate_limit_type")
                if existing_type:
                    logger.debug(
                        f"任务 '{task_name}' 已有限流配置（可能来自 WebUI），跳过代码中的默认配置"
                    )
                    return

            config_dict = config.to_dict()
            flat_config = {}
            rate_type = config_dict.get("type")
            flat_config["rate_limit_type"] = rate_type

            if rate_type == "qps":
                flat_config["rate_limit_qps"] = str(config_dict.get("qps", 10))
                flat_config["rate_limit_window_size"] = str(config_dict.get("window_size", 1.0))
            elif rate_type == "concurrency":
                flat_config["rate_limit_max_concurrency"] = str(config_dict.get("max_concurrency", 10))

            redis_client.hset(config_key, mapping=flat_config)
            redis_client.sadd(index_key, task_name)
            redis_client.incr(version_key)

            logger.debug(f"Registered rate limit config for task '{task_name}': {config}")
        except Exception as e:
            logger.error(f"Failed to register rate limit config for task '{task_name}': {e}")

    @staticmethod
    def unregister_rate_limit_config(
        redis_client,
        task_name: str,
        redis_prefix: str = "jettask"
    ):
        try:
            config_key = f"{redis_prefix}:CONFIG:TASK:{task_name}"
            index_key = f"{redis_prefix}:CONFIG:TASK:INDEX"
            version_key = f"{redis_prefix}:CONFIG:VERSION"

            redis_client.hdel(config_key, "rate_limit_type", "rate_limit_qps",
                             "rate_limit_window_size", "rate_limit_max_concurrency")

            remaining = redis_client.hlen(config_key)
            if remaining == 0:
                redis_client.delete(config_key)
                redis_client.srem(index_key, task_name)

            redis_client.incr(version_key)

            logger.debug(f"Removed rate limit config for task '{task_name}'")
            return True
        except Exception as e:
            logger.error(f"Failed to remove rate limit config for task '{task_name}': {e}")
            return False


__all__ = ['RateLimiterManager']
