"""
统一的 Redis Pub/Sub 管理器

提供统一的 Pub/Sub 订阅管理、错误处理和自动重连机制。

注意：
- 连接池级别的配置（socket_timeout, health_check_interval）应使用
  connector.py 的 get_async_redis_pool_for_pubsub()
- 本模块只处理应用层的订阅逻辑、消息回调和断线重连

使用示例:

    # 创建订阅
    subscription = await PubSubManager.subscribe(
        redis_client=redis,
        channel="my_channel",
        on_message=my_callback,
        name="my_listener"
    )

    # 停止订阅
    await subscription.stop()

    # 或使用上下文管理器
    async with PubSubManager.subscribe(...) as subscription:
        await asyncio.sleep(60)  # 监听 60 秒
"""

import asyncio
import json
import logging
from typing import Any, Callable, Optional, Union, Awaitable
from dataclasses import dataclass

import redis.asyncio as redis

logger = logging.getLogger(__name__)


@dataclass
class PubSubConfig:
    """Pub/Sub 配置"""
    initial_retry_delay: float = 1.0  
    max_retry_delay: float = 30.0  
    retry_backoff_multiplier: float = 2.0  

    mark_connection: bool = True


class PubSubSubscription:
    """
    Pub/Sub 订阅实例

    管理单个 channel 的订阅生命周期，包括：
    - 自动重连（应用层，当 listen() 抛异常时重新订阅）
    - 连接标记（防止被空闲连接清理）

    注意：连接池级别的健康检查由 Redis 连接池处理，
    参见 connector.py 的 get_async_redis_pool_for_pubsub()
    """

    def __init__(
        self,
        redis_client: redis.Redis,
        channel: str,
        on_message: Callable[[dict], Union[None, Awaitable[None]]],
        name: str = None,
        config: PubSubConfig = None,
        parse_json: bool = True,
    ):
        self._redis = redis_client
        self._channel = channel
        self._on_message = on_message
        self._name = name or channel
        self._config = config or PubSubConfig()
        self._parse_json = parse_json

        self._pubsub: Optional[redis.client.PubSub] = None
        self._running = False
        self._listener_task: Optional[asyncio.Task] = None

    @property
    def channel(self) -> str:
        return self._channel

    @property
    def is_running(self) -> bool:
        return self._running

    async def start(self):
        if self._running:
            logger.debug(f"[PubSub:{self._name}] Already running")
            return

        self._running = True
        self._pubsub = await self._create_and_subscribe()
        self._listener_task = asyncio.create_task(self._listen_loop())

        logger.debug(f"[PubSub:{self._name}] Started on channel: {self._channel}")

    async def stop(self):
        if not self._running:
            return

        self._running = False

        if self._listener_task:
            self._listener_task.cancel()
            try:
                await self._listener_task
            except asyncio.CancelledError:
                pass
            self._listener_task = None

        if self._pubsub:
            try:
                await self._pubsub.unsubscribe(self._channel)
                await self._pubsub.close()
            except Exception as e:
                logger.debug(f"[PubSub:{self._name}] Error closing pubsub: {e}")
            self._pubsub = None

        logger.debug(f"[PubSub:{self._name}] Stopped")

    async def _create_and_subscribe(self) -> redis.client.PubSub:
        if self._pubsub:
            try:
                await self._pubsub.close()
            except Exception:
                pass

        pubsub = self._redis.pubsub()
        await pubsub.subscribe(self._channel)

        if self._config.mark_connection:
            self._mark_pubsub_connection(pubsub)

        logger.debug(f"[PubSub:{self._name}] Subscribed to channel: {self._channel}")
        return pubsub

    def _mark_pubsub_connection(self, pubsub: redis.client.PubSub):
        if hasattr(pubsub, 'connection') and pubsub.connection:
            pubsub.connection._is_pubsub_connection = True
            logger.debug(
                f"[PubSub:{self._name}] Marked connection {id(pubsub.connection)} "
                "to prevent idle cleanup"
            )

    async def _listen_loop(self):
        retry_delay = self._config.initial_retry_delay

        while self._running:
            try:
                async for message in self._pubsub.listen():
                    if not self._running:
                        break

                    if message['type'] == 'message':
                        await self._handle_message(message)

                retry_delay = self._config.initial_retry_delay

            except asyncio.CancelledError:
                logger.debug(f"[PubSub:{self._name}] Listen loop cancelled")
                break

            except Exception as e:
                if not self._running:
                    break

                error_msg = str(e).lower()
                is_connection_error = any(
                    keyword in error_msg
                    for keyword in ['connection', 'closed', 'buffer', 'timeout', 'eof']
                )

                if is_connection_error:
                    logger.warning(
                        f"[PubSub:{self._name}] Connection lost: {e}, "
                        f"reconnecting in {retry_delay}s..."
                    )
                else:
                    logger.error(f"[PubSub:{self._name}] Error in listen loop: {e}")

                await asyncio.sleep(retry_delay)

                try:
                    self._pubsub = await self._create_and_subscribe()
                    logger.info(f"[PubSub:{self._name}] Reconnected successfully")
                    retry_delay = self._config.initial_retry_delay
                except Exception as reconnect_error:
                    logger.error(
                        f"[PubSub:{self._name}] Reconnect failed: {reconnect_error}"
                    )
                    retry_delay = min(
                        retry_delay * self._config.retry_backoff_multiplier,
                        self._config.max_retry_delay
                    )

        logger.debug(f"[PubSub:{self._name}] Listen loop exited")

    async def _handle_message(self, message: dict):
        try:
            data = message['data']

            if self._parse_json:
                try:
                    data = json.loads(data)
                except (json.JSONDecodeError, TypeError):
                    pass

            if asyncio.iscoroutinefunction(self._on_message):
                await self._on_message(data)
            else:
                self._on_message(data)

        except Exception as e:
            logger.error(f"[PubSub:{self._name}] Error handling message: {e}")

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

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


class PubSubManager:
    """
    Pub/Sub 管理器

    提供统一的订阅创建和管理接口
    """

    @staticmethod
    async def subscribe(
        redis_client: redis.Redis,
        channel: str,
        on_message: Callable[[dict], Union[None, Awaitable[None]]],
        name: str = None,
        config: PubSubConfig = None,
        parse_json: bool = True,
        auto_start: bool = True,
    ) -> PubSubSubscription:
        subscription = PubSubSubscription(
            redis_client=redis_client,
            channel=channel,
            on_message=on_message,
            name=name,
            config=config,
            parse_json=parse_json,
        )

        if auto_start:
            await subscription.start()

        return subscription

    @staticmethod
    async def publish(
        redis_client: redis.Redis,
        channel: str,
        message: Any,
        serialize_json: bool = True,
    ) -> int:
        if serialize_json and not isinstance(message, str):
            message = json.dumps(message)

        return await redis_client.publish(channel, message)


class MultiChannelSubscription:
    """
    多 Channel 订阅管理器

    用于同时订阅多个 channel，共享一个 PubSub 连接
    """

    def __init__(
        self,
        redis_client: redis.Redis,
        name: str = "multi",
        config: PubSubConfig = None,
        parse_json: bool = True,
    ):
        self._redis = redis_client
        self._name = name
        self._config = config or PubSubConfig()
        self._parse_json = parse_json

        self._pubsub: Optional[redis.client.PubSub] = None
        self._running = False
        self._listener_task: Optional[asyncio.Task] = None

        self._handlers: dict[str, Callable] = {}

    async def subscribe(
        self,
        channel: str,
        on_message: Callable[[dict], Union[None, Awaitable[None]]],
    ):
        self._handlers[channel] = on_message

        if self._pubsub:
            await self._pubsub.subscribe(channel)
            logger.debug(f"[PubSub:{self._name}] Added subscription: {channel}")

    async def unsubscribe(self, channel: str):
        if channel in self._handlers:
            del self._handlers[channel]

            if self._pubsub:
                await self._pubsub.unsubscribe(channel)
                logger.debug(f"[PubSub:{self._name}] Removed subscription: {channel}")

    async def start(self):
        if self._running:
            return

        self._running = True
        self._pubsub = await self._create_and_subscribe()
        self._listener_task = asyncio.create_task(self._listen_loop())

        logger.debug(f"[PubSub:{self._name}] Started with {len(self._handlers)} channels")

    async def stop(self):
        if not self._running:
            return

        self._running = False

        if self._listener_task:
            self._listener_task.cancel()
            try:
                await self._listener_task
            except asyncio.CancelledError:
                pass

        if self._pubsub:
            try:
                await self._pubsub.close()
            except Exception:
                pass

        logger.debug(f"[PubSub:{self._name}] Stopped")

    async def _create_and_subscribe(self) -> redis.client.PubSub:
        if self._pubsub:
            try:
                await self._pubsub.close()
            except Exception:
                pass

        pubsub = self._redis.pubsub()

        for channel in self._handlers.keys():
            await pubsub.subscribe(channel)

        if self._config.mark_connection and hasattr(pubsub, 'connection') and pubsub.connection:
            pubsub.connection._is_pubsub_connection = True

        return pubsub

    async def _listen_loop(self):
        retry_delay = self._config.initial_retry_delay

        while self._running:
            try:
                async for message in self._pubsub.listen():
                    if not self._running:
                        break

                    if message['type'] == 'message':
                        channel = message['channel']
                        if isinstance(channel, bytes):
                            channel = channel.decode('utf-8')

                        handler = self._handlers.get(channel)
                        if handler:
                            await self._handle_message(message, handler)

                retry_delay = self._config.initial_retry_delay

            except asyncio.CancelledError:
                break
            except Exception as e:
                if not self._running:
                    break

                logger.warning(f"[PubSub:{self._name}] Error: {e}, reconnecting in {retry_delay}s")
                await asyncio.sleep(retry_delay)

                try:
                    self._pubsub = await self._create_and_subscribe()
                    retry_delay = self._config.initial_retry_delay
                except Exception:
                    retry_delay = min(
                        retry_delay * self._config.retry_backoff_multiplier,
                        self._config.max_retry_delay
                    )

    async def _handle_message(self, message: dict, handler: Callable):
        try:
            data = message['data']
            if self._parse_json:
                try:
                    data = json.loads(data)
                except (json.JSONDecodeError, TypeError):
                    pass

            if asyncio.iscoroutinefunction(handler):
                await handler(data)
            else:
                handler(data)
        except Exception as e:
            logger.error(f"[PubSub:{self._name}] Error handling message: {e}")

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

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


__all__ = [
    'PubSubConfig',
    'PubSubSubscription',
    'PubSubManager',
    'MultiChannelSubscription',
]
