"""Resource cleanup manager for diarization services."""

import logging
import gc
import psutil
import torch
import tempfile
import shutil
from pathlib import Path
from typing import Dict, List, Optional, Set, Any
from dataclasses import dataclass
from datetime import datetime, timezone, timedelta
import threading
import time

logger = logging.getLogger(__name__)


@dataclass
class CleanupStats:
    """Statistics about cleanup operations."""
    
    memory_freed_mb: float
    gpu_memory_freed_mb: float
    temp_files_removed: int
    cache_entries_cleared: int
    cleanup_duration_seconds: float
    timestamp: datetime


class ResourceCleanupManager:
    """Manages automatic cleanup of resources to prevent memory leaks and optimize performance."""
    
    def __init__(self, cleanup_interval_seconds: int = 60, max_temp_files: int = 100):
        """Initialize the cleanup manager.
        
        Args:
            cleanup_interval_seconds: Interval between automatic cleanups
            max_temp_files: Maximum number of temporary files to keep
        """
        self.cleanup_interval_seconds = cleanup_interval_seconds
        self.max_temp_files = max_temp_files
        self.temp_directories: Set[Path] = set()
        self.cache_directories: Set[Path] = set()
        self.model_references: Set[Any] = set()
        self.last_cleanup = datetime.now(timezone.utc)
        self.cleanup_stats: List[CleanupStats] = []
        
        # Start background cleanup thread
        self._stop_event = threading.Event()
        self._cleanup_thread = threading.Thread(target=self._background_cleanup, daemon=True)
        self._cleanup_thread.start()
        
        logger.info(f"Resource cleanup manager initialized with {cleanup_interval_seconds}s interval")
    
    def register_temp_directory(self, directory: Path):
        """Register a temporary directory for cleanup.
        
        Args:
            directory: Path to temporary directory
        """
        self.temp_directories.add(directory)
        logger.debug(f"Registered temp directory: {directory}")
    
    def register_cache_directory(self, directory: Path):
        """Register a cache directory for cleanup.
        
        Args:
            directory: Path to cache directory
        """
        self.cache_directories.add(directory)
        logger.debug(f"Registered cache directory: {directory}")
    
    def register_model_reference(self, model: Any):
        """Register a model reference for cleanup.
        
        Args:
            model: Model object to track
        """
        self.model_references.add(model)
        logger.debug(f"Registered model reference: {type(model).__name__}")
    
    def cleanup_memory(self) -> Dict[str, float]:
        """Clean up memory resources.
        
        Returns:
            Dictionary with memory cleanup statistics
        """
        start_time = time.time()
        
        # Get initial memory usage
        initial_memory = psutil.virtual_memory()
        initial_memory_mb = initial_memory.used / (1024 * 1024)
        
        # Get initial GPU memory if available
        initial_gpu_memory_mb = 0
        if torch.cuda.is_available():
            try:
                initial_gpu_memory_mb = torch.cuda.memory_allocated() / (1024 * 1024)
            except (AttributeError, RuntimeError):
                pass
        
        # Clear model references
        self.model_references.clear()
        
        # Force garbage collection
        collected = gc.collect()
        
        # Clear PyTorch cache
        if torch.cuda.is_available():
            try:
                torch.cuda.empty_cache()
            except (AttributeError, RuntimeError):
                pass
        
        # Get final memory usage
        final_memory = psutil.virtual_memory()
        final_memory_mb = final_memory.used / (1024 * 1024)
        
        # Calculate memory freed
        memory_freed_mb = initial_memory_mb - final_memory_mb
        
        # Calculate GPU memory freed
        gpu_memory_freed_mb = 0
        if torch.cuda.is_available():
            try:
                final_gpu_memory_mb = torch.cuda.memory_allocated() / (1024 * 1024)
                gpu_memory_freed_mb = initial_gpu_memory_mb - final_gpu_memory_mb
            except (AttributeError, RuntimeError):
                pass
        
        cleanup_duration = time.time() - start_time
        
        logger.info(f"Memory cleanup completed: {memory_freed_mb:.1f}MB RAM, "
                   f"{gpu_memory_freed_mb:.1f}MB GPU, {collected} objects collected")
        
        return {
            "memory_freed_mb": memory_freed_mb,
            "gpu_memory_freed_mb": gpu_memory_freed_mb,
            "objects_collected": collected,
            "cleanup_duration_seconds": cleanup_duration
        }
    
    def cleanup_temp_files(self) -> Dict[str, int]:
        """Clean up temporary files.
        
        Returns:
            Dictionary with file cleanup statistics
        """
        start_time = time.time()
        files_removed = 0
        directories_removed = 0
        
        # Clean up registered temp directories
        for temp_dir in list(self.temp_directories):
            try:
                if temp_dir.exists():
                    # Remove files older than 1 hour
                    cutoff_time = datetime.now(timezone.utc) - timedelta(hours=1)
                    
                    for file_path in temp_dir.rglob("*"):
                        if file_path.is_file():
                            try:
                                file_mtime = datetime.fromtimestamp(file_path.stat().st_mtime, tz=timezone.utc)
                                if file_mtime < cutoff_time:
                                    file_path.unlink()
                                    files_removed += 1
                            except (OSError, PermissionError):
                                pass  # Skip files that can't be removed
                    
                    # Remove empty directories
                    for dir_path in reversed(list(temp_dir.rglob("*"))):
                        if dir_path.is_dir() and not any(dir_path.iterdir()):
                            try:
                                dir_path.rmdir()
                                directories_removed += 1
                            except (OSError, PermissionError):
                                pass
                
            except Exception as e:
                logger.warning(f"Failed to cleanup temp directory {temp_dir}: {e}")
        
        cleanup_duration = time.time() - start_time
        
        logger.info(f"Temp file cleanup completed: {files_removed} files, "
                   f"{directories_removed} directories removed")
        
        return {
            "files_removed": files_removed,
            "directories_removed": directories_removed,
            "cleanup_duration_seconds": cleanup_duration
        }
    
    def cleanup_cache(self) -> Dict[str, int]:
        """Clean up cache directories.
        
        Returns:
            Dictionary with cache cleanup statistics
        """
        start_time = time.time()
        cache_entries_cleared = 0
        
        # Clean up registered cache directories
        for cache_dir in list(self.cache_directories):
            try:
                if cache_dir.exists():
                    # Remove cache entries older than 24 hours
                    cutoff_time = datetime.now(timezone.utc) - timedelta(hours=24)
                    
                    for file_path in cache_dir.rglob("*"):
                        if file_path.is_file():
                            try:
                                file_mtime = datetime.fromtimestamp(file_path.stat().st_mtime, tz=timezone.utc)
                                if file_mtime < cutoff_time:
                                    file_path.unlink()
                                    cache_entries_cleared += 1
                            except (OSError, PermissionError):
                                pass
                
            except Exception as e:
                logger.warning(f"Failed to cleanup cache directory {cache_dir}: {e}")
        
        cleanup_duration = time.time() - start_time
        
        logger.info(f"Cache cleanup completed: {cache_entries_cleared} entries cleared")
        
        return {
            "cache_entries_cleared": cache_entries_cleared,
            "cleanup_duration_seconds": cleanup_duration
        }
    
    def perform_full_cleanup(self) -> CleanupStats:
        """Perform a full cleanup of all resources.
        
        Returns:
            Cleanup statistics
        """
        start_time = time.time()
        
        # Clean up memory
        memory_stats = self.cleanup_memory()
        
        # Clean up temp files
        temp_stats = self.cleanup_temp_files()
        
        # Clean up cache
        cache_stats = self.cleanup_cache()
        
        cleanup_duration = time.time() - start_time
        
        # Create cleanup stats
        stats = CleanupStats(
            memory_freed_mb=memory_stats["memory_freed_mb"],
            gpu_memory_freed_mb=memory_stats["gpu_memory_freed_mb"],
            temp_files_removed=temp_stats["files_removed"],
            cache_entries_cleared=cache_stats["cache_entries_cleared"],
            cleanup_duration_seconds=cleanup_duration,
            timestamp=datetime.now(timezone.utc)
        )
        
        self.cleanup_stats.append(stats)
        
        # Keep only last 100 stats
        if len(self.cleanup_stats) > 100:
            self.cleanup_stats = self.cleanup_stats[-100:]
        
        self.last_cleanup = stats.timestamp
        
        logger.info(f"Full cleanup completed in {cleanup_duration:.2f}s")
        
        return stats
    
    def _background_cleanup(self):
        """Background thread for periodic cleanup."""
        while not self._stop_event.is_set():
            try:
                # Wait for cleanup interval
                self._stop_event.wait(self.cleanup_interval_seconds)
                
                if not self._stop_event.is_set():
                    # Perform cleanup
                    self.perform_full_cleanup()
                    
            except Exception as e:
                logger.error(f"Background cleanup failed: {e}")
    
    def get_cleanup_stats(self, limit: int = 10) -> List[CleanupStats]:
        """Get recent cleanup statistics.
        
        Args:
            limit: Maximum number of stats to return
            
        Returns:
            List of recent cleanup statistics
        """
        return self.cleanup_stats[-limit:] if self.cleanup_stats else []
    
    def get_memory_usage(self) -> Dict[str, float]:
        """Get current memory usage.
        
        Returns:
            Dictionary with memory usage information
        """
        memory = psutil.virtual_memory()
        
        result = {
            "total_memory_gb": memory.total / (1024**3),
            "available_memory_gb": memory.available / (1024**3),
            "used_memory_gb": memory.used / (1024**3),
            "memory_percent": memory.percent
        }
        
        # Add GPU memory if available
        if torch.cuda.is_available():
            try:
                gpu_memory_allocated = torch.cuda.memory_allocated() / (1024**3)
                gpu_memory_reserved = torch.cuda.memory_reserved() / (1024**3)
                gpu_memory_total = torch.cuda.get_device_properties(0).total_memory / (1024**3)
                
                result.update({
                    "gpu_memory_allocated_gb": gpu_memory_allocated,
                    "gpu_memory_reserved_gb": gpu_memory_reserved,
                    "gpu_memory_total_gb": gpu_memory_total,
                    "gpu_memory_percent": (gpu_memory_allocated / gpu_memory_total) * 100
                })
            except (AttributeError, RuntimeError):
                pass
        
        return result
    
    def should_perform_cleanup(self) -> bool:
        """Check if cleanup should be performed based on memory usage.
        
        Returns:
            True if cleanup should be performed
        """
        memory_usage = self.get_memory_usage()
        
        # Cleanup if memory usage is high
        if memory_usage["memory_percent"] > 80:
            return True
        
        # Cleanup if GPU memory usage is high
        if "gpu_memory_percent" in memory_usage and memory_usage["gpu_memory_percent"] > 80:
            return True
        
        # Cleanup if it's been a while since last cleanup
        time_since_last = datetime.now(timezone.utc) - self.last_cleanup
        if time_since_last.total_seconds() > self.cleanup_interval_seconds * 2:
            return True
        
        return False
    
    def shutdown(self):
        """Shutdown the cleanup manager."""
        logger.info("Shutting down resource cleanup manager")
        
        # Stop background thread
        self._stop_event.set()
        if self._cleanup_thread.is_alive():
            self._cleanup_thread.join(timeout=5)
        
        # Perform final cleanup
        self.perform_full_cleanup()
        
        logger.info("Resource cleanup manager shutdown complete")
