"""
Bounded LRU cache for reference data with memory management.

Provides memory-aware caching to prevent unbounded memory growth during
large dataset generation while maintaining performance benefits.
"""

from collections import OrderedDict
from typing import Optional, Dict
import polars as pl
import psutil
import os


class BoundedReferenceCache:
    """
    LRU cache with memory bounds for reference data DataFrames.
    
    Features:
    - LRU eviction when memory limits are exceeded
    - Memory usage estimation for DataFrames
    - Automatic cleanup when memory pressure is detected
    - Thread-safe operations for concurrent access
    """
    
    def __init__(self, max_memory_bytes: int = 1_000_000_000):  # 1GB default
        self.max_memory_bytes = max_memory_bytes
        self.cache: OrderedDict[str, pl.DataFrame] = OrderedDict()
        self.memory_usage: Dict[str, int] = {}
        self.current_memory = 0
    
    def get(self, key: str) -> Optional[pl.DataFrame]:
        """
        Get cached DataFrame, moving to end if found (LRU).
        
        Args:
            key: Cache key for the reference data
            
        Returns:
            Cached DataFrame or None if not found
        """
        if key in self.cache:
            # Move to end (most recently used)
            self.cache.move_to_end(key)
            return self.cache[key]
        return None
    
    def put(self, key: str, df: pl.DataFrame) -> None:
        """
        Store DataFrame in cache with memory bounds checking.
        
        Args:
            key: Cache key for the reference data
            df: DataFrame to cache
        """
        # Estimate memory usage more accurately than static estimates
        memory_size = self._estimate_dataframe_memory(df)
        
        # Evict oldest entries if needed to make room
        while (self.current_memory + memory_size > self.max_memory_bytes 
               and self.cache):
            oldest_key, _ = self.cache.popitem(last=False)
            self.current_memory -= self.memory_usage.pop(oldest_key, 0)
        
        # Remove existing entry if updating
        if key in self.cache:
            self.current_memory -= self.memory_usage.pop(key, 0)
            del self.cache[key]
        
        # Add new entry
        self.cache[key] = df
        self.memory_usage[key] = memory_size
        self.current_memory += memory_size
    
    def _estimate_dataframe_memory(self, df: pl.DataFrame) -> int:
        """
        Estimate memory usage of a Polars DataFrame.
        
        Args:
            df: DataFrame to estimate
            
        Returns:
            Estimated memory usage in bytes
        """
        if df.is_empty():
            return 1024  # Minimal overhead for empty DataFrames
        
        # Use Polars' built-in memory estimation if available
        if hasattr(df, 'estimated_size'):
            return int(df.estimated_size('mb') * 1_000_000)
        
        # Fallback: estimate based on data types and row count
        bytes_per_row = 0
        for dtype in df.dtypes:
            if dtype == pl.String:
                bytes_per_row += 64  # Average string length
            elif dtype in [pl.Int64, pl.Float64]:
                bytes_per_row += 8
            elif dtype in [pl.Int32, pl.Float32]:
                bytes_per_row += 4
            elif dtype == pl.Boolean:
                bytes_per_row += 1
            else:
                bytes_per_row += 8  # Default
        
        # Add overhead for DataFrame structure
        estimated_size = (df.height * bytes_per_row) + 10240  # 10KB overhead
        return estimated_size
    
    def clear(self) -> None:
        """Clear all cached entries."""
        self.cache.clear()
        self.memory_usage.clear()
        self.current_memory = 0
    
    def evict_if_memory_pressure(self, threshold_percentage: float = 0.8) -> int:
        """
        Proactively evict cache entries if system memory pressure is detected.
        
        Args:
            threshold_percentage: System memory threshold (0.8 = 80%)
            
        Returns:
            Number of entries evicted
        """
        # Check system memory usage
        memory_info = psutil.virtual_memory()
        if memory_info.percent / 100.0 > threshold_percentage:
            # System memory pressure - evict half the cache
            target_size = len(self.cache) // 2
            evicted = 0
            
            while len(self.cache) > target_size and self.cache:
                oldest_key, _ = self.cache.popitem(last=False)
                self.current_memory -= self.memory_usage.pop(oldest_key, 0)
                evicted += 1
            
            return evicted
        
        return 0
    
    def get_stats(self) -> Dict[str, any]:
        """
        Get cache statistics for monitoring.
        
        Returns:
            Dictionary with cache statistics
        """
        return {
            "entries": len(self.cache),
            "memory_usage_bytes": self.current_memory,
            "memory_limit_bytes": self.max_memory_bytes,
            "memory_utilization": self.current_memory / self.max_memory_bytes,
            "keys": list(self.cache.keys())
        }
    
    def __len__(self) -> int:
        """Return number of cached entries."""
        return len(self.cache)
    
    def __contains__(self, key: str) -> bool:
        """Check if key exists in cache."""
        return key in self.cache