import functools
import hashlib
import io
import logging
import os
import pickle  # nosec
import socket
from collections.abc import Callable
from contextlib import suppress
from typing import ParamSpec, TypeVar, cast

import diskcache
import diskcache.core
from diskcache import Lock
from platformdirs import user_cache_dir
from redis import Redis

LOGGER = logging.getLogger(__name__)
REDIS_CACHE_ENDPOINT = os.environ.get("CACHE_URL")
REDIS_CLIENT: Redis | None = None

P = ParamSpec("P")
R = TypeVar("R")
TFun = Callable[P, R]

if REDIS_CACHE_ENDPOINT:
    try:
        host, _port = REDIS_CACHE_ENDPOINT.split(":", maxsplit=1)
        port = int(_port)

        # Try to connect to redis endpoint
        socket.create_connection((host, port), timeout=2)

        REDIS_CLIENT = cast(
            "Redis",
            Redis(
                host=host,
                port=port,
                ssl=True,
                username="batch-rw-user",
                password=os.environ.get("CACHE_BATCH_RW_USER_PASSWORD"),
            ),
        )
    except (OSError, ValueError):
        REDIS_CLIENT = None

DISK_CACHE = diskcache.Cache(user_cache_dir("fluid-labels", "fluidattacks"))


# 2 weeks
TTL = 604800 * 2

# Whitelist of allowed modules and classes for pickling/unpickling
_PICKLE_ALLOWED_CLASSES: frozenset[tuple[str, str]] = frozenset(
    {
        # Built-in types
        ("builtins", "dict"),
        ("builtins", "list"),
        ("builtins", "set"),
        ("builtins", "frozenset"),
        ("builtins", "tuple"),
        ("builtins", "bytes"),
        ("builtins", "bytearray"),
        # Collections
        ("collections", "OrderedDict"),
        ("collections", "defaultdict"),
        # datetime types
        ("datetime", "datetime"),
        ("datetime", "date"),
        ("datetime", "time"),
        ("datetime", "timedelta"),
    }
)


class RestrictedUnpickler(pickle.Unpickler):  # noqa: S301 - mitigated via _PICKLE_ALLOWED_CLASSES allowlist
    """Unpickler that only allows specific safe classes."""

    def find_class(self, module: str, name: str) -> type:
        if (module, name) in _PICKLE_ALLOWED_CLASSES:
            return cast("type", super().find_class(module, name))

        err_msg = f"Forbidden class in cache: {module}.{name}"
        raise pickle.UnpicklingError(err_msg)

    def load(self) -> object:
        return cast("object", super().load())


def _pickle_loads(data: object) -> object:
    """Deserialize bytes data using RestrictedUnpickler.

    Returns deserialized object if safe, otherwise raises exception
    """
    if not isinstance(data, bytes):
        raise pickle.UnpicklingError

    return RestrictedUnpickler(io.BytesIO(data)).load()


def _pickle_dumps(item: object) -> bytes:
    """Serialize and validate an object using RestrictedUnpickler.

    Returns serialized bytes if safe, otherwise raises exception
    """
    data = pickle.dumps(item)
    RestrictedUnpickler(io.BytesIO(data)).load()
    return data


def make_hashable(item: object) -> str:
    serialized_object = pickle.dumps(item)

    return hashlib.sha256(serialized_object).hexdigest()


def safe_generate_cache_key(
    func: Callable[P, R],
    args: object,
    kwargs: object,
) -> str | None:
    try:
        key = f"{func.__module__}.{func.__name__}-"
        key += str(make_hashable((args, kwargs)))
        return hashlib.sha256(key.encode()).hexdigest()
    except Exception:
        LOGGER.exception("Failed to generate cache key for function: %s", func.__name__)
        return None


def safe_save_to_cache(result: object, cache_key: str, function_name: str) -> None:
    try:
        data = _pickle_dumps(result)
        with suppress(diskcache.core.Timeout), Lock(DISK_CACHE, "dual_cache"):
            DISK_CACHE.set(cache_key, data, expire=TTL, retry=True)  # Save with TTL
        if REDIS_CLIENT:
            REDIS_CLIENT.setex(cache_key, TTL, data)  # Save with TTL
    except Exception:
        LOGGER.exception("Failed to pickle result for function: %s", function_name)


_SERIALIZATION_ERROR = object()


def safe_retrieve_from_cache(cache_key: str) -> object:
    with suppress(Exception):
        # Try to recover from disk cache
        if cached_data := DISK_CACHE.get(cache_key):
            return _pickle_loads(cached_data)

        # Try to recover from Redis cache (ElastiCache) if available
        if REDIS_CLIENT and (cached_data := REDIS_CLIENT.get(cache_key)):
            result = _pickle_loads(cached_data)
            # Save to disk cache
            with Lock(DISK_CACHE, "dual_cache"):
                DISK_CACHE.set(cache_key, cached_data, expire=TTL, retry=True)
            return result
    # Anything other than a cache hit returns a serialization error
    return _SERIALIZATION_ERROR


def dual_cache(func: Callable[P, R]) -> Callable[P, R]:
    @functools.wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        if os.environ.get("SKIP_FA_LABELS_CACHE_USE"):
            return func(*args, **kwargs)

        cache_key = safe_generate_cache_key(func, args, kwargs)
        if cache_key:
            cached_result = safe_retrieve_from_cache(cache_key)
            # None or falsy values might be stored on the cache, so this check is necessary
            if cached_result is not _SERIALIZATION_ERROR:
                return cast("R", cached_result)

        # Run the function and store the result in both caches if serializable
        result = func(*args, **kwargs)
        if cache_key:
            safe_save_to_cache(result, cache_key, func.__name__)
        return result

    return wrapper
