# -*- coding: utf-8 -*-
"""
@File    : rxcore.py
@Author  : Jay
@Date    : 2026/3/19 15:32
@Desc    : Redis connection pool manager supporting sync/async clients,
          distributed locks, and rate-limited pools.
"""
import threading
from contextvars import ContextVar
from types import MappingProxyType
from typing import Optional, Union, Dict, Any

from pydantic import BaseModel, Field, PrivateAttr
from redis import ConnectionPool as SyncConnectionPool, Redis as SyncRedis
from redis.asyncio import ConnectionPool as AsyncConnectionPool, Redis as AsyncRedis

from .awaitable import is_coroutine_context
from .distributed_lock import AsyncRedisAtomicMultiLock
from .rate_limited_pool import Limited
from .rtypes import RedisConfigModel

# Locks & ContextVar for async Redis client pool
_async_pool_rlock = threading.RLock()
_sync_pool_rlock = threading.RLock()
_async_client_pool: ContextVar[Optional[Dict[int, "AsyncRedis"]]] = ContextVar(
    "_async_client_pool", default=None
)


class RedisXsyncStruct(BaseModel):
    """Data structure to hold Redis configuration and client pool."""

    cfgs: Union[RedisConfigModel] = Field(
        description="Redis connection configuration, compatible with redis.Redis parameters"
    )
    pool: Dict[int, Union["RedisXsync", Any]] = Field(
        default_factory=dict, description="Redis client instances per DB"
    )

    # Internal lock, not included in BaseModel serialization
    _lock: threading.Lock = PrivateAttr(default_factory=threading.Lock)

    class Config:
        frozen = True  # Make the model immutable

    @property
    def readonly_pool(self):
        """Return a read-only view of the client pool."""
        return MappingProxyType(self.pool)

    def set_pool(self, db: int, value: Union["RedisXsync", "AsyncRedis", "SyncRedis"]):
        """Thread-safe add client to pool if not exists.

        Args:
            db (int): Redis database index.
            value: Redis client instance.
        """
        with self._lock:
            if db not in self.pool:
                self.pool[db] = value

    def get_pool(self, db: int) -> Union["RedisXsync", "AsyncRedis", "SyncRedis"]:
        """Thread-safe get Redis client from pool.

        Args:
            db (int): Redis database index.

        Returns:
            Redis client instance or None if not exists.
        """
        with self._lock:
            return self.pool.get(db, None)


class RedisXsync(RedisConfigModel):
    """Unified Redis client manager supporting sync/async clients,
    distributed locks, and rate-limited pools.
    """

    def __init__(self, /, host: str = None, port: int = None, db: Union[str, int] = None,
                 password: Optional[str] = None, decode_responses: bool = False, **kwargs):
        """Initialize RedisXsync instance with optional connection overrides.

        Args:
            host (str): Redis server host.
            port (int): Redis server port.
            db (str|int): Redis database index.
            password (str): Redis password.
            decode_responses (bool): Decode bytes to str if True.
            **kwargs: Additional RedisConfigModel fields.
        """
        super().__init__(**kwargs)
        self._normalize_db = lambda dbx: self.db if dbx < 0 else dbx

        self.__rxs = RedisXsyncStruct(
            cfgs=RedisConfigModel(
                host=host,
                port=port,
                db=self._normalize_db(db),
                password=password,
                decode_responses=decode_responses,
                **kwargs
            )
        )

    @property
    def rxs(self):
        return self.__rxs

    @rxs.setter
    def rxs(self, rxs):
        self.__rxs = rxs

    def _create_async_redis_client(self, db=-1):
        """Create or reuse an async Redis client for a given DB.

        Args:
            db (int): Database index.

        Returns:
            AsyncRedis: Async Redis client instance.
        """
        db = self._normalize_db(dbx=db)

        pool = _async_client_pool.get({})
        if db in pool:
            return pool[db]

        with _async_pool_rlock:
            pool = _async_client_pool.get({})
            if db not in pool:
                connection_pool = AsyncConnectionPool(**self._update_cfg(db=db))
                pool[db] = AsyncRedis(connection_pool=connection_pool)
                _async_client_pool.set(pool)

            return pool[db]

    def _create_sync_redis_client(self, db=-1):
        """Create or reuse a sync Redis client for a given DB.

        Args:
            db (int): Database index.

        Returns:
            SyncRedis: Sync Redis client instance.
        """
        db = self._normalize_db(dbx=db)

        sr = self.rxs.get_pool(db)
        if sr:
            return sr

        with _sync_pool_rlock:
            sr = self.rxs.get_pool(db)
            if sr:
                return sr

            connection_pool = SyncConnectionPool(**self._update_cfg(db=db))
            sr = SyncRedis(connection_pool=connection_pool)
            self.rxs.set_pool(db, sr)
            return sr

    def _update_cfg(self, db=-1):
        """Generate Redis connection configuration for the given DB.

        Args:
            db (int): Redis database index.

        Returns:
            dict: Redis connection parameters.
        """
        cfg = self.rxs.cfgs.model_copy(update={"db": db}, deep=True)
        return cfg.model_dump(exclude_unset=True)

    async def __aenter__(self, db=-1) -> AsyncRedis:
        return self._create_async_redis_client(db=db)

    async def __aexit__(self, exc_type, exc, tb):
        pass

    def __enter__(self, db=-1):
        return self._create_sync_redis_client(db=db)

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

    def __call__(self, db=-1, forced_sync=False, **override):
        """Convenience method to get connected client depending on context.

        Args:
            db (int): Database index.
            forced_sync (bool): Force synchronous client.
            **override: Extra Redis config overrides.

        Returns:
            Redis client instance (sync or async).
        """
        is_coro = is_coroutine_context()
        return self.get_connected(db=db, is_coro=is_coro, **override)

    def get_connected(self, db=-1, is_coro=False, **override):
        """Get a Redis client from pool, creating one if necessary.

        Args:
            db (int): Database index.
            is_coro (bool): Return async client if True, else sync client.
            **override: Extra Redis config overrides.

        Returns:
            Redis client instance.
        """
        db = self._normalize_db(dbx=db)
        return self._create_async_redis_client(db) if is_coro else self._create_sync_redis_client(db)

    def AsyncRedisAtomicMultiLock(
            self,
            *keys: str,
            db=-1,
            ttl_ms: int = 5000,
            blocking: bool = False,
            blocking_timeout: int = 5,
            sleep: float = 0.1
    ):
        """Factory for distributed atomic multi-key lock.

        Returns:
            function: Factory function to create AsyncRedisAtomicMultiLock instances.
        """

        db = self._normalize_db(dbx=db)
        return AsyncRedisAtomicMultiLock(
            redis_xsync=self,
            redis_db=db,
            *keys,
            ttl_ms=ttl_ms,
            blocking=blocking,
            blocking_timeout=blocking_timeout,
            sleep=sleep,
        )

    def Limited(self, redis_key: str = None) -> Limited:
        """Create a Limited instance (DB0 only).

        Args:
            redis_key (str, optional): Redis key prefix for the rate-limited pool.

        Returns:
            Limited: A Limited instance bound to this redis client.
        """
        return Limited(redis_xsync=self, redis_key=redis_key)


__all__ = ["RedisXsync"]
