import hashlib
from typing import Any, Awaitable, Dict, Optional, Union, cast

import redis.asyncio as aioredis
from redis.asyncio.cluster import RedisCluster

from rate_limiter.exceptions import StorageBackendError


class RedisStorageBackend:
    """Redis storage backend with atomic operations. Supports both standalone and cluster mode."""

    def __init__(self, redis_url: str, key_prefix: str = "rateon", cluster_mode: bool = False):
        """
        Initialize Redis storage backend.

        Args:
            redis_url: Redis connection URL. For cluster mode, provide a single node URL
                      (client will auto-discover other nodes). Multiple URLs (comma-separated)
                      are supported for redundancy but only one is required.
            key_prefix: Prefix for all Redis keys
            cluster_mode: Enable Redis Cluster mode
        """
        self.redis_url = redis_url
        self.key_prefix = key_prefix
        self.cluster_mode = cluster_mode
        self._client: Optional[Union[aioredis.Redis, RedisCluster]] = None
        self._scripts: Dict[str, str] = {}  # Cache for Lua scripts

    async def _get_client(self) -> Union[aioredis.Redis, RedisCluster]:
        """Get or create Redis client (standalone or cluster)."""
        if self._client is None:
            try:
                if self.cluster_mode:
                    # Parse startup node(s) from URL
                    # Only one node is needed - client will discover the rest of the cluster
                    nodes = self._parse_cluster_nodes(self.redis_url)
                    self._client = RedisCluster(
                        startup_nodes=nodes,
                        decode_responses=False,  # We'll handle encoding ourselves
                        socket_connect_timeout=1,
                        socket_timeout=1,
                        skip_full_coverage_check=True,  # Allow partial cluster coverage
                    )
                    # Initialize cluster connection (discovers all nodes automatically)
                    await self._client.initialize()
                else:
                    self._client = await aioredis.from_url(
                        self.redis_url,
                        decode_responses=False,  # We'll handle encoding ourselves
                        socket_connect_timeout=1,
                        socket_timeout=1,
                    )
            except Exception as e:
                raise StorageBackendError(f"Failed to connect to Redis: {e}") from e
        return self._client

    def _parse_cluster_nodes(self, url_or_nodes: str) -> list[dict]:
        """
        Parse cluster startup node(s) from URL string.

        For Redis Cluster, only one node is required - the client will automatically
        discover all other nodes in the cluster. Multiple nodes can be provided for
        redundancy (if one is down, others can be tried).

        Args:
            url_or_nodes: Single Redis URL or comma-separated URLs for redundancy

        Returns:
            List of node dictionaries with 'host' and 'port' keys
        """
        nodes = []
        # Split by comma if multiple URLs provided (for redundancy)
        node_strings = url_or_nodes.split(",") if "," in url_or_nodes else [url_or_nodes]
        
        for node_str in node_strings:
            node_str = node_str.strip()
            if not node_str:
                continue

            # Parse URL format: redis://host:port or host:port
            if node_str.startswith("redis://"):
                node_str = node_str[8:]
            elif node_str.startswith("rediss://"):
                node_str = node_str[9:]

            if ":" in node_str:
                host, port = node_str.rsplit(":", 1)
                nodes.append({"host": host, "port": int(port)})
            else:
                # Default port
                nodes.append({"host": node_str, "port": 6379})

        if not nodes:
            raise ValueError("No valid cluster nodes found. Provide at least one node URL.")

        return nodes

    def _make_key(self, key: str) -> bytes:
        """
        Create Redis key with prefix and sanitization.

        For cluster mode, uses hash tags to ensure keys are in the same slot.

        Args:
            key: Original key

        Returns:
            Sanitized Redis key as bytes
        """
        # Sanitize key to prevent injection
        safe_key = hashlib.sha256(key.encode()).hexdigest()[:32]
        
        if self.cluster_mode:
            # Use hash tags for cluster mode to ensure keys are in the same slot
            # This is critical for Lua scripts which require all keys in the same slot
            full_key = f"{self.key_prefix}:{{{safe_key}}}"
        else:
            full_key = f"{self.key_prefix}:{safe_key}"
        
        return full_key.encode()

    async def incr(self, key: str, ttl: int) -> int:
        """
        Increment counter atomically and return new value.

        Args:
            key: Storage key
            ttl: Time to live in seconds

        Returns:
            New counter value
        """
        try:
            client = await self._get_client()
            redis_key = self._make_key(key)

            # Use Lua script for atomic increment + TTL
            script = """
            local current = redis.call('INCR', KEYS[1])
            if current == 1 then
                redis.call('EXPIRE', KEYS[1], ARGV[1])
            end
            return current
            """
            # Handle both standalone and cluster modes
            # Type cast needed because redis-py type stubs may not correctly infer awaitable return type
            eval_result = client.eval(script, 1, redis_key, ttl)
            result = await cast(Awaitable[Any], eval_result)
            return int(result) if result else 0
        except Exception as e:
            raise StorageBackendError(f"Redis INCR failed: {e}") from e

    async def get(self, key: str) -> int:
        """
        Get current counter value.

        Args:
            key: Storage key

        Returns:
            Current counter value (0 if key doesn't exist)
        """
        try:
            client = await self._get_client()
            redis_key = self._make_key(key)
            result = await client.get(redis_key)
            return int(result) if result else 0
        except Exception as e:
            raise StorageBackendError(f"Redis GET failed: {e}") from e

    async def set(self, key: str, value: int, ttl: int) -> None:
        """
        Set counter value.

        Args:
            key: Storage key
            value: Counter value
            ttl: Time to live in seconds
        """
        try:
            client = await self._get_client()
            redis_key = self._make_key(key)
            await client.setex(redis_key, ttl, value)
        except Exception as e:
            raise StorageBackendError(f"Redis SET failed: {e}") from e

    async def delete(self, key: str) -> None:
        """
        Delete a key.

        Args:
            key: Storage key
        """
        try:
            client = await self._get_client()
            redis_key = self._make_key(key)
            await client.delete(redis_key)
        except Exception as e:
            raise StorageBackendError(f"Redis DELETE failed: {e}") from e

    async def exists(self, key: str) -> bool:
        """
        Check if key exists.

        Args:
            key: Storage key

        Returns:
            True if key exists, False otherwise
        """
        try:
            client = await self._get_client()
            redis_key = self._make_key(key)
            result = await client.exists(redis_key)
            return bool(result)
        except Exception as e:
            raise StorageBackendError(f"Redis EXISTS failed: {e}") from e

    async def ttl(self, key: str) -> int:
        """
        Get remaining TTL for a key.

        Args:
            key: Storage key

        Returns:
            Remaining TTL in seconds (-1 if key doesn't exist, -2 if no expiry)
        """
        try:
            client = await self._get_client()
            redis_key = self._make_key(key)
            result = await client.ttl(redis_key)
            return int(result) if result is not None else -1
        except Exception as e:
            raise StorageBackendError(f"Redis TTL failed: {e}") from e

    async def close(self) -> None:
        """Close Redis connection."""
        if self._client:
            await self._client.aclose()
            self._client = None

    async def execute_script(self, script: str, keys: list[str], args: list[Any]) -> Any:
        """
        Execute Lua script with keys and arguments.

        For cluster mode, ensures all keys use hash tags to be in the same slot.

        Args:
            script: Lua script content
            keys: List of keys
            args: List of arguments

        Returns:
            Script execution result
        """
        try:
            client = await self._get_client()
            # Hash script to cache it
            script_hash = hashlib.sha1(script.encode()).hexdigest()

            # Convert keys to bytes (hash tags are added in _make_key for cluster mode)
            redis_keys = [self._make_key(k) for k in keys]

            # Convert args to strings/bytes
            redis_args = [str(arg).encode() if not isinstance(arg, bytes) else arg for arg in args]

            if self.cluster_mode:
                # For cluster mode, use eval with script content directly
                # EVALSHA doesn't work well with cluster as scripts need to be loaded on each node
                eval_result = client.eval(script, len(redis_keys), *redis_keys, *redis_args)
                result = await cast(Awaitable[Any], eval_result)
            else:
                # Use EVALSHA if script is cached, otherwise use EVAL
                try:
                    evalsha_result = client.evalsha(
                        script_hash, len(redis_keys), *redis_keys, *redis_args
                    )
                    result = await cast(Awaitable[Any], evalsha_result)
                except aioredis.ResponseError:
                    # Script not cached, use EVAL and cache it
                    eval_result = client.eval(script, len(redis_keys), *redis_keys, *redis_args)
                    result = await cast(Awaitable[Any], eval_result)
                    # Cache script for future use
                    await client.script_load(script)
                    self._scripts[script_hash] = script

            return result
        except Exception as e:
            raise StorageBackendError(f"Redis script execution failed: {e}") from e
