# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/cache/resource_cache.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Resource Cache Implementation.
This module implements a simple in-memory cache with TTL expiration for caching
resource content in the MCP Gateway. Features:
- TTL-based expiration
- Maximum size limit with LRU eviction
- Thread-safe operations

Examples:
    >>> from mcpgateway.cache.resource_cache import ResourceCache
    >>> from unittest.mock import patch
    >>> cache = ResourceCache(max_size=2, ttl=1)
    >>> cache.set('a', 1)
    >>> cache.get('a')
    1

    Test TTL expiration using mocked time (no actual sleep):

    >>> with patch("time.time") as mock_time:
    ...     mock_time.return_value = 1000
    ...     cache2 = ResourceCache(max_size=2, ttl=1)
    ...     cache2.set('x', 100)
    ...     cache2.get('x')  # Before expiration
    ...     mock_time.return_value = 1002  # Advance past TTL
    ...     cache2.get('x') is None  # After expiration
    100
    True

    Test LRU eviction:

    >>> cache.set('a', 1)
    >>> cache.set('b', 2)
    >>> cache.set('c', 3)  # LRU eviction
    >>> sorted(cache._cache.keys())
    ['b', 'c']
    >>> cache.delete('b')
    >>> cache.get('b') is None
    True
    >>> cache.clear()
    >>> cache.get('a') is None
    True
"""

# Standard
import asyncio
from collections import OrderedDict
from dataclasses import dataclass
import heapq
import threading
import time
from typing import Any, Optional

# First-Party
from mcpgateway.services.logging_service import LoggingService

# Initialize logging service first
logging_service = LoggingService()
logger = logging_service.get_logger(__name__)


@dataclass
class CacheEntry:
    """Cache entry with expiration."""

    value: Any
    expires_at: float


class ResourceCache:
    """
    Resource content cache with TTL expiration.

    Attributes:
        max_size: Maximum number of entries
        ttl: Time-to-live in seconds
        _cache: Cache storage
        _lock: Threading lock for thread safety

    Examples:
        >>> from mcpgateway.cache.resource_cache import ResourceCache
        >>> from unittest.mock import patch
        >>> cache = ResourceCache(max_size=2, ttl=1)
        >>> cache.set('a', 1)
        >>> cache.get('a')
        1

        Test TTL expiration using mocked time (no actual sleep):

        >>> with patch("time.time") as mock_time:
        ...     mock_time.return_value = 1000
        ...     cache2 = ResourceCache(max_size=2, ttl=1)
        ...     cache2.set('x', 100)
        ...     cache2.get('x')  # Before expiration
        ...     mock_time.return_value = 1002  # Advance past TTL
        ...     cache2.get('x') is None  # After expiration
        100
        True

        Test LRU eviction:

        >>> cache.set('a', 1)
        >>> cache.set('b', 2)
        >>> cache.set('c', 3)  # LRU eviction
        >>> sorted(cache._cache.keys())
        ['b', 'c']
        >>> cache.delete('b')
        >>> cache.get('b') is None
        True
        >>> cache.clear()
        >>> cache.get('a') is None
        True
    """

    def __init__(self, max_size: int = 1000, ttl: int = 3600):
        """Initialize cache.

        Args:
            max_size: Maximum number of entries
            ttl: Time-to-live in seconds
        """
        self.max_size = max_size
        self.ttl = ttl
        self._cache: OrderedDict[str, CacheEntry] = OrderedDict()
        # Use a threading lock for thread-safe operations across sync methods
        # and the background cleanup thread.
        self._lock = threading.Lock()
        # Min-heap of (expires_at, key) for efficient expiration cleanup
        self._expiry_heap: list[tuple[float, str]] = []
        self._cleanup_task: Optional[asyncio.Task] = None

    async def initialize(self) -> None:
        """Initialize cache service."""
        logger.info("Initializing resource cache")
        # Start cleanup task and store reference to prevent GC
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    async def shutdown(self) -> None:
        """Shutdown cache service."""
        logger.info("Shutting down resource cache")
        if self._cleanup_task and not self._cleanup_task.done():
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
        self.clear()

    def get(self, key: str) -> Optional[Any]:
        """
        Get value from cache.

        Args:
            key: Cache key

        Returns:
            Cached value or None if not found/expired

        Examples:
            >>> from mcpgateway.cache.resource_cache import ResourceCache
            >>> from unittest.mock import patch

            >>> # Normal get
            >>> cache = ResourceCache(max_size=2, ttl=1)
            >>> cache.set('a', 1)
            >>> cache.get('a')
            1

            >>> # Test expiration deterministically using mock time
            >>> with patch("time.time") as mock_time:
            ...     mock_time.return_value = 1000
            ...     short_cache = ResourceCache(max_size=2, ttl=0.1)
            ...     short_cache.set('b', 2)
            ...     short_cache.get('b')
            ...     # Advance time past TTL
            ...     mock_time.return_value = 1000.2
            ...     short_cache.get('b') is None
            2
            True
        """
        with self._lock:
            if key not in self._cache:
                return None

            entry = self._cache[key]
            now = time.time()

            # Check expiration
            if now > entry.expires_at:
                del self._cache[key]
                return None

            self._cache.move_to_end(key)

            return entry.value

    def set(self, key: str, value: Any) -> None:
        """
        Set value in cache.

        Args:
            key: Cache key
            value: Value to cache

        Examples:
            >>> from mcpgateway.cache.resource_cache import ResourceCache
            >>> cache = ResourceCache(max_size=2, ttl=1)
            >>> cache.set('a', 1)
            >>> cache.get('a')
            1
        """
        expires_at = time.time() + self.ttl
        with self._lock:
            if key in self._cache:
                self._cache.move_to_end(key)
            elif len(self._cache) >= self.max_size:
                # Evict LRU
                self._cache.popitem(last=False)

            # Add / update entry
            self._cache[key] = CacheEntry(value=value, expires_at=expires_at)
            # Push expiry into heap; stale heap entries are ignored later
            heapq.heappush(self._expiry_heap, (expires_at, key))

    def delete(self, key: str) -> None:
        """
        Delete value from cache.

        Args:
            key: Cache key to delete

        Examples:
            >>> from mcpgateway.cache.resource_cache import ResourceCache
            >>> cache = ResourceCache()
            >>> cache.set('a', 1)
            >>> cache.delete('a')
            >>> cache.get('a') is None
            True
        """
        with self._lock:
            self._cache.pop(key, None)
            # We don't remove entries from the heap here; they'll be ignored
            # by the cleanup when popped if missing or timestamp differs.

    def clear(self) -> None:
        """
        Clear all cached entries.

        Examples:
            >>> from mcpgateway.cache.resource_cache import ResourceCache
            >>> cache = ResourceCache()
            >>> cache.set('a', 1)
            >>> cache.clear()
            >>> cache.get('a') is None
            True
        """
        with self._lock:
            self._cache.clear()
            self._expiry_heap.clear()

    async def _cleanup_loop(self) -> None:
        """Background task to clean expired entries efficiently.

        Uses a min-heap of expiration timestamps to avoid scanning the
        entire cache on each run. The actual cleanup work runs under the
        same threading lock as sync methods by delegating to a thread via
        `asyncio.to_thread` so we don't block the event loop.
        """

        async def _run_once() -> None:
            """Execute a single cleanup pass, catching and logging any errors."""
            try:
                await asyncio.to_thread(self._cleanup_once)
            except Exception as e:
                logger.error(f"Cache cleanup error: {e}")

        while True:
            await _run_once()
            await asyncio.sleep(60)  # Run every minute

    def _cleanup_once(self) -> None:
        """Synchronous cleanup routine executed in a thread.

        Pops entries from the expiry heap until the next non-expired
        timestamp is reached. Each popped entry is validated against
        the current cache entry to avoid removing updated entries.
        Also compacts the heap if it grows too large relative to cache size.
        """
        now = time.time()
        removed = 0
        needs_compaction = False

        with self._lock:
            while self._expiry_heap and self._expiry_heap[0][0] <= now:
                expires_at, key = heapq.heappop(self._expiry_heap)
                entry = self._cache.get(key)
                # If entry is present and timestamps match, remove it
                if entry is not None and entry.expires_at == expires_at:
                    del self._cache[key]
                    removed += 1

            # Check if heap needs compaction (done outside lock)
            needs_compaction = len(self._expiry_heap) > 2 * self.max_size

        if removed:
            logger.debug(f"Cleaned {removed} expired cache entries")

        # Compact heap outside the main lock to minimize contention
        if needs_compaction:
            self._compact_heap()

    def _compact_heap(self) -> None:
        """Rebuild the expiry heap with only valid (current) entries.

        Called when heap grows too large due to stale entries from
        key updates or deletions. Minimizes lock contention by doing
        the O(n) heapify outside the lock.
        """
        # Snapshot current entries under lock (fast dict iteration)
        with self._lock:
            entries = [(entry.expires_at, key) for key, entry in self._cache.items()]
            old_size = len(self._expiry_heap)
            # Track max expiry in snapshot to identify entries added during compaction
            max_snapshot_expiry = max((e[0] for e in entries), default=0.0)

        # Build heap outside lock - O(n) work doesn't block get/set
        heapq.heapify(entries)

        # Swap back under lock, preserving entries added during compaction
        with self._lock:
            # Keep heap entries with expiry > max_snapshot_expiry (added via set() during compaction)
            new_entries = [(exp, k) for exp, k in self._expiry_heap if exp > max_snapshot_expiry]
            self._expiry_heap = entries
            for entry in new_entries:
                heapq.heappush(self._expiry_heap, entry)

        logger.debug(f"Compacted expiry heap: {old_size} -> {len(self._expiry_heap)} entries")

    def __len__(self) -> int:
        """
        Get the number of entries in cache.

        Args:
            None

        Returns:
            int: Number of entries in cache

        Examples:
            >>> from mcpgateway.cache.resource_cache import ResourceCache
            >>> cache = ResourceCache(max_size=2, ttl=1)
            >>> cache.set('a', 1)
            >>> len(cache)
            1
        """
        with self._lock:
            return len(self._cache)
