"""Caching for subscription results."""

import asyncio
import hashlib
import pickle
import time
from collections.abc import AsyncGenerator
from dataclasses import dataclass
from functools import wraps
from typing import Any, Dict, Optional


@dataclass
class CacheEntry:
    """A cached subscription result."""

    value: Any
    timestamp: float
    ttl: float

    def is_expired(self) -> bool:
        """Check if cache entry is expired."""
        return time.time() - self.timestamp > self.ttl


class SubscriptionCache:
    """Caches subscription results to reduce load."""

    def __init__(self):
        self._cache: Dict[str, CacheEntry] = {}
        self._locks: Dict[str, asyncio.Lock] = {}
        self._cleanup_task: Optional[asyncio.Task] = None

    async def start(self):
        """Start cache cleanup task."""
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    async def stop(self):
        """Stop cache cleanup task."""
        if self._cleanup_task:
            self._cleanup_task.cancel()
            await asyncio.gather(self._cleanup_task, return_exceptions=True)

    def _make_key(self, func_name: str, args: Dict[str, Any]) -> str:
        """Generate cache key from function and arguments."""
        key_data = {"func": func_name, "args": args}
        key_bytes = pickle.dumps(key_data, protocol=pickle.HIGHEST_PROTOCOL)
        return hashlib.sha256(key_bytes).hexdigest()

    async def get_or_generate(
        self, key: str, generator: AsyncGenerator, ttl: float
    ) -> AsyncGenerator[Any, None]:
        """Get cached values or generate new ones."""
        # Check cache
        if key in self._cache:
            entry = self._cache[key]
            if not entry.is_expired():
                # Return cached value
                yield entry.value
                return

        # Ensure only one generator per key
        if key not in self._locks:
            self._locks[key] = asyncio.Lock()

        async with self._locks[key]:
            # Double-check cache
            if key in self._cache and not self._cache[key].is_expired():
                yield self._cache[key].value
                return

            # Generate new value
            async for value in generator:
                # Cache the value
                self._cache[key] = CacheEntry(
                    value=value, timestamp=time.time(), ttl=ttl
                )
                yield value

    async def _cleanup_loop(self):
        """Periodically clean expired entries."""
        while True:
            try:
                await asyncio.sleep(60)  # Every minute

                expired = []
                for key, entry in self._cache.items():
                    if entry.is_expired():
                        expired.append(key)

                for key in expired:
                    del self._cache[key]
                    if key in self._locks:
                        del self._locks[key]

            except asyncio.CancelledError:
                break
            except Exception as e:
                print(f"Cache cleanup error: {e}")


def cache(ttl: float = 5.0):
    """
    Decorator to cache subscription results.

    Usage:
        @subscription
        @cache(ttl=10)  # Cache for 10 seconds
        async def expensive_stats(info):
            ...
    """

    def decorator(func):
        func._cache_ttl = ttl

        @wraps(func)
        async def wrapper(info, **kwargs):
            # Get cache from context
            sub_cache = None
            if hasattr(info, "context") and info.context:
                sub_cache = info.context.get("subscription_cache")

            if not sub_cache:
                # No cache, execute directly
                async for value in func(info, **kwargs):
                    yield value
                return

            # Generate cache key
            cache_key = sub_cache._make_key(func.__name__, kwargs)

            # Use cached values or generate
            generator = func(info, **kwargs)
            async for value in sub_cache.get_or_generate(cache_key, generator, ttl):
                yield value

        return wrapper

    return decorator
