import functools
import logging
import socket
import ssl
import sys
from collections.abc import Callable
from datetime import timedelta
from typing import Any, Union, cast

import redis
from redis import RedisError
from redis.backoff import ExponentialWithJitterBackoff  # type: ignore
from redis.cache import CacheConfig
from redis.client import PubSub
from redis.cluster import ClusterNode, RedisCluster
from redis.connection import Connection, SSLConnection
from redis.exceptions import ConnectionError, TimeoutError
from redis.retry import Retry
from redis.sentinel import Sentinel
from typing_extensions import TypedDict

from configs import dify_config
from dify_app import DifyApp
from extensions.redis_names import (
    normalize_redis_key_prefix,
    serialize_redis_name,
    serialize_redis_name_arg,
    serialize_redis_name_args,
)
from libs.broadcast_channel.channel import BroadcastChannel as BroadcastChannelProtocol
from libs.broadcast_channel.redis.pubsub_channel import BroadcastChannel as RedisBroadcastChannel
from libs.broadcast_channel.redis.sharded_channel import ShardedRedisBroadcastChannel
from libs.broadcast_channel.redis.streams_channel import StreamsBroadcastChannel

logger = logging.getLogger(__name__)


_normalize_redis_key_prefix = normalize_redis_key_prefix
_serialize_redis_name = serialize_redis_name
_serialize_redis_name_arg = serialize_redis_name_arg
_serialize_redis_name_args = serialize_redis_name_args


class RedisClientWrapper:
    """
    A wrapper class for the Redis client that addresses the issue where the global
    `redis_client` variable cannot be updated when a new Redis instance is returned
    by Sentinel.

    This class allows for deferred initialization of the Redis client, enabling the
    client to be re-initialized with a new instance when necessary. This is particularly
    useful in scenarios where the Redis instance may change dynamically, such as during
    a failover in a Sentinel-managed Redis setup.

    Attributes:
        _client: The actual Redis client instance. It remains None until
                 initialized with the `initialize` method.

    Methods:
        initialize(client): Initializes the Redis client if it hasn't been initialized already.
        __getattr__(item): Delegates attribute access to the Redis client, raising an error
                           if the client is not initialized.
    """

    _client: Union[redis.Redis, RedisCluster, None]

    def __init__(self) -> None:
        self._client = None

    def initialize(self, client: Union[redis.Redis, RedisCluster]) -> None:
        if self._client is None:
            self._client = client

    def _require_client(self) -> redis.Redis | RedisCluster:
        if self._client is None:
            raise RuntimeError("Redis client is not initialized. Call init_app first.")
        return self._client

    def _get_prefix(self) -> str:
        return dify_config.REDIS_KEY_PREFIX

    def get(self, name: str | bytes) -> Any:
        return self._require_client().get(_serialize_redis_name_arg(name, self._get_prefix()))

    def set(
        self,
        name: str | bytes,
        value: Any,
        ex: int | None = None,
        px: int | None = None,
        nx: bool = False,
        xx: bool = False,
        keepttl: bool = False,
        get: bool = False,
        exat: int | None = None,
        pxat: int | None = None,
    ) -> Any:
        return self._require_client().set(
            _serialize_redis_name_arg(name, self._get_prefix()),
            value,
            ex=ex,
            px=px,
            nx=nx,
            xx=xx,
            keepttl=keepttl,
            get=get,
            exat=exat,
            pxat=pxat,
        )

    def setex(self, name: str | bytes, time: int | timedelta, value: Any) -> Any:
        return self._require_client().setex(_serialize_redis_name_arg(name, self._get_prefix()), time, value)

    def setnx(self, name: str | bytes, value: Any) -> Any:
        return self._require_client().setnx(_serialize_redis_name_arg(name, self._get_prefix()), value)

    def delete(self, *names: str | bytes) -> Any:
        return self._require_client().delete(*_serialize_redis_name_args(names, self._get_prefix()))

    def incr(self, name: str | bytes, amount: int = 1) -> Any:
        return self._require_client().incr(_serialize_redis_name_arg(name, self._get_prefix()), amount)

    def expire(
        self,
        name: str | bytes,
        time: int | timedelta,
        nx: bool = False,
        xx: bool = False,
        gt: bool = False,
        lt: bool = False,
    ) -> Any:
        return self._require_client().expire(
            _serialize_redis_name_arg(name, self._get_prefix()),
            time,
            nx=nx,
            xx=xx,
            gt=gt,
            lt=lt,
        )

    def exists(self, *names: str | bytes) -> Any:
        return self._require_client().exists(*_serialize_redis_name_args(names, self._get_prefix()))

    def ttl(self, name: str | bytes) -> Any:
        return self._require_client().ttl(_serialize_redis_name_arg(name, self._get_prefix()))

    def getdel(self, name: str | bytes) -> Any:
        return self._require_client().getdel(_serialize_redis_name_arg(name, self._get_prefix()))

    def lock(
        self,
        name: str,
        timeout: float | None = None,
        sleep: float = 0.1,
        blocking: bool = True,
        blocking_timeout: float | None = None,
        thread_local: bool = True,
    ) -> Any:
        return self._require_client().lock(
            _serialize_redis_name(name, self._get_prefix()),
            timeout=timeout,
            sleep=sleep,
            blocking=blocking,
            blocking_timeout=blocking_timeout,
            thread_local=thread_local,
        )

    def hset(self, name: str | bytes, *args: Any, **kwargs: Any) -> Any:
        return self._require_client().hset(_serialize_redis_name_arg(name, self._get_prefix()), *args, **kwargs)

    def hgetall(self, name: str | bytes) -> Any:
        return self._require_client().hgetall(_serialize_redis_name_arg(name, self._get_prefix()))

    def hkeys(self, name: str | bytes) -> Any:
        return self._require_client().hkeys(_serialize_redis_name_arg(name, self._get_prefix()))

    def hexists(self, name: str | bytes, key: str | bytes) -> Any:
        return self._require_client().hexists(_serialize_redis_name_arg(name, self._get_prefix()), key)

    def hdel(self, name: str | bytes, *keys: str | bytes) -> Any:
        return self._require_client().hdel(_serialize_redis_name_arg(name, self._get_prefix()), *keys)

    def hlen(self, name: str | bytes) -> Any:
        return self._require_client().hlen(_serialize_redis_name_arg(name, self._get_prefix()))

    def zadd(
        self,
        name: str | bytes,
        mapping: dict[str | bytes | int | float, float | int | str | bytes],
        nx: bool = False,
        xx: bool = False,
        ch: bool = False,
        incr: bool = False,
        gt: bool = False,
        lt: bool = False,
    ) -> Any:
        return self._require_client().zadd(
            _serialize_redis_name_arg(name, self._get_prefix()),
            cast(Any, mapping),
            nx=nx,
            xx=xx,
            ch=ch,
            incr=incr,
            gt=gt,
            lt=lt,
        )

    def zremrangebyscore(self, name: str | bytes, min: float | str, max: float | str) -> Any:
        return self._require_client().zremrangebyscore(_serialize_redis_name_arg(name, self._get_prefix()), min, max)

    def zcard(self, name: str | bytes) -> Any:
        return self._require_client().zcard(_serialize_redis_name_arg(name, self._get_prefix()))

    def pubsub(self) -> PubSub:
        return self._require_client().pubsub()

    def pipeline(self, transaction: bool = True, shard_hint: str | None = None) -> Any:
        return self._require_client().pipeline(transaction=transaction, shard_hint=shard_hint)

    def __getattr__(self, item: str) -> Any:
        return getattr(self._require_client(), item)


redis_client: RedisClientWrapper = RedisClientWrapper()
_pubsub_redis_client: redis.Redis | RedisCluster | None = None


class RedisSSLParamsDict(TypedDict):
    ssl_cert_reqs: int
    ssl_ca_certs: str | None
    ssl_certfile: str | None
    ssl_keyfile: str | None


class RedisHealthParamsDict(TypedDict):
    retry: Retry
    socket_timeout: float | None
    socket_connect_timeout: float | None
    health_check_interval: int | None
    socket_keepalive: bool
    socket_keepalive_options: dict[int, int]


class RedisClusterHealthParamsDict(TypedDict):
    retry: Retry
    socket_timeout: float | None
    socket_connect_timeout: float | None
    socket_keepalive: bool
    socket_keepalive_options: dict[int, int]


class RedisBaseParamsDict(TypedDict):
    username: str | None
    password: str | None
    db: int
    encoding: str
    encoding_errors: str
    decode_responses: bool
    protocol: int
    cache_config: CacheConfig | None
    retry: Retry
    socket_timeout: float | None
    socket_connect_timeout: float | None
    health_check_interval: int | None
    socket_keepalive: bool
    socket_keepalive_options: dict[int, int]


def _get_ssl_configuration() -> tuple[type[Union[Connection, SSLConnection]], dict[str, Any]]:
    """Get SSL configuration for Redis connection."""
    if not dify_config.REDIS_USE_SSL:
        return Connection, {}

    cert_reqs_map = {
        "CERT_NONE": ssl.CERT_NONE,
        "CERT_OPTIONAL": ssl.CERT_OPTIONAL,
        "CERT_REQUIRED": ssl.CERT_REQUIRED,
    }
    ssl_cert_reqs = cert_reqs_map.get(dify_config.REDIS_SSL_CERT_REQS, ssl.CERT_NONE)

    ssl_kwargs = {
        "ssl_cert_reqs": ssl_cert_reqs,
        "ssl_ca_certs": dify_config.REDIS_SSL_CA_CERTS,
        "ssl_certfile": dify_config.REDIS_SSL_CERTFILE,
        "ssl_keyfile": dify_config.REDIS_SSL_KEYFILE,
    }

    return SSLConnection, ssl_kwargs


def _get_cache_configuration() -> CacheConfig | None:
    """Get client-side cache configuration if enabled."""
    if not dify_config.REDIS_ENABLE_CLIENT_SIDE_CACHE:
        return None

    resp_protocol = dify_config.REDIS_SERIALIZATION_PROTOCOL
    if resp_protocol < 3:
        raise ValueError("Client side cache is only supported in RESP3")

    return CacheConfig()


def _get_retry_policy() -> Retry:
    """Build the shared retry policy for Redis connections."""
    return Retry(
        backoff=ExponentialWithJitterBackoff(
            base=dify_config.REDIS_RETRY_BACKOFF_BASE,
            cap=dify_config.REDIS_RETRY_BACKOFF_CAP,
        ),
        retries=dify_config.REDIS_RETRY_RETRIES,
        supported_errors=(
            ConnectionError,
            TimeoutError,
            BrokenPipeError,
            OSError,
        ),
    )


def _get_connection_health_params() -> RedisHealthParamsDict:
    """Get connection health and retry parameters for standalone and Sentinel Redis clients."""
    socket_keepalive_options: dict[int, int] = {}
    if sys.platform == "linux":
        socket_keepalive_options[socket.TCP_KEEPIDLE] = dify_config.REDIS_KEEPALIVE_IDLE
        socket_keepalive_options[socket.TCP_KEEPINTVL] = dify_config.REDIS_KEEPALIVE_INTERVAL
        socket_keepalive_options[socket.TCP_KEEPCNT] = dify_config.REDIS_KEEPALIVE_COUNT
    elif sys.platform == "darwin":
        socket_keepalive_options[socket.TCP_KEEPALIVE] = dify_config.REDIS_KEEPALIVE_IDLE

    return RedisHealthParamsDict(
        retry=_get_retry_policy(),
        socket_timeout=dify_config.REDIS_SOCKET_TIMEOUT,
        socket_connect_timeout=dify_config.REDIS_SOCKET_CONNECT_TIMEOUT,
        health_check_interval=dify_config.REDIS_HEALTH_CHECK_INTERVAL,
        socket_keepalive=dify_config.REDIS_KEEPALIVE,
        socket_keepalive_options=socket_keepalive_options,
    )


def _get_cluster_connection_health_params() -> RedisClusterHealthParamsDict:
    """Get retry and timeout parameters for Redis Cluster clients.

    RedisCluster does not support ``health_check_interval`` as a constructor
    keyword (it is silently stripped by ``cleanup_kwargs``), so it is excluded
    here. Only ``retry``, ``socket_timeout``, and ``socket_connect_timeout``
    are passed through.
    """
    health_params = _get_connection_health_params()
    result: RedisClusterHealthParamsDict = {
        "retry": health_params["retry"],
        "socket_timeout": health_params["socket_timeout"],
        "socket_connect_timeout": health_params["socket_connect_timeout"],
        "socket_keepalive": health_params["socket_keepalive"],
        "socket_keepalive_options": health_params["socket_keepalive_options"],
    }
    return result


def _get_base_redis_params() -> RedisBaseParamsDict:
    """Get base Redis connection parameters including retry and health policy."""
    return RedisBaseParamsDict(
        username=dify_config.REDIS_USERNAME,
        password=dify_config.REDIS_PASSWORD or None,
        db=dify_config.REDIS_DB,
        encoding="utf-8",
        encoding_errors="strict",
        decode_responses=False,
        protocol=dify_config.REDIS_SERIALIZATION_PROTOCOL,
        cache_config=_get_cache_configuration(),
        **_get_connection_health_params(),
    )


def _create_sentinel_client(redis_params: RedisBaseParamsDict) -> Union[redis.Redis, RedisCluster]:
    """Create Redis client using Sentinel configuration."""
    if not dify_config.REDIS_SENTINELS:
        raise ValueError("REDIS_SENTINELS must be set when REDIS_USE_SENTINEL is True")

    if not dify_config.REDIS_SENTINEL_SERVICE_NAME:
        raise ValueError("REDIS_SENTINEL_SERVICE_NAME must be set when REDIS_USE_SENTINEL is True")

    sentinel_hosts = [(node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")]

    health_params = _get_connection_health_params()

    sentinel_kwargs = {
        "socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,
        "username": dify_config.REDIS_SENTINEL_USERNAME,
        "password": dify_config.REDIS_SENTINEL_PASSWORD,
        "socket_keepalive": health_params["socket_keepalive"],
        "socket_keepalive_options": health_params["socket_keepalive_options"],
    }

    if dify_config.REDIS_MAX_CONNECTIONS:
        sentinel_kwargs["max_connections"] = dify_config.REDIS_MAX_CONNECTIONS

    sentinel = Sentinel(
        sentinel_hosts,
        sentinel_kwargs=sentinel_kwargs,
    )

    params: dict[str, Any] = {**redis_params}
    master: redis.Redis = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **params)
    return master


def _create_cluster_client() -> Union[redis.Redis, RedisCluster]:
    """Create Redis cluster client."""
    if not dify_config.REDIS_CLUSTERS:
        raise ValueError("REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True")

    nodes = [
        ClusterNode(host=node.split(":")[0], port=int(node.split(":")[1]))
        for node in dify_config.REDIS_CLUSTERS.split(",")
    ]

    cluster_kwargs: dict[str, Any] = {
        "startup_nodes": nodes,
        "password": dify_config.REDIS_CLUSTERS_PASSWORD,
        "protocol": dify_config.REDIS_SERIALIZATION_PROTOCOL,
        "cache_config": _get_cache_configuration(),
        **_get_cluster_connection_health_params(),
    }
    if dify_config.REDIS_MAX_CONNECTIONS:
        cluster_kwargs["max_connections"] = dify_config.REDIS_MAX_CONNECTIONS
    cluster: RedisCluster = RedisCluster(**cluster_kwargs)
    return cluster


def _create_standalone_client(redis_params: RedisBaseParamsDict) -> Union[redis.Redis, RedisCluster]:
    """Create standalone Redis client."""
    connection_class, ssl_kwargs = _get_ssl_configuration()

    params: dict[str, Any] = {
        **redis_params,
        "host": dify_config.REDIS_HOST,
        "port": dify_config.REDIS_PORT,
        "connection_class": connection_class,
    }

    if dify_config.REDIS_MAX_CONNECTIONS:
        params["max_connections"] = dify_config.REDIS_MAX_CONNECTIONS

    if ssl_kwargs:
        params.update(ssl_kwargs)

    pool = redis.ConnectionPool(**params)
    client: redis.Redis = redis.Redis(connection_pool=pool)
    return client


def _create_pubsub_client(pubsub_url: str, use_clusters: bool) -> redis.Redis | RedisCluster:
    max_conns = dify_config.REDIS_MAX_CONNECTIONS

    if use_clusters:
        health_params = _get_cluster_connection_health_params()
        kwargs: dict[str, Any] = {**health_params}
        if max_conns:
            kwargs["max_connections"] = max_conns
        return RedisCluster.from_url(pubsub_url, **kwargs)

    standalone_health_params: dict[str, Any] = dict(_get_connection_health_params())
    kwargs = {**standalone_health_params}
    if max_conns:
        kwargs["max_connections"] = max_conns
    return redis.Redis.from_url(pubsub_url, **kwargs)


def init_app(app: DifyApp):
    """Initialize Redis client and attach it to the app."""
    global redis_client

    # Determine Redis mode and create appropriate client
    if dify_config.REDIS_USE_SENTINEL:
        redis_params = _get_base_redis_params()
        client = _create_sentinel_client(redis_params)
    elif dify_config.REDIS_USE_CLUSTERS:
        client = _create_cluster_client()
    else:
        redis_params = _get_base_redis_params()
        client = _create_standalone_client(redis_params)

    # Initialize the wrapper and attach to app
    redis_client.initialize(client)
    app.extensions["redis"] = redis_client

    global _pubsub_redis_client
    _pubsub_redis_client = client
    if dify_config.normalized_pubsub_redis_url:
        _pubsub_redis_client = _create_pubsub_client(
            dify_config.normalized_pubsub_redis_url, dify_config.PUBSUB_REDIS_USE_CLUSTERS
        )


def get_pubsub_broadcast_channel() -> BroadcastChannelProtocol:
    assert _pubsub_redis_client is not None, "PubSub redis Client should be initialized here."
    if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "sharded":
        return ShardedRedisBroadcastChannel(_pubsub_redis_client)
    if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "streams":
        return StreamsBroadcastChannel(
            _pubsub_redis_client,
            retention_seconds=dify_config.PUBSUB_STREAMS_RETENTION_SECONDS,
        )
    return RedisBroadcastChannel(_pubsub_redis_client)


def redis_fallback[T](default_return: T | None = None):  # type: ignore
    """
    decorator to handle Redis operation exceptions and return a default value when Redis is unavailable.

    Args:
        default_return: The value to return when a Redis operation fails. Defaults to None.
    """

    def decorator[**P, R](func: Callable[P, R]) -> Callable[P, R | T | None]:
        @functools.wraps(func)
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R | T | None:
            try:
                return func(*args, **kwargs)
            except RedisError as e:
                func_name = getattr(func, "__name__", "Unknown")
                logger.warning("Redis operation failed in %s: %s", func_name, str(e), exc_info=True)
                return default_return

        return wrapper

    return decorator