"""
Performance optimizer for M3 MacBook optimization.

Provides M3-specific optimizations for transcription and batch processing,
including dynamic worker scaling and memory management.
"""

import asyncio
import logging
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Protocol
from functools import lru_cache

from .performance import (
    ResourceMonitor,
    SystemResources,
    M3OptimizationConfig,
    PerformanceMetrics
)

logger = logging.getLogger(__name__)


class PerformanceOptimizerProtocol(Protocol):
    """Protocol for performance optimization services."""
    def get_m3_optimization_config(self) -> M3OptimizationConfig: ...
    def optimize_worker_count(self) -> int: ...
    def optimize_batch_size(self, target_size: int) -> int: ...
    async def optimize_processing_pipeline(
        self, 
        items: List[Any], 
        processor_func: Callable[[Any], Any],
        batch_size: Optional[int] = None
    ) -> List[Any]: ...


class PerformanceOptimizer:
    """M3-optimized performance optimizer for batch processing."""
    
    def __init__(
        self,
        max_workers: int = 8,
        memory_limit_mb: float = 2048.0,
        cpu_limit_percent: float = 90.0,
        resource_monitor: Optional[ResourceMonitor] = None
    ):
        self.max_workers = max_workers
        self.memory_limit_mb = memory_limit_mb
        self.cpu_limit_percent = cpu_limit_percent
        
        # Resource monitoring
        self.resource_monitor = resource_monitor or ResourceMonitor()
        
        # M3 optimization settings
        self.m3_config = self._get_default_m3_config()
        
        # Performance tracking
        self.performance_history: List[PerformanceMetrics] = []
        
        logger.info(f"PerformanceOptimizer initialized with {max_workers} max workers")
    
    def _get_default_m3_config(self) -> M3OptimizationConfig:
        """Get default M3 optimization configuration."""
        return M3OptimizationConfig(
            use_metal_acceleration=True,
            optimize_memory_layout=True,
            parallel_processing_enabled=True,
            chunk_size_mb=50,
            max_concurrent_operations=8,
            memory_alignment_bytes=64
        )
    
    def get_m3_optimization_config(self) -> M3OptimizationConfig:
        """Get current M3 optimization configuration."""
        return self.m3_config
    
    def update_m3_config(self, config: M3OptimizationConfig) -> None:
        """Update M3 optimization configuration."""
        self.m3_config = config
        logger.info("M3 optimization configuration updated")
    
    def optimize_worker_count(self) -> int:
        """Calculate optimal worker count based on system resources."""
        return self.resource_monitor.get_optimal_worker_count(self.max_workers)
    
    def optimize_batch_size(self, target_size: int) -> int:
        """Optimize batch size based on available memory and M3 config."""
        resources = self.resource_monitor.get_system_resources()
        
        # Calculate available memory for processing
        available_memory_mb = resources.memory_available_mb * 0.8  # Use 80% of available
        
        # Estimate memory per item (conservative estimate)
        estimated_memory_per_item_mb = 50.0  # 50MB per item estimate
        
        # Calculate optimal batch size
        optimal_batch_size = int(available_memory_mb / estimated_memory_per_item_mb)
        
        # Apply M3-specific optimizations
        if self.m3_config.optimize_memory_layout:
            # M3 has better memory bandwidth, can handle larger batches
            optimal_batch_size = int(optimal_batch_size * 1.2)
        
        # Ensure batch size is within reasonable bounds
        optimal_batch_size = max(1, min(optimal_batch_size, target_size))
        
        logger.debug(f"Optimized batch size: {optimal_batch_size} "
                    f"(target: {target_size}, available memory: {available_memory_mb:.1f}MB)")
        
        return optimal_batch_size
    
    async def optimize_processing_pipeline(
        self,
        items: List[Any],
        processor_func: Callable[[Any], Any],
        batch_size: Optional[int] = None
    ) -> List[Any]:
        """Optimize processing pipeline with M3-specific enhancements."""
        if not items:
            return []
        
        # Start resource monitoring
        await self.resource_monitor.start_monitoring()
        
        try:
            # Optimize parameters
            optimal_workers = self.optimize_worker_count()
            optimal_batch_size = batch_size or self.optimize_batch_size(len(items))
            
            logger.info(f"Starting optimized processing with {optimal_workers} workers, "
                       f"batch size: {optimal_batch_size}")
            
            # Record start metrics
            start_time = time.time()
            start_resources = self.resource_monitor.get_system_resources()
            
            # Process in batches with optimized workers
            results = await self._process_with_optimization(
                items, processor_func, optimal_workers, optimal_batch_size
            )
            
            # Record end metrics
            end_time = time.time()
            end_resources = self.resource_monitor.get_system_resources()
            
            # Create performance metrics
            metrics = PerformanceMetrics(
                operation="optimized_processing",
                duration_seconds=end_time - start_time,
                memory_peak_mb=end_resources.memory_percent,
                cpu_peak_percent=end_resources.cpu_percent,
                throughput_items_per_second=len(items) / (end_time - start_time),
                error_count=0,  # TODO: Track actual errors
                success_count=len(results),
                total_count=len(items)
            )
            
            self.performance_history.append(metrics)
            
            logger.info(f"Processing completed: {len(results)} items in "
                       f"{metrics.duration_seconds:.2f}s "
                       f"({metrics.throughput_items_per_second:.2f} items/s)")
            
            return results
            
        finally:
            await self.resource_monitor.stop_monitoring()
    
    async def _process_with_optimization(
        self,
        items: List[Any],
        processor_func: Callable[[Any], Any],
        worker_count: int,
        batch_size: int
    ) -> List[Any]:
        """Process items with optimized worker pool and batching."""
        # Create batches
        batches = [items[i:i + batch_size] for i in range(0, len(items), batch_size)]
        
        # Create semaphore to limit concurrent workers
        semaphore = asyncio.Semaphore(worker_count)
        
        async def process_batch(batch: List[Any]) -> List[Any]:
            async with semaphore:
                batch_results = []
                for item in batch:
                    try:
                        if asyncio.iscoroutinefunction(processor_func):
                            result = await processor_func(item)
                        else:
                            result = processor_func(item)
                        batch_results.append(result)
                    except Exception as e:
                        logger.error(f"Error processing item: {e}")
                        # Continue with other items
                        batch_results.append(None)
                return batch_results
        
        # Process all batches concurrently
        batch_tasks = [process_batch(batch) for batch in batches]
        batch_results = await asyncio.gather(*batch_tasks, return_exceptions=True)
        
        # Flatten results
        results = []
        for batch_result in batch_results:
            if isinstance(batch_result, Exception):
                logger.error(f"Batch processing error: {batch_result}")
                # Add None for failed batches
                results.extend([None] * batch_size)
            else:
                results.extend(batch_result)
        
        return results[:len(items)]  # Ensure we return exactly the right number
    
    def get_performance_history(self) -> List[PerformanceMetrics]:
        """Get performance history."""
        return self.performance_history.copy()
    
    def get_average_performance(self) -> Dict[str, float]:
        """Calculate average performance metrics."""
        if not self.performance_history:
            return {}
        
        total_operations = len(self.performance_history)
        avg_duration = sum(m.duration_seconds for m in self.performance_history) / total_operations
        avg_memory = sum(m.memory_peak_mb for m in self.performance_history) / total_operations
        avg_cpu = sum(m.cpu_peak_percent for m in self.performance_history) / total_operations
        avg_throughput = sum(m.throughput_items_per_second for m in self.performance_history) / total_operations
        
        return {
            "avg_duration_seconds": avg_duration,
            "avg_memory_peak_mb": avg_memory,
            "avg_cpu_peak_percent": avg_cpu,
            "avg_throughput_items_per_second": avg_throughput,
            "total_operations": total_operations
        }
    
    @lru_cache(maxsize=32)
    def get_optimal_chunk_size(self, file_size_mb: float) -> int:
        """Get optimal chunk size for file processing based on M3 config."""
        # M3 has better memory bandwidth, can handle larger chunks
        base_chunk_size = self.m3_config.chunk_size_mb
        
        # Adjust based on file size
        if file_size_mb < 100:
            return base_chunk_size
        elif file_size_mb < 500:
            return base_chunk_size * 2
        else:
            return base_chunk_size * 3
    
    def get_memory_optimization_tips(self) -> List[str]:
        """Get memory optimization tips for M3."""
        return [
            "Use Metal acceleration for video processing",
            "Enable memory layout optimization",
            "Use aligned memory allocations (64-byte alignment)",
            "Process large files in chunks",
            "Monitor memory usage and adjust batch sizes dynamically"
        ]
    
    async def cleanup(self) -> None:
        """Cleanup resources."""
        await self.resource_monitor.cleanup()
        logger.info("PerformanceOptimizer cleanup completed")


def create_performance_optimizer(
    max_workers: int = 8,
    memory_limit_mb: float = 2048.0,
    cpu_limit_percent: float = 90.0
) -> PerformanceOptimizer:
    """Create a performance optimizer instance."""
    return PerformanceOptimizer(
        max_workers=max_workers,
        memory_limit_mb=memory_limit_mb,
        cpu_limit_percent=cpu_limit_percent
    )
