Trixy3/trixy_core/assets/asset_cache.py

"""
Asset Caching System for Trixy Application

This module provides a comprehensive caching system for assets to improve
performance and reduce I/O operations in the Trixy asset management system.

Features:
- Memory-efficient LRU cache with configurable size limits
- Compressed caching to reduce memory usage
- Cache persistence across application restarts
- Thread-safe operations for multi-satellite support
- Cache warming and preloading capabilities
- Detailed cache statistics and monitoring
- Cache policies and eviction strategies
- Integration with asset validation system

Cache Types:
- Memory cache: Fast in-memory storage for frequently accessed assets
- Disk cache: Persistent storage for large or less frequently used assets
- Compressed cache: Reduced memory footprint using compression
- Metadata cache: Quick access to asset metadata without full loading

Usage:
    from trixy_core.assets.asset_cache import AssetCache, CachePolicy
    
    # Create cache with policy
    policy = CachePolicy(max_entries=100, max_memory_bytes=50*1024*1024)
    cache = AssetCache(policy=policy)
    
    # Store and retrieve assets
    cache.put("audio/success.wav", asset_data, metadata)
    cached_data, metadata = cache.get("audio/success.wav")
    
    # Cache management
    cache.clear()
    stats = cache.get_stats()
"""

import os
import json
import gzip
import pickle
import threading
import time
import weakref
from pathlib import Path
from typing import Dict, Any, Optional, List, Set, Union, Tuple, NamedTuple
from dataclasses import dataclass, field
from enum import Enum
import hashlib
from collections import OrderedDict
import tempfile
import shutil


class CacheEntryStatus(Enum):
    """Status of a cache entry."""
    LOADING = "loading"
    CACHED = "cached"
    EXPIRED = "expired"
    INVALID = "invalid"
    EVICTED = "evicted"


class CacheEvictionStrategy(Enum):
    """Cache eviction strategies."""
    LRU = "lru"  # Least Recently Used
    LFU = "lfu"  # Least Frequently Used
    FIFO = "fifo"  # First In, First Out
    TTL = "ttl"  # Time To Live based
    SIZE = "size"  # Based on entry size


@dataclass
class CachePolicy:
    """Configuration policy for asset caching."""
    max_entries: int = 100  # Maximum number of cached entries
    max_memory_bytes: int = 100 * 1024 * 1024  # Maximum memory usage (100MB)
    max_entry_size: int = 10 * 1024 * 1024  # Maximum size per entry (10MB)
    default_ttl_seconds: float = 3600.0  # Default time-to-live (1 hour)
    eviction_strategy: CacheEvictionStrategy = CacheEvictionStrategy.LRU
    compress: bool = False  # Whether to compress cached data
    compression_threshold: int = 1024  # Minimum size to compress (bytes)
    compression_level: int = 6  # Compression level (1-9, 6 is good balance)
    enable_persistence: bool = True  # Enable disk persistence
    persistence_path: Optional[str] = None  # Path for cache persistence
    enable_stats: bool = True  # Enable statistics collection
    warmup_on_startup: bool = False  # Warm cache on startup
    preload_patterns: List[str] = field(default_factory=list)  # Patterns to preload


class CacheEntry:
    """Represents a single cache entry."""
    
    def __init__(
        self,
        key: str,
        data: Any,
        metadata: Optional[Dict[str, Any]] = None,
        ttl_seconds: Optional[float] = None,
        compressed: bool = False
    ):
        self.key = key
        self.data = data
        self.metadata = metadata or {}
        self.compressed = compressed
        self.created_time = time.time()
        self.last_access_time = self.created_time
        self.access_count = 0
        self.size_bytes = self._calculate_size()
        self.status = CacheEntryStatus.CACHED
        
        # TTL handling
        if ttl_seconds is not None:
            self.expires_at = self.created_time + ttl_seconds
        else:
            self.expires_at = None
        
        # Compression
        if compressed and not self.compressed:
            self._compress_data()
    
    def _calculate_size(self) -> int:
        """Calculate the memory size of this cache entry."""
        try:
            # Estimate size using pickle serialization
            return len(pickle.dumps(self.data)) + len(pickle.dumps(self.metadata))
        except Exception:
            # Fallback estimation
            if hasattr(self.data, '__len__'):
                return len(self.data)
            return 1024  # Default estimate
    
    def _compress_data(self):
        """Compress the cached data."""
        try:
            if not self.compressed:
                compressed_data = gzip.compress(pickle.dumps(self.data))
                self.data = compressed_data
                self.compressed = True
                self.size_bytes = len(compressed_data) + len(pickle.dumps(self.metadata))
        except Exception as e:
            # Compression failed, keep original data
            pass
    
    def _decompress_data(self):
        """Decompress the cached data."""
        try:
            if self.compressed:
                decompressed_data = pickle.loads(gzip.decompress(self.data))
                return decompressed_data
            return self.data
        except Exception:
            # Decompression failed, return None to indicate error
            return None
    
    def access(self) -> Any:
        """Access the cached data, updating access statistics."""
        self.last_access_time = time.time()
        self.access_count += 1
        
        if self.compressed:
            return self._decompress_data()
        return self.data
    
    def is_expired(self) -> bool:
        """Check if this cache entry has expired."""
        if self.expires_at is None:
            return False
        return time.time() > self.expires_at
    
    def is_valid(self) -> bool:
        """Check if this cache entry is valid."""
        return self.status == CacheEntryStatus.CACHED and not self.is_expired()
    
    def update_metadata(self, metadata: Dict[str, Any]):
        """Update the metadata for this entry."""
        self.metadata.update(metadata)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert cache entry to dictionary for serialization."""
        return {
            'key': self.key,
            'metadata': self.metadata,
            'compressed': self.compressed,
            'created_time': self.created_time,
            'last_access_time': self.last_access_time,
            'access_count': self.access_count,
            'size_bytes': self.size_bytes,
            'expires_at': self.expires_at,
            'status': self.status.value
        }


@dataclass
class CacheStats:
    """Cache statistics."""
    hits: int = 0
    misses: int = 0
    evictions: int = 0
    expirations: int = 0
    entries: int = 0
    memory_usage: int = 0
    compression_ratio: float = 0.0
    hit_rate: float = 0.0
    average_access_time: float = 0.0
    
    def update_hit_rate(self):
        """Update the hit rate calculation."""
        total = self.hits + self.misses
        if total > 0:
            self.hit_rate = self.hits / total
        else:
            self.hit_rate = 0.0


def pprint(message: str) -> None:
    """
    Asset cache logging function.
    """
    print(f"[ASSET_CACHE] {message}")


class AssetCache:
    """
    High-performance asset caching system with configurable policies.
    
    Provides memory-efficient caching with compression, persistence,
    and multiple eviction strategies for optimal performance.
    """
    
    def __init__(
        self,
        policy: Optional[CachePolicy] = None,
        name: str = "default"
    ):
        """
        Initialize the asset cache.
        
        Args:
            policy: Cache policy configuration
            name: Name identifier for this cache instance
        """
        self.policy = policy or CachePolicy()
        self.name = name
        
        # Cache storage
        self._cache: OrderedDict[str, CacheEntry] = OrderedDict()
        self._lock = threading.RLock()
        
        # Statistics
        self._stats = CacheStats() if self.policy.enable_stats else None
        
        # Persistence
        self._persistence_path = None
        if self.policy.enable_persistence:
            self._setup_persistence()
        
        # Startup time
        self._startup_time = time.time()
        
        pprint(f"Asset cache '{name}' initialized with policy: "
               f"max_entries={policy.max_entries}, "
               f"max_memory={policy.max_memory_bytes//1024//1024}MB, "
               f"compression={policy.compress}")
        
        # Load persisted cache if available
        if self.policy.enable_persistence:
            self._load_persisted_cache()
        
        # Warm up cache if configured
        if self.policy.warmup_on_startup:
            self._warmup_cache()
    
    def _setup_persistence(self):
        """Setup cache persistence directory."""
        if self.policy.persistence_path:
            self._persistence_path = Path(self.policy.persistence_path)
        else:
            # Use temporary directory by default
            cache_dir = Path(tempfile.gettempdir()) / "trixy_asset_cache"
            self._persistence_path = cache_dir / self.name
        
        try:
            self._persistence_path.mkdir(parents=True, exist_ok=True)
            pprint(f"Cache persistence directory: {self._persistence_path}")
        except Exception as e:
            pprint(f"Warning: Could not create persistence directory: {e}")
            self._persistence_path = None
    
    def _load_persisted_cache(self):
        """Load cache entries from persistence."""
        if not self._persistence_path:
            return
        
        try:
            metadata_file = self._persistence_path / "cache_metadata.json"
            if metadata_file.exists():
                with open(metadata_file, 'r') as f:
                    persisted_metadata = json.load(f)
                
                loaded_count = 0
                for entry_data in persisted_metadata.get('entries', []):
                    try:
                        key = entry_data['key']
                        data_file = self._persistence_path / f"{self._key_to_filename(key)}.dat"
                        
                        if data_file.exists():
                            with open(data_file, 'rb') as f:
                                data = pickle.load(f)
                            
                            # Recreate cache entry
                            entry = CacheEntry(
                                key=key,
                                data=data,
                                metadata=entry_data.get('metadata', {}),
                                compressed=entry_data.get('compressed', False)
                            )
                            
                            # Restore entry properties
                            entry.created_time = entry_data.get('created_time', time.time())
                            entry.last_access_time = entry_data.get('last_access_time', time.time())
                            entry.access_count = entry_data.get('access_count', 0)
                            entry.size_bytes = entry_data.get('size_bytes', entry.size_bytes)
                            entry.expires_at = entry_data.get('expires_at')
                            
                            # Check if still valid
                            if not entry.is_expired():
                                self._cache[key] = entry
                                loaded_count += 1
                    except Exception as entry_error:
                        pprint(f"Warning: Could not load cache entry {key}: {entry_error}")
                
                pprint(f"Loaded {loaded_count} entries from cache persistence")
                
        except Exception as e:
            pprint(f"Warning: Could not load persisted cache: {e}")
    
    def _persist_cache(self):
        """Persist cache entries to disk."""
        if not self._persistence_path or not self.policy.enable_persistence:
            return
        
        try:
            # Save metadata
            metadata = {
                'cache_name': self.name,
                'created_time': self._startup_time,
                'last_saved': time.time(),
                'entries': []
            }
            
            # Save individual entries
            with self._lock:
                for key, entry in self._cache.items():
                    if entry.is_valid():
                        # Save entry data
                        data_file = self._persistence_path / f"{self._key_to_filename(key)}.dat"
                        with open(data_file, 'wb') as f:
                            data_to_save = entry.data
                            if entry.compressed:
                                # Data is already compressed
                                data_to_save = entry.data
                            pickle.dump(data_to_save, f)
                        
                        # Add to metadata
                        metadata['entries'].append(entry.to_dict())
            
            # Save metadata
            metadata_file = self._persistence_path / "cache_metadata.json"
            with open(metadata_file, 'w') as f:
                json.dump(metadata, f, indent=2)
            
            pprint(f"Persisted {len(metadata['entries'])} cache entries")
            
        except Exception as e:
            pprint(f"Warning: Could not persist cache: {e}")
    
    def _key_to_filename(self, key: str) -> str:
        """Convert cache key to safe filename."""
        # Create a hash of the key for safe filename
        return hashlib.md5(key.encode()).hexdigest()
    
    def _warmup_cache(self):
        """Warm up the cache with preload patterns."""
        if not self.policy.preload_patterns:
            return
        
        pprint(f"Warming up cache with {len(self.policy.preload_patterns)} patterns")
        # This would integrate with AssetManager to preload assets
        # For now, just log the intent
    
    def get(self, key: str) -> Tuple[Optional[Any], Optional[Dict[str, Any]]]:
        """
        Get an item from the cache.
        
        Args:
            key: Cache key
            
        Returns:
            Tuple of (data, metadata) if found, (None, None) if not found
        """
        with self._lock:
            entry = self._cache.get(key)
            
            if entry is None:
                self._record_miss()
                return None, None
            
            if not entry.is_valid():
                # Entry is expired or invalid
                self._remove_entry(key, reason="expired")
                self._record_miss()
                return None, None
            
            # Move to end for LRU
            if self.policy.eviction_strategy == CacheEvictionStrategy.LRU:
                self._cache.move_to_end(key)
            
            self._record_hit()
            data = entry.access()
            
            if data is None:
                # Decompression or access failed
                self._remove_entry(key, reason="corrupted")
                self._record_miss()
                return None, None
            
            return data, entry.metadata.copy()
    
    def put(
        self,
        key: str,
        data: Any,
        metadata: Optional[Dict[str, Any]] = None,
        ttl_seconds: Optional[float] = None
    ) -> bool:
        """
        Put an item into the cache.
        
        Args:
            key: Cache key
            data: Data to cache
            metadata: Optional metadata
            ttl_seconds: Time-to-live in seconds
            
        Returns:
            True if successfully cached, False otherwise
        """
        if not key or data is None:
            return False
        
        # Check if data is too large
        try:
            estimated_size = len(pickle.dumps(data))
            if estimated_size > self.policy.max_entry_size:
                pprint(f"Entry too large to cache: {key} ({estimated_size} bytes)")
                return False
        except Exception:
            # If we can't serialize it, we can't cache it
            return False
        
        with self._lock:
            # Determine if we should compress
            should_compress = (
                self.policy.compress and
                estimated_size >= self.policy.compression_threshold
            )
            
            # Use provided TTL or default
            ttl = ttl_seconds if ttl_seconds is not None else self.policy.default_ttl_seconds
            
            # Create cache entry
            entry = CacheEntry(
                key=key,
                data=data,
                metadata=metadata,
                ttl_seconds=ttl,
                compressed=should_compress
            )
            
            # Check if we need to make room
            self._ensure_capacity_for_entry(entry)
            
            # Add to cache
            if key in self._cache:
                # Update existing entry
                old_entry = self._cache[key]
                self._update_stats_for_removal(old_entry)
            
            self._cache[key] = entry
            
            # Move to end for LRU
            if self.policy.eviction_strategy == CacheEvictionStrategy.LRU:
                self._cache.move_to_end(key)
            
            self._update_stats_for_addition(entry)
            
            pprint(f"Cached asset: {key} ({entry.size_bytes} bytes, compressed: {entry.compressed})")
            return True
    
    def remove(self, key: str) -> bool:
        """
        Remove an item from the cache.
        
        Args:
            key: Cache key to remove
            
        Returns:
            True if item was removed, False if not found
        """
        with self._lock:
            return self._remove_entry(key, reason="manual")
    
    def _remove_entry(self, key: str, reason: str = "unknown") -> bool:
        """Remove an entry from the cache."""
        entry = self._cache.pop(key, None)
        if entry:
            self._update_stats_for_removal(entry)
            
            if reason == "expired":
                self._record_expiration()
            elif reason in ["eviction", "capacity"]:
                self._record_eviction()
            
            return True
        return False
    
    def clear(self) -> int:
        """
        Clear all items from the cache.
        
        Returns:
            Number of items removed
        """
        with self._lock:
            count = len(self._cache)
            self._cache.clear()
            
            if self._stats:
                self._stats.entries = 0
                self._stats.memory_usage = 0
            
            pprint(f"Cache cleared: {count} entries removed")
            return count
    
    def _ensure_capacity_for_entry(self, entry: CacheEntry):
        """Ensure there's capacity for a new entry."""
        # Check entry count limit
        while len(self._cache) >= self.policy.max_entries:
            self._evict_one_entry()
        
        # Check memory limit
        current_memory = self._calculate_memory_usage()
        while (current_memory + entry.size_bytes) > self.policy.max_memory_bytes:
            if not self._evict_one_entry():
                # Couldn't evict anything, can't add this entry
                break
            current_memory = self._calculate_memory_usage()
    
    def _evict_one_entry(self) -> bool:
        """Evict one entry based on the eviction strategy."""
        if not self._cache:
            return False
        
        key_to_evict = None
        
        if self.policy.eviction_strategy == CacheEvictionStrategy.LRU:
            # Remove least recently used (first item in OrderedDict)
            key_to_evict = next(iter(self._cache))
        
        elif self.policy.eviction_strategy == CacheEvictionStrategy.LFU:
            # Remove least frequently used
            min_access_count = float('inf')
            for key, entry in self._cache.items():
                if entry.access_count < min_access_count:
                    min_access_count = entry.access_count
                    key_to_evict = key
        
        elif self.policy.eviction_strategy == CacheEvictionStrategy.FIFO:
            # Remove oldest entry
            oldest_time = float('inf')
            for key, entry in self._cache.items():
                if entry.created_time < oldest_time:
                    oldest_time = entry.created_time
                    key_to_evict = key
        
        elif self.policy.eviction_strategy == CacheEvictionStrategy.TTL:
            # Remove entry with earliest expiration
            earliest_expiry = float('inf')
            for key, entry in self._cache.items():
                if entry.expires_at and entry.expires_at < earliest_expiry:
                    earliest_expiry = entry.expires_at
                    key_to_evict = key
        
        elif self.policy.eviction_strategy == CacheEvictionStrategy.SIZE:
            # Remove largest entry
            max_size = 0
            for key, entry in self._cache.items():
                if entry.size_bytes > max_size:
                    max_size = entry.size_bytes
                    key_to_evict = key
        
        if key_to_evict:
            self._remove_entry(key_to_evict, reason="eviction")
            return True
        
        return False
    
    def _calculate_memory_usage(self) -> int:
        """Calculate total memory usage of the cache."""
        return sum(entry.size_bytes for entry in self._cache.values())
    
    def cleanup_expired(self) -> int:
        """
        Remove expired entries from the cache.
        
        Returns:
            Number of expired entries removed
        """
        with self._lock:
            expired_keys = [
                key for key, entry in self._cache.items()
                if entry.is_expired()
            ]
            
            for key in expired_keys:
                self._remove_entry(key, reason="expired")
            
            if expired_keys:
                pprint(f"Cleaned up {len(expired_keys)} expired cache entries")
            
            return len(expired_keys)
    
    def get_stats(self) -> Dict[str, Any]:
        """
        Get cache statistics.
        
        Returns:
            Dictionary containing cache statistics
        """
        if not self._stats:
            return {}
        
        with self._lock:
            self._stats.entries = len(self._cache)
            self._stats.memory_usage = self._calculate_memory_usage()
            self._stats.update_hit_rate()
            
            # Calculate compression ratio
            if self.policy.compress:
                total_compressed = sum(
                    entry.size_bytes for entry in self._cache.values()
                    if entry.compressed
                )
                total_uncompressed = sum(
                    len(pickle.dumps(entry.data)) for entry in self._cache.values()
                    if not entry.compressed
                )
                
                if total_uncompressed > 0:
                    self._stats.compression_ratio = total_compressed / total_uncompressed
            
            return {
                'hits': self._stats.hits,
                'misses': self._stats.misses,
                'hit_rate': self._stats.hit_rate,
                'evictions': self._stats.evictions,
                'expirations': self._stats.expirations,
                'entries': self._stats.entries,
                'memory_usage_bytes': self._stats.memory_usage,
                'memory_usage_mb': self._stats.memory_usage / (1024 * 1024),
                'compression_ratio': self._stats.compression_ratio,
                'uptime_seconds': time.time() - self._startup_time,
                'policy': {
                    'max_entries': self.policy.max_entries,
                    'max_memory_mb': self.policy.max_memory_bytes / (1024 * 1024),
                    'eviction_strategy': self.policy.eviction_strategy.value,
                    'compression_enabled': self.policy.compress,
                }
            }
    
    def get_entry_info(self, key: str) -> Optional[Dict[str, Any]]:
        """
        Get information about a specific cache entry.
        
        Args:
            key: Cache key
            
        Returns:
            Dictionary with entry information or None if not found
        """
        with self._lock:
            entry = self._cache.get(key)
            if entry:
                return {
                    'key': entry.key,
                    'size_bytes': entry.size_bytes,
                    'compressed': entry.compressed,
                    'created_time': entry.created_time,
                    'last_access_time': entry.last_access_time,
                    'access_count': entry.access_count,
                    'expires_at': entry.expires_at,
                    'is_expired': entry.is_expired(),
                    'is_valid': entry.is_valid(),
                    'status': entry.status.value,
                    'metadata': entry.metadata.copy()
                }
        return None
    
    def list_keys(self) -> List[str]:
        """
        Get a list of all cache keys.
        
        Returns:
            List of cache keys
        """
        with self._lock:
            return list(self._cache.keys())
    
    def _record_hit(self):
        """Record a cache hit."""
        if self._stats:
            self._stats.hits += 1
    
    def _record_miss(self):
        """Record a cache miss."""
        if self._stats:
            self._stats.misses += 1
    
    def _record_eviction(self):
        """Record a cache eviction."""
        if self._stats:
            self._stats.evictions += 1
    
    def _record_expiration(self):
        """Record a cache expiration."""
        if self._stats:
            self._stats.expirations += 1
    
    def _update_stats_for_addition(self, entry: CacheEntry):
        """Update statistics when adding an entry."""
        # Stats are updated in real-time through other methods
        pass
    
    def _update_stats_for_removal(self, entry: CacheEntry):
        """Update statistics when removing an entry."""
        # Stats are updated in real-time through other methods
        pass
    
    def shutdown(self):
        """Shutdown the cache and save persistence."""
        pprint(f"Shutting down asset cache '{self.name}'...")
        
        # Persist cache if enabled
        if self.policy.enable_persistence:
            self._persist_cache()
        
        # Clear cache
        with self._lock:
            self._cache.clear()
        
        pprint(f"Asset cache '{self.name}' shutdown complete")
    
    def __enter__(self):
        """Context manager entry."""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.shutdown()


# Convenience functions

def create_memory_cache(
    max_entries: int = 100,
    max_memory_mb: int = 100,
    enable_compression: bool = False
) -> AssetCache:
    """
    Create a simple memory-only cache.
    
    Args:
        max_entries: Maximum number of entries
        max_memory_mb: Maximum memory usage in MB
        enable_compression: Whether to enable compression
        
    Returns:
        AssetCache instance
    """
    policy = CachePolicy(
        max_entries=max_entries,
        max_memory_bytes=max_memory_mb * 1024 * 1024,
        compress=enable_compression,
        enable_persistence=False
    )
    
    return AssetCache(policy=policy, name="memory")


def create_persistent_cache(
    cache_dir: str,
    max_entries: int = 500,
    max_memory_mb: int = 200,
    enable_compression: bool = True
) -> AssetCache:
    """
    Create a cache with disk persistence.
    
    Args:
        cache_dir: Directory for cache persistence
        max_entries: Maximum number of entries
        max_memory_mb: Maximum memory usage in MB
        enable_compression: Whether to enable compression
        
    Returns:
        AssetCache instance
    """
    policy = CachePolicy(
        max_entries=max_entries,
        max_memory_bytes=max_memory_mb * 1024 * 1024,
        compress=enable_compression,
        enable_persistence=True,
        persistence_path=cache_dir,
        warmup_on_startup=True
    )
    
    return AssetCache(policy=policy, name="persistent")


# Module exports
__all__ = [
    'AssetCache',
    'CachePolicy',
    'CacheEntry',
    'CacheEntryStatus',
    'CacheEvictionStrategy',
    'CacheStats',
    'create_memory_cache',
    'create_persistent_cache',
    'pprint'
]