"""Client wrapper for a redis connection."""

import json
import logging
from abc import ABC, abstractmethod
from typing import Any, AnyStr, Dict, Generic, List, Optional, Type, TypeVar

import redis.asyncio as redis
from ddtrace.trace import tracer
from fakeredis.aioredis import FakeRedis
from pydantic import BaseModel, TypeAdapter
from redis.exceptions import RedisError

from pdp import config

logger = logging.getLogger(__name__)

T = TypeVar("T")
TypeAdapterT = TypeVar("TypeAdapterT")


class PPCacheItemSerializer(ABC):
    """Base cache item serializer class."""

    @abstractmethod
    def load(self, cached_data: AnyStr) -> Any:
        """Abstract deserializer method."""
        raise NotImplementedError("Base class methods are not implemented.")

    @abstractmethod
    def dump(self, data: Any) -> Any:
        """Abstract serializer method."""
        raise NotImplementedError("Base class methods are not implemented.")


class NullSerializer(PPCacheItemSerializer):
    """Performs no serializations or deserialization.

    Used by the cache administration endpoints to read ANY value from cache.
    """

    @tracer.wrap("NullSerializer.load")
    def load(self, cached_data: AnyStr) -> Optional[str]:
        """Return the cached string as-is, without deserialization."""
        if not cached_data:
            return None

        if isinstance(cached_data, bytes):
            # redis commands return bytes
            result: str = str(cached_data, encoding="utf-8")
        else:
            result = cached_data

        return result

    @tracer.wrap("NullSerializer.dump")
    def dump(self, data: Any) -> Any:
        """Not supported. Use JSONSerializer instead."""
        raise NotImplementedError("NullSerializer does not support dump()")


class JSONSerializer(PPCacheItemSerializer):
    """Plain old JSON serializer."""

    @tracer.wrap("JSONSerializer.load")
    def load(self, cached_data: AnyStr) -> Optional[Dict[str, Any]]:
        """Convert json-string to Dictionary."""
        if not cached_data:
            return None

        result: Dict[str, Any] = json.loads(cached_data)
        return result

    @tracer.wrap("JSONSerializer.dump")
    def dump(self, data: Dict[str, Any]) -> Optional[str]:
        """Convert dictionary to json-string."""
        if not data:
            return None
        return json.dumps(data)


class PydanticSchemaSerializer(PPCacheItemSerializer):
    """Serializer for schema objects derived from pydantic.BaseModel.

    Examples (r = RedisConnector())
        # GET
        result = await r.get(key=cache_key,
                             serializer=PydanticSchemaSerializer(RolesResponse))

        # SET
        await r.set(key=cache_key,
                    item=result,
                    serializer=PydanticSchemaSerializer(RolesResponse))
    """

    def __init__(self, schema: Type[BaseModel]):
        """Set BaseModel schema used by serializer."""
        self.schema = schema

    @tracer.wrap("PydanticSchemaSerializer.load")
    def load(self, cached_data: AnyStr) -> Optional[BaseModel]:
        """Convert json-string to Pydantic schema object."""
        if not cached_data:
            return None

        return self.schema.model_validate_json(cached_data)

    @tracer.wrap("PydanticSchemaSerializer.dump")
    def dump(self, data: BaseModel) -> Optional[str]:
        """Convert Pydantic schema object to json-string."""
        if not data:
            return None
        return self.schema.model_dump_json(data)


# NOTE:
# TypeAdapterT represents the typevar 'T' in TypeAdapter[T]
# For example:
#   RoleListValidator is a TypeAdapter(List[Role]), then T is List[Role]
#   so for TypeAdapterSerializer(RoleListValidator), TypeAdapterT is List[Role]
class TypeAdapterSerializer(PPCacheItemSerializer, Generic[TypeAdapterT]):
    """Serializer for schema objects derived from pydantic.TypeAdapter.

    Examples (r = RedisConnector())
        # GET
        tenants = await r.get(key=cache_key,
                              serializer=TypeAdapterSerializer(TenantRolesMapValidator))

        # SET
        await redis_connector.set(key=cache_key,
                                  item=tenants,
                                  serializer=TypeAdapterSerializer(TenantRolesMapValidator))
    """

    def __init__(
        self,
        type_adapter: TypeAdapter[TypeAdapterT],
        json_validation_enabled: bool = True,
    ):
        """Set TypeAdapter class used by the Serializer.

        Args:
            type_adapter: Pydantic TypeAdapter
            json_validation_enabled: If true, attempt to deserialize the dump_json()
             output to validate it matches TypeAdapterT. Defaults to true.
        """
        self.type_adapter = type_adapter
        self.json_validation_enabled = json_validation_enabled

    @tracer.wrap("TypeAdapterSerializer.load")
    def load(self, cached_data: AnyStr) -> Optional[TypeAdapterT]:
        """Convert the cached_data to a TypeAdapter schema object."""
        if not cached_data:
            return None
        return self.type_adapter.validate_json(cached_data)

    @tracer.wrap("TypeAdapterSerializer.dump")
    def dump(self, data: TypeAdapterT) -> Optional[bytes]:
        """Convert a TypeAdapter schema object to a json-string."""
        if not data:
            return None

        json_data = self.type_adapter.dump_json(data)

        if self.json_validation_enabled:
            # Should raise pydantic_core._pydantic_core.ValidationError
            # if there's a problem. Unfortunately, `dump_json(round_trip=True)`
            # does not work as expected.
            assert self.load(json_data) == data

        return json_data


@tracer.wrap()
def _safe_load(
    key: AnyStr,
    cached_data: str,
    serializer: PPCacheItemSerializer,
) -> Optional[T]:
    """Attempt to load the item, handle exceptions, and log errors."""
    result = None
    try:
        result = serializer.load(cached_data)
    except (ValueError, TypeError):
        # Return 'None' for load errors. Callers should consider this a cache miss.
        # Both JSONDecodeError and pydantic.ValidationError are subtypes of ValueError.
        logger.error(
            "Load failed for key '%s' with serializer type: '%s'",
            key,
            type(serializer),
            exc_info=True,
        )
    except Exception:
        logger.error(
            "Unexpected load error for key '%s' with serializer type: '%s'",
            key,
            type(serializer),
            exc_info=True,
        )
    return result


class RedisConnector:
    """Connect to Redis.

    When adding new command support to this RedisConnector class, wrap any
    calls to self.client.method_name with try...except RedisError... Redis is used
    as a cache, and in the event of any RedisError, treat it as a cache-miss.

    The exception is when the command is added to support infra-related endpoints.
    Thus, the following do not follow this try...except pattern:
        * ping
        * dbsize
        * getdel
    """

    def __init__(self, redis_url: str, use_redis_cache: bool = False):
        """Create a redis or fake redis connector."""
        logger.info("Redis URL: %s", redis_url)
        if use_redis_cache:
            _client: redis.client.Redis = redis.Redis.from_url(
                url=redis_url,
                socket_connect_timeout=config.REDIS_SOCKET_CONNECT_TIMEOUT,
                socket_timeout=config.REDIS_SOCKET_TIMEOUT,
            )
        else:
            _client: FakeRedis = FakeRedis()  # type:ignore

        self._client = _client
        self._url = redis_url
        self._use_redis_cache = use_redis_cache

    @property
    def client(self) -> redis.Redis:
        """Returns the internal redis client."""
        return self._client

    @tracer.wrap()
    async def ping(self) -> bool:
        """
        Ping the Redis server.

        For more information see https://redis.io/commands/ping
        """
        try:
            response = await self.client.ping()
        except RedisError as error:
            logger.error(
                f"Failed to ping Redis at {self._url} with use_redis_cache = {self._use_redis_cache}",  # noqa: E501
                exc_info=error,
            )
            return False
        return bool(response)

    @tracer.wrap()
    async def dbsize(self) -> int:
        """
        Return the number of keys in the current database.

        For more information see https://redis.io/commands/dbsize
        """
        response = await self.client.dbsize()
        return int(response)

    @tracer.wrap()
    async def mget(
        self, keys: List[str], serializer: PPCacheItemSerializer
    ) -> List[Optional[T]]:
        """
        Fetch cached entries using Redis.MGET and deserialize to a schema type.

        Args:
              keys: List of redis keys to fetch
              serializer: Object to deserialize the cached JSON-formatted strings
                to Python Objects

        Returns:
            A list of deserialized cached items. List entries can be 'None'
            if the deserialization fails.
        """
        if not keys:
            return []

        decoded_entries: List[Optional[T]] = []
        try:
            cached_data_entries = await self.client.mget(keys=keys)
        except RedisError:
            logger.error(
                "[RedisError] mget failed, returning array of Nones", exc_info=True
            )
            return [None] * len(keys)

        if not cached_data_entries:
            return [None] * len(keys)

        for key, cached_data in zip(keys, cached_data_entries):
            result = _safe_load(key=key, cached_data=cached_data, serializer=serializer)
            decoded_entries.append(result)

        return decoded_entries

    @tracer.wrap()
    async def get(self, key: str, serializer: PPCacheItemSerializer) -> Optional[T]:
        """
        Fetch one cached entry and deserialize to a schema type.

        Args:
            key: Redis key to fetch
            serializer: Object to deserialize the cached JSON-formatted strings
              to Python Objects

        Returns:
            A deserialized cached item
        """
        result: List[Optional[T]] = await self.mget([key], serializer)
        if not result:
            return None

        return result[0]

    @tracer.wrap()
    async def getdel(
        self, keys: List[str], serializer: PPCacheItemSerializer
    ) -> List[Optional[T]]:
        """
        Get the values for the keys and delete the entries.

        Args:
              keys: List of redis keys to delete
              serializer: Object to deserialize the cached JSON-formatted strings
                to Python Objects

        Returns:
            A list of deserialized cached items. Deleted entries can be 'None'
            if the deserialization fails.
        """
        if not keys:
            return []

        decoded_entries: List[Optional[T]] = []
        for key in keys:
            cached_data = await self.client.getdel(name=key)
            result = _safe_load(key=key, cached_data=cached_data, serializer=serializer)
            decoded_entries.append(result)

        return decoded_entries

    @tracer.wrap()
    async def set(
        self,
        key: str,
        item: T,
        serializer: PPCacheItemSerializer,
        ttl: Optional[int] = config.REDIS_CACHE_TTL,
        redis_client: Optional[redis.Redis] = None,
    ) -> bool:
        """
        Serialize the item and write a cached entry using Redis.SET.

        Args:
            key: Redis key to store the value
            item: Python object to store
            serializer: Object to serialize the Python object to a
                JSON-formatted strings
            ttl: Cache EXPIRE time in seconds.
            redis_client: Redis Client to use for set, e.g. Pipeline.
                If left empty/None, will use the RedisConnector's client

        Returns:
            True if the SET command succeeds. Otherwise, False.
        """
        if not item:
            return False

        if not redis_client:
            redis_client = self.client

        json_data = serializer.dump(item)
        try:
            result = await redis_client.set(key, json_data, ex=ttl)
        except RedisError:
            logger.error("[RedisError] set failed, returning False", exc_info=True)
            return False
        return bool(result)

    @tracer.wrap()
    async def mset(
        self,
        key_item_dict: Dict[str, Any],
        serializer: PPCacheItemSerializer,
        ttl: Optional[int] = config.REDIS_CACHE_TTL,
    ) -> Dict[str, bool]:
        """Write multiple cached entries using Redis.SET.

        Args:
            key_item_dict (Dict[str, Any]): Collection of items to save
            serializer (PPCacheItemSerializer): Serializer to transform item for storage
            ttl (int): Time-to-live for saved items

        Returns:
            Dict of cached keys (key) and the results of the SET commands (value).
            {
                "key1": True,
                "key2": False,
            }
        """
        dict_results = {}
        try:
            for k, v in key_item_dict.items():
                # We are calling set method from this class, which calls Redis.SET
                result = await self.set(k, v, serializer=serializer, ttl=ttl)
                dict_results[k] = result
        except RedisError:
            logger.error("[RedisError] set failed, returning False", exc_info=True)
            dict_results[k] = False
        return dict_results

    @tracer.wrap()
    async def mset_with_pipeline(
        self,
        key_item_dict: Dict[str, Any],
        serializer: PPCacheItemSerializer,
        ttl: Optional[int] = config.REDIS_CACHE_TTL,
    ) -> Dict[str, bool]:
        """Write multiple cached entries using Redis.SET and Pipelining.

        Args:
            key_item_dict (Dict[str, Any]): Collection of items to save
            serializer (PPCacheItemSerializer): Serializer to transform item for storage
            ttl (int): Time-to-live for saved items

        Returns:
            Dict of cached keys (key) and the results of the SET commands (value).
            {
                "key1": True,
                "key2": False,
            }
        """
        dict_results = {}
        falsy_dict_results = dict(
            zip(
                key_item_dict.keys(),
                [False for _ in range(len(key_item_dict.keys()))],
            )
        )
        pipeline = self.client.pipeline(transaction=False)
        try:
            for k, v in key_item_dict.items():
                # We are calling set method from this class, which calls Redis.SET
                await self.set(
                    k,
                    v,
                    serializer=serializer,
                    ttl=ttl,
                    redis_client=pipeline,
                )
            result = await pipeline.execute(raise_on_error=True)
            assert len(result) == len(key_item_dict)
            dict_results = dict(zip(key_item_dict.keys(), result))
        except RedisError:
            logger.error(
                "[RedisError] Pipeline Execute Failed, returning False",
                exc_info=True,
            )
            dict_results = falsy_dict_results
        except AssertionError:
            logger.error(
                "[AssertionError] received mismatched result length", exc_info=True
            )
            dict_results = falsy_dict_results
        finally:
            await pipeline.reset()  # type: ignore[no-untyped-call]
        return dict_results

    @tracer.wrap()
    async def delete(self, *keys: str) -> int:
        """
        Delete an item for this key using Redis.DELETE.

        Args:
            keys: List of keys to delete from the cache

        Returns:
            The number items deleted cache entries.
        """
        # cleanup keys, in case an entry is None.
        keys = tuple(k for k in keys if k)

        if not keys:
            return 0

        try:
            num_affected_entries = await self.client.delete(*keys)
        except RedisError:
            logger.error("[RedisError] delete failed, returning 0", exc_info=True)
            return 0
        return int(num_affected_entries)

    @tracer.wrap()
    async def list(self, pattern: str) -> List[str]:
        """
        Find all keys matching the pattern using Redis.KEYS.

        Args:
            pattern: Wildcard pattern to match keys. (e.g. prefix_*)

        Returns:
            A list of keys matching the pattern
        """
        if not pattern:
            return []

        matching_keys = []
        try:
            async for key in self.client.scan_iter(match=pattern):
                # convert byte-string to str.
                matching_keys.append(str(key, encoding="utf-8"))
        except RedisError:
            logger.error("[RedisError] scan_iter failed, returning []", exc_info=True)
            return []

        return matching_keys

    @tracer.wrap()
    async def delete_all_matching_pattern(self, pattern: str) -> List[str]:
        """
        Find all keys matching the pattern and delete the matching cache entries.

        Example:
            # Bust caches for identity_uuid
            keys_deleted = await r.delete_all_matching_pattern(f"*{identity_uuid}*")

        Args:
            pattern: Wildcard pattern to match keys. (e.g. prefix_*)

        Returns:
            A list of keys matching the pattern.
            Raises ValueError if the number of matched keys does not
            match the number of deleted keys.
        """
        matching_keys = await self.list(pattern=pattern)

        if not matching_keys:
            return []

        num_affected_entries = await self.delete(*matching_keys)

        if num_affected_entries != len(matching_keys):
            logger.warning(
                "[bludgeon] Expected %s deleted keys but DELETE returned: '%s",
                num_affected_entries,
                len(matching_keys),
            )

        return matching_keys


class RedisConnectorFactoryError(Exception):
    """Indicates a working RedisConnector could not be created.

    This includes when ping fails for a fakeredis-based RedisConnector.
    """

    pass


async def redis_connector_factory() -> RedisConnector:
    """Create a redis connector, with fakeredis as backup, to avoid a SPOF."""
    redis_connector = RedisConnector(
        redis_url=config.REDIS_URL, use_redis_cache=config.CACHE_USE_REDIS
    )
    ping = await redis_connector.ping()
    if ping:
        logger.info(
            f"Initialized RedisConnector at {config.REDIS_URL} with use_redis_cache = {config.CACHE_USE_REDIS}"  # noqa: E501
        )
    else:
        logger.error("Switching to use fakeredis.")
        redis_connector = RedisConnector(redis_url="fakeredis", use_redis_cache=False)
        ping = await redis_connector.ping()
        if ping:
            logger.info("Initialized RedisConnector with fakeredis")
        else:
            raise RedisConnectorFactoryError(
                "Failed to initialize RedisConnector with fakeredis"
            )
    return redis_connector