src/infrastructure/query_cache.py

"""
Query result caching system for database optimization.

This module provides intelligent caching of database query results
to reduce database load and improve response times.
"""

import hashlib
from typing import Any, Dict, List, Optional, Callable, TypeVar, Union
from datetime import datetime, timedelta
import asyncio
from functools import wraps

from sqlalchemy.sql import Select
from sqlalchemy.ext.asyncio import AsyncSession

from src.core import get_logger
from src.services.cache_service import cache_service
from src.core.json_utils import dumps, loads

logger = get_logger(__name__)

T = TypeVar('T')


class QueryCache:
    """
    Intelligent query result caching system.
    
    Features:
    - Automatic cache key generation
    - Configurable TTL per query type
    - Cache invalidation strategies
    - Performance metrics
    """
    
    def __init__(self):
        """Initialize query cache."""
        self._cache = cache_service
        self._ttl_config = {
            # Table-specific TTLs (in seconds)
            "investigations": 300,      # 5 minutes
            "contracts": 3600,          # 1 hour
            "users": 1800,              # 30 minutes
            "anomalies": 600,           # 10 minutes
            "agent_messages": 120,      # 2 minutes
            "chat_sessions": 60,        # 1 minute
            "default": 300              # 5 minutes default
        }
        
        # Cache statistics
        self._stats = {
            "hits": 0,
            "misses": 0,
            "invalidations": 0,
            "errors": 0
        }
    
    def _generate_cache_key(
        self,
        query: Union[str, Select],
        params: Optional[Dict[str, Any]] = None,
        prefix: str = "query"
    ) -> str:
        """Generate a unique cache key for a query."""
        # Convert query to string
        query_str = str(query.compile(compile_kwargs={"literal_binds": True})) if hasattr(query, 'compile') else str(query)
        
        # Include parameters in key
        if params:
            params_str = dumps(sorted(params.items()))
        else:
            params_str = ""
        
        # Create hash of query + params
        key_data = f"{query_str}:{params_str}"
        key_hash = hashlib.sha256(key_data.encode()).hexdigest()[:16]
        
        return f"db:{prefix}:{key_hash}"
    
    def _get_ttl_for_query(self, query: Union[str, Select]) -> int:
        """Determine TTL based on query type."""
        query_str = str(query).lower()
        
        # Check for table names in query
        for table, ttl in self._ttl_config.items():
            if table in query_str:
                return ttl
        
        return self._ttl_config["default"]
    
    async def get_or_fetch(
        self,
        query: Union[str, Select],
        fetch_func: Callable,
        params: Optional[Dict[str, Any]] = None,
        ttl: Optional[int] = None,
        prefix: str = "query"
    ) -> Any:
        """
        Get query result from cache or fetch from database.
        
        Args:
            query: SQL query
            fetch_func: Async function to fetch data if not cached
            params: Query parameters
            ttl: Cache TTL (auto-determined if not provided)
            prefix: Cache key prefix
            
        Returns:
            Query result
        """
        # Generate cache key
        cache_key = self._generate_cache_key(query, params, prefix)
        
        # Try to get from cache
        cached_result = await self._cache.get(cache_key)
        if cached_result is not None:
            self._stats["hits"] += 1
            logger.debug(f"Cache hit for query: {cache_key}")
            return cached_result
        
        # Cache miss - fetch from database
        self._stats["misses"] += 1
        logger.debug(f"Cache miss for query: {cache_key}")
        
        try:
            # Fetch data
            result = await fetch_func()
            
            # Determine TTL
            if ttl is None:
                ttl = self._get_ttl_for_query(query)
            
            # Cache the result
            await self._cache.set(
                cache_key,
                result,
                ttl=ttl,
                compress=len(dumps(result)) > 1024  # Compress if > 1KB
            )
            
            return result
            
        except Exception as e:
            self._stats["errors"] += 1
            logger.error(f"Error in cache fetch: {e}")
            raise
    
    async def invalidate(
        self,
        pattern: Optional[str] = None,
        table: Optional[str] = None,
        prefix: str = "query"
    ):
        """
        Invalidate cached queries.
        
        Args:
            pattern: Pattern to match cache keys
            table: Table name to invalidate
            prefix: Cache key prefix
        """
        self._stats["invalidations"] += 1
        
        if pattern:
            # Invalidate by pattern
            invalidated = await self._invalidate_by_pattern(f"db:{prefix}:{pattern}*")
            logger.info(f"Invalidated {invalidated} cache entries matching pattern: {pattern}")
            
        elif table:
            # Invalidate all queries for a table
            invalidated = await self._invalidate_by_pattern(f"db:*{table}*")
            logger.info(f"Invalidated {invalidated} cache entries for table: {table}")
            
        else:
            # Invalidate all query cache
            invalidated = await self._invalidate_by_pattern(f"db:{prefix}:*")
            logger.info(f"Invalidated {invalidated} cache entries with prefix: {prefix}")
    
    async def _invalidate_by_pattern(self, pattern: str) -> int:
        """Invalidate cache entries matching a pattern."""
        # Note: This is a simplified implementation
        # In production, use Redis SCAN to find matching keys
        count = 0
        
        try:
            # For now, we'll track invalidations
            logger.debug(f"Invalidating cache pattern: {pattern}")
            count = 1  # Placeholder
        except Exception as e:
            logger.error(f"Error invalidating cache: {e}")
        
        return count
    
    def get_stats(self) -> Dict[str, Any]:
        """Get cache statistics."""
        total_requests = self._stats["hits"] + self._stats["misses"]
        hit_rate = (
            self._stats["hits"] / total_requests
            if total_requests > 0 else 0
        )
        
        return {
            **self._stats,
            "total_requests": total_requests,
            "hit_rate": hit_rate
        }


# Global query cache instance
query_cache = QueryCache()


def cached_query(
    ttl: Optional[int] = None,
    key_prefix: str = "query",
    invalidate_on: Optional[List[str]] = None
):
    """
    Decorator for caching database queries.
    
    Args:
        ttl: Cache TTL in seconds
        key_prefix: Prefix for cache key
        invalidate_on: List of table names that invalidate this cache
        
    Example:
        @cached_query(ttl=300, invalidate_on=["users"])
        async def get_user_by_id(session: AsyncSession, user_id: int):
            result = await session.execute(
                select(User).where(User.id == user_id)
            )
            return result.scalar_one_or_none()
    """
    def decorator(func: Callable) -> Callable:
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Extract session and create a cache key from function name and args
            session = None
            for arg in args:
                if isinstance(arg, AsyncSession):
                    session = arg
                    break
            
            # Generate cache key from function and arguments
            cache_key_parts = [
                func.__name__,
                *[str(arg) for arg in args if not isinstance(arg, AsyncSession)],
                *[f"{k}={v}" for k, v in sorted(kwargs.items())]
            ]
            cache_key = ":".join(cache_key_parts)
            
            # Use query cache
            async def fetch_func():
                return await func(*args, **kwargs)
            
            return await query_cache.get_or_fetch(
                query=cache_key,  # Use function signature as "query"
                fetch_func=fetch_func,
                ttl=ttl,
                prefix=key_prefix
            )
        
        # Store invalidation configuration
        if invalidate_on:
            wrapper._invalidate_on = invalidate_on
        
        return wrapper
    return decorator


class CachedRepository:
    """
    Base repository class with built-in caching support.
    
    Example:
        class UserRepository(CachedRepository):
            def __init__(self, session: AsyncSession):
                super().__init__(session, "users")
            
            @cached_query(ttl=1800)
            async def get_by_id(self, user_id: int):
                # Implementation
    """
    
    def __init__(self, session: AsyncSession, table_name: str):
        """
        Initialize cached repository.
        
        Args:
            session: Database session
            table_name: Name of the table for cache invalidation
        """
        self.session = session
        self.table_name = table_name
        self._cache = query_cache
    
    async def invalidate_cache(self, pattern: Optional[str] = None):
        """Invalidate cache for this repository."""
        await self._cache.invalidate(
            table=self.table_name if not pattern else None,
            pattern=pattern
        )
    
    async def after_insert(self, entity: Any):
        """Hook called after insert - invalidates relevant cache."""
        await self.invalidate_cache()
    
    async def after_update(self, entity: Any):
        """Hook called after update - invalidates relevant cache."""
        await self.invalidate_cache()
    
    async def after_delete(self, entity: Any):
        """Hook called after delete - invalidates relevant cache."""
        await self.invalidate_cache()