"""Memory optimization for domain adaptation system.

This module provides memory optimization features for the domain adaptation system,
including adapter swapping, shared parameters, and memory management.
"""

import logging
import os
import gc
import time
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass
from collections import OrderedDict

import torch
import psutil

logger = logging.getLogger(__name__)


@dataclass
class MemoryStats:
    """Memory usage statistics."""
    rss_mb: float
    vms_mb: float
    percent: float
    gpu_memory_mb: Optional[float] = None
    gpu_memory_percent: Optional[float] = None


class AdapterCache:
    """LRU cache for domain adapters to optimize memory usage."""
    
    def __init__(self, max_size: int = 3, max_memory_mb: int = 2048):
        """Initialize the adapter cache.
        
        Args:
            max_size: Maximum number of adapters to keep in memory
            max_memory_mb: Maximum memory usage in MB
        """
        self.max_size = max_size
        self.max_memory_mb = max_memory_mb
        self.cache: OrderedDict[str, Any] = OrderedDict()
        self.adapter_sizes: Dict[str, int] = {}
        
        logger.info(f"Adapter cache initialized: max_size={max_size}, max_memory={max_memory_mb}MB")
    
    def get(self, domain_name: str) -> Optional[Any]:
        """Get an adapter from cache.
        
        Args:
            domain_name: Name of the domain
            
        Returns:
            Cached adapter or None if not found
        """
        if domain_name in self.cache:
            # Move to end (most recently used)
            self.cache.move_to_end(domain_name)
            logger.debug(f"Cache hit for domain: {domain_name}")
            return self.cache[domain_name]
        
        logger.debug(f"Cache miss for domain: {domain_name}")
        return None
    
    def put(self, domain_name: str, adapter: Any, size_mb: int) -> None:
        """Add an adapter to cache.
        
        Args:
            domain_name: Name of the domain
            adapter: Adapter model to cache
            size_mb: Size of the adapter in MB
        """
        # Remove if already exists
        if domain_name in self.cache:
            self.cache.pop(domain_name)
            self.adapter_sizes.pop(domain_name, 0)
        
        # Check memory constraints
        current_memory = sum(self.adapter_sizes.values())
        if current_memory + size_mb > self.max_memory_mb:
            self._evict_adapters(size_mb)
        
        # Check size constraints
        if len(self.cache) >= self.max_size:
            self._evict_lru()
        
        # Add to cache
        self.cache[domain_name] = adapter
        self.adapter_sizes[domain_name] = size_mb
        
        logger.debug(f"Cached adapter for domain: {domain_name} (size: {size_mb}MB)")
    
    def _evict_lru(self) -> None:
        """Evict least recently used adapter."""
        if self.cache:
            domain_name = next(iter(self.cache))
            self.cache.pop(domain_name)
            self.adapter_sizes.pop(domain_name, 0)
            logger.debug(f"Evicted LRU adapter: {domain_name}")
    
    def _evict_adapters(self, required_mb: int) -> None:
        """Evict adapters to free required memory."""
        current_memory = sum(self.adapter_sizes.values())
        target_memory = current_memory + required_mb - self.max_memory_mb
        
        evicted_memory = 0
        domains_to_evict = []
        
        for domain_name in self.cache:
            if evicted_memory >= target_memory:
                break
            domains_to_evict.append(domain_name)
            evicted_memory += self.adapter_sizes.get(domain_name, 0)
        
        for domain_name in domains_to_evict:
            self.cache.pop(domain_name)
            self.adapter_sizes.pop(domain_name, 0)
            logger.debug(f"Evicted adapter for memory: {domain_name}")
    
    def clear(self) -> None:
        """Clear all cached adapters."""
        self.cache.clear()
        self.adapter_sizes.clear()
        logger.info("Adapter cache cleared")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get cache statistics.
        
        Returns:
            Dictionary with cache statistics
        """
        return {
            "size": len(self.cache),
            "max_size": self.max_size,
            "memory_used_mb": sum(self.adapter_sizes.values()),
            "max_memory_mb": self.max_memory_mb,
            "domains": list(self.cache.keys())
        }


class DomainMemoryOptimizer:
    """Memory optimization manager for domain adaptation system."""
    
    def __init__(self, cache_size: int = 3, max_memory_mb: int = 2048):
        """Initialize the memory optimizer.
        
        Args:
            cache_size: Maximum number of adapters to keep in memory
            max_memory_mb: Maximum memory usage in MB
        """
        self.cache = AdapterCache(cache_size, max_memory_mb)
        self.swap_dir = Path("models/adapters/swap")
        self.swap_dir.mkdir(parents=True, exist_ok=True)
        
        logger.info("Domain memory optimizer initialized")
    
    def get_memory_stats(self) -> MemoryStats:
        """Get current memory usage statistics.
        
        Returns:
            Memory usage statistics
        """
        process = psutil.Process()
        memory_info = process.memory_info()
        
        stats = MemoryStats(
            rss_mb=memory_info.rss / (1024 * 1024),
            vms_mb=memory_info.vms / (1024 * 1024),
            percent=process.memory_percent()
        )
        
        # Add GPU memory if available
        if torch.cuda.is_available():
            gpu_memory = torch.cuda.memory_allocated() / (1024 * 1024)
            gpu_memory_percent = (gpu_memory / torch.cuda.get_device_properties(0).total_memory) * 100
            stats.gpu_memory_mb = gpu_memory
            stats.gpu_memory_percent = gpu_memory_percent
        
        return stats
    
    def estimate_adapter_size(self, adapter: Any) -> int:
        """Estimate the memory size of an adapter in MB.
        
        Args:
            adapter: Adapter model to measure
            
        Returns:
            Estimated size in MB
        """
        try:
            # Count parameters
            total_params = sum(p.numel() for p in adapter.parameters())
            
            # Estimate size (assuming float16 for most parameters)
            size_bytes = total_params * 2  # 2 bytes per parameter for float16
            size_mb = size_bytes / (1024 * 1024)
            
            return int(size_mb)
        except Exception as e:
            logger.warning(f"Failed to estimate adapter size: {e}")
            return 100  # Default estimate
    
    def swap_adapter_to_disk(self, domain_name: str, adapter: Any) -> str:
        """Swap an adapter to disk to free memory.
        
        Args:
            domain_name: Name of the domain
            adapter: Adapter to swap
            
        Returns:
            Path to the swapped adapter file
        """
        swap_path = self.swap_dir / f"{domain_name}_swapped.pt"
        
        try:
            logger.info(f"Swapping adapter to disk: {domain_name}")
            
            # Save adapter state
            torch.save(adapter.state_dict(), swap_path)
            
            # Clear CUDA cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            logger.info(f"Adapter swapped to disk: {swap_path}")
            return str(swap_path)
            
        except Exception as e:
            logger.error(f"Failed to swap adapter to disk: {e}")
            raise
    
    def load_adapter_from_disk(self, domain_name: str, swap_path: str, base_model: Any) -> Any:
        """Load an adapter from disk.
        
        Args:
            domain_name: Name of the domain
            swap_path: Path to the swapped adapter file
            base_model: Base model to attach the adapter to
            
        Returns:
            Loaded adapter
        """
        try:
            logger.info(f"Loading adapter from disk: {domain_name}")
            
            # Load adapter state
            state_dict = torch.load(swap_path, map_location='cpu')
            
            # Create new adapter and load state
            from .domain_adaptation import DomainConfig
            config = DomainConfig(name=domain_name)
            
            from peft import LoraConfig, get_peft_model
            lora_config = LoraConfig(
                r=config.rank,
                lora_alpha=config.alpha,
                target_modules=config.target_modules,
                lora_dropout=config.dropout,
                bias="none",
                task_type="SEQ_2_SEQ_LM"
            )
            
            adapter = get_peft_model(base_model, lora_config)
            adapter.load_state_dict(state_dict)
            
            logger.info(f"Adapter loaded from disk: {domain_name}")
            return adapter
            
        except Exception as e:
            logger.error(f"Failed to load adapter from disk: {e}")
            raise
    
    def optimize_memory_usage(self, current_adapters: Dict[str, Any], base_model: Any) -> Dict[str, Any]:
        """Optimize memory usage by swapping adapters as needed.
        
        Args:
            current_adapters: Currently loaded adapters
            base_model: Base model for creating new adapters
            
        Returns:
            Optimized adapter dictionary
        """
        memory_stats = self.get_memory_stats()
        logger.info(f"Current memory usage: {memory_stats.rss_mb:.1f}MB")
        
        # Check if memory usage is high
        if memory_stats.rss_mb > 4000:  # 4GB threshold
            logger.info("High memory usage detected, optimizing...")
            
            # Sort adapters by last access time (if available)
            adapters_to_swap = list(current_adapters.keys())
            
            # Swap adapters to disk
            for domain_name in adapters_to_swap:
                if domain_name in current_adapters:
                    adapter = current_adapters[domain_name]
                    swap_path = self.swap_adapter_to_disk(domain_name, adapter)
                    
                    # Remove from memory
                    del current_adapters[domain_name]
                    
                    # Store swap path for later retrieval
                    self.cache.put(domain_name, swap_path, 0)  # Size 0 for swap paths
            
            # Force garbage collection
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            logger.info("Memory optimization completed")
        
        return current_adapters
    
    def cleanup_swap_files(self) -> None:
        """Clean up temporary swap files."""
        try:
            for swap_file in self.swap_dir.glob("*_swapped.pt"):
                swap_file.unlink()
                logger.debug(f"Cleaned up swap file: {swap_file}")
            
            logger.info("Swap files cleanup completed")
        except Exception as e:
            logger.warning(f"Failed to cleanup swap files: {e}")
    
    def get_optimization_stats(self) -> Dict[str, Any]:
        """Get memory optimization statistics.
        
        Returns:
            Dictionary with optimization statistics
        """
        memory_stats = self.get_memory_stats()
        cache_stats = self.cache.get_stats()
        
        return {
            "memory_usage": {
                "rss_mb": memory_stats.rss_mb,
                "vms_mb": memory_stats.vms_mb,
                "percent": memory_stats.percent,
                "gpu_memory_mb": memory_stats.gpu_memory_mb,
                "gpu_memory_percent": memory_stats.gpu_memory_percent
            },
            "cache_stats": cache_stats,
            "swap_files": len(list(self.swap_dir.glob("*_swapped.pt")))
        }
