"""
Memory-efficient batch loader for image processing.
Reduces memory usage by processing images in configurable batches with optional preprocessing.
"""

import gc
import torch
from PIL import Image
from pathlib import Path
from typing import List, Dict, Any, Iterator, Tuple, Optional
import psutil
import os

from .logger_config import get_logger

logger = get_logger(__name__)


class MemoryEfficientBatchLoader:
    """Memory-efficient image batch loader with preprocessing options."""
    
    def __init__(self, 
                 batch_size: int = 8,
                 max_image_size: Optional[Tuple[int, int]] = (1024, 1024),
                 enable_memory_monitoring: bool = True,
                 memory_threshold_percent: float = 85.0):
        """
        Initialize the batch loader.
        
        Args:
            batch_size: Number of images to process per batch
            max_image_size: Maximum image dimensions (width, height). If None, no resizing
            enable_memory_monitoring: Whether to monitor memory usage
            memory_threshold_percent: Memory usage threshold to trigger cleanup
        """
        self.batch_size = batch_size
        self.max_image_size = max_image_size
        self.enable_memory_monitoring = enable_memory_monitoring
        self.memory_threshold_percent = memory_threshold_percent
        
        logger.info(f"BatchLoader initialized: batch_size={batch_size}, "
                   f"max_size={max_image_size}, memory_monitoring={enable_memory_monitoring}")
    
    def _get_memory_usage(self) -> float:
        """Get current memory usage percentage."""
        return psutil.virtual_memory().percent
    
    def _cleanup_memory(self):
        """Force garbage collection and clear GPU cache if available."""
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        if self.enable_memory_monitoring:
            memory_usage = self._get_memory_usage()
            logger.debug(f"Memory cleanup performed. Current usage: {memory_usage:.1f}%")
    
    def _check_memory_threshold(self):
        """Check if memory usage exceeds threshold and cleanup if needed."""
        if not self.enable_memory_monitoring:
            return
            
        memory_usage = self._get_memory_usage()
        if memory_usage > self.memory_threshold_percent:
            logger.warning(f"Memory usage ({memory_usage:.1f}%) exceeds threshold "
                          f"({self.memory_threshold_percent}%). Performing cleanup...")
            self._cleanup_memory()
    
    def _preprocess_image(self, image_path: str) -> Optional[Image.Image]:
        """
        Load and preprocess image to reduce memory usage.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            PIL Image object or None if failed
        """
        try:
            image = Image.open(image_path)
            
            # Convert to RGB if needed (removes alpha channel)
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Resize if max_image_size is specified
            if self.max_image_size:
                # Calculate resize ratio to maintain aspect ratio
                width, height = image.size
                max_w, max_h = self.max_image_size
                
                if width > max_w or height > max_h:
                    ratio = min(max_w / width, max_h / height)
                    new_size = (int(width * ratio), int(height * ratio))
                    image = image.resize(new_size, Image.Resampling.LANCZOS)
                    logger.debug(f"Resized {Path(image_path).name} from {width}x{height} to {new_size[0]}x{new_size[1]}")
            
            return image
            
        except Exception as e:
            logger.error(f"Failed to preprocess image {image_path}: {e}")
            return None
    
    def _get_file_size_mb(self, file_path: str) -> float:
        """Get file size in MB."""
        try:
            return os.path.getsize(file_path) / (1024 * 1024)
        except:
            return 0.0
    
    def create_batches(self, image_paths: List[str]) -> Iterator[List[str]]:
        """
        Create batches of image paths.
        
        Args:
            image_paths: List of image file paths
            
        Yields:
            Batches of image paths
        """
        for i in range(0, len(image_paths), self.batch_size):
            batch = image_paths[i:i + self.batch_size]
            yield batch
    
    def process_batch(self, image_paths: List[str], 
                     scoring_function, 
                     include_file_info: bool = True) -> Dict[str, Dict[str, Any]]:
        """
        Process a batch of images with memory management.
        
        Args:
            image_paths: List of image paths in the batch
            scoring_function: Function to score images (should accept list of PIL Images)
            include_file_info: Whether to include file size and path info
            
        Returns:
            Dictionary with results for each image in the batch
        """
        results = {}
        
        # Check memory before processing batch
        self._check_memory_threshold()
        
        # Preprocess all images in batch
        processed_images = []
        valid_paths = []
        
        for path in image_paths:
            image = self._preprocess_image(path)
            if image is not None:
                processed_images.append(image)
                valid_paths.append(path)
            else:
                # Handle failed image
                image_name = Path(path).name
                results[image_name] = {
                    'error': f'Failed to load/preprocess image: {path}',
                    'file_path': path if include_file_info else None,
                    'file_size_mb': self._get_file_size_mb(path) if include_file_info else None
                }
        
        if processed_images:
            try:
                # Score the batch
                scores = scoring_function(processed_images)
                
                # Combine results
                for i, (path, score) in enumerate(zip(valid_paths, scores)):
                    image_name = Path(path).name
                    result = {
                        'deqa_score': score,
                    }
                    
                    if include_file_info:
                        result.update({
                            'file_path': path,
                            'file_size_mb': round(self._get_file_size_mb(path), 2)
                        })
                    
                    results[image_name] = result
                
                logger.debug(f"Processed batch of {len(processed_images)} images")
                
            except Exception as e:
                logger.error(f"Error scoring batch: {e}")
                # Mark all as failed
                for path in valid_paths:
                    image_name = Path(path).name
                    results[image_name] = {
                        'error': f'Scoring failed: {str(e)}',
                        'file_path': path if include_file_info else None,
                        'file_size_mb': self._get_file_size_mb(path) if include_file_info else None
                    }
            
            finally:
                # Clear processed images from memory
                del processed_images
                self._cleanup_memory()
        
        return results
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """Get current memory statistics."""
        memory = psutil.virtual_memory()
        stats = {
            'memory_percent': memory.percent,
            'memory_available_gb': memory.available / (1024**3),
            'memory_used_gb': memory.used / (1024**3),
            'memory_total_gb': memory.total / (1024**3)
        }
        
        if torch.cuda.is_available():
            gpu_memory = torch.cuda.memory_allocated() / (1024**3)
            gpu_cached = torch.cuda.memory_reserved() / (1024**3)
            stats.update({
                'gpu_memory_allocated_gb': gpu_memory,
                'gpu_memory_cached_gb': gpu_cached
            })
        
        return stats


class BatchProcessor:
    """High-level processor that uses MemoryEfficientBatchLoader."""
    
    def __init__(self, 
                 batch_size: int = 8,
                 max_image_size: Optional[Tuple[int, int]] = (1024, 1024),
                 enable_memory_monitoring: bool = True):
        """
        Initialize the batch processor.
        
        Args:
            batch_size: Number of images to process per batch
            max_image_size: Maximum image dimensions for preprocessing
            enable_memory_monitoring: Whether to monitor memory usage
        """
        self.loader = MemoryEfficientBatchLoader(
            batch_size=batch_size,
            max_image_size=max_image_size,
            enable_memory_monitoring=enable_memory_monitoring
        )
    
    def process_images_in_batches(self, 
                                 image_paths: List[str], 
                                 scorer,
                                 progress_callback=None) -> Dict[str, Dict[str, Any]]:
        """
        Process all images in batches.
        
        Args:
            image_paths: List of all image paths to process
            scorer: Object with score_multiple_images method
            progress_callback: Optional function to call with progress updates
            
        Returns:
            Combined results from all batches
        """
        all_results = {}
        total_images = len(image_paths)
        processed_images = 0
        
        logger.info(f"Processing {total_images} images in batches of {self.loader.batch_size}")
        
        # Show initial memory stats
        if self.loader.enable_memory_monitoring:
            initial_stats = self.loader.get_memory_stats()
            logger.info(f"Initial memory usage: {initial_stats['memory_percent']:.1f}% "
                       f"({initial_stats['memory_used_gb']:.1f}GB/{initial_stats['memory_total_gb']:.1f}GB)")
        
        for batch_idx, batch_paths in enumerate(self.loader.create_batches(image_paths)):
            logger.info(f"Processing batch {batch_idx + 1} ({len(batch_paths)} images)...")
            
            # Create scoring function that works with the batch loader
            def scoring_function(images):
                return scorer.score_multiple_images_from_pil(images)
            
            # Process batch
            batch_results = self.loader.process_batch(batch_paths, scoring_function)
            all_results.update(batch_results)
            
            processed_images += len(batch_paths)
            
            # Progress callback
            if progress_callback:
                progress_callback(processed_images, total_images)
            
            # Log progress and memory stats
            progress_percent = (processed_images / total_images) * 100
            logger.info(f"Progress: {processed_images}/{total_images} ({progress_percent:.1f}%)")
            
            if self.loader.enable_memory_monitoring:
                current_stats = self.loader.get_memory_stats()
                logger.debug(f"Current memory usage: {current_stats['memory_percent']:.1f}%")
        
        logger.info(f"Completed processing {total_images} images")
        
        # Final memory cleanup
        self.loader._cleanup_memory()
        
        return all_results
