"""
Performance optimization services for Trax platform.

Provides resource monitoring, M3 optimization, and performance benchmarking
for optimal transcription and batch processing performance.
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Callable, Dict, List, Optional, Protocol
from pathlib import Path

import psutil

from ..base.services import BaseService

logger = logging.getLogger(__name__)


@dataclass
class SystemResources:
    """Current system resource usage."""
    memory_percent: float
    cpu_percent: float
    disk_percent: float
    memory_available_mb: float
    cpu_count: int
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class PerformanceMetrics:
    """Performance metrics for operations."""
    operation: str
    duration_seconds: float
    memory_peak_mb: float
    cpu_peak_percent: float
    throughput_items_per_second: float
    error_count: int
    success_count: int
    total_count: int
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


@dataclass
class M3OptimizationConfig:
    """M3-specific optimization configuration."""
    use_metal_acceleration: bool = True
    optimize_memory_layout: bool = True
    parallel_processing_enabled: bool = True
    chunk_size_mb: int = 50
    max_concurrent_operations: int = 8
    memory_alignment_bytes: int = 64


@dataclass
class ResourceThresholds:
    """Resource usage thresholds for optimization."""
    memory_critical: float = 85.0
    cpu_critical: float = 90.0
    disk_critical: float = 90.0
    memory_warning: float = 70.0
    cpu_warning: float = 80.0
    disk_warning: float = 80.0


class ResourceMonitorProtocol(Protocol):
    """Protocol for resource monitoring services."""
    async def start_monitoring(self) -> None: ...
    async def stop_monitoring(self) -> None: ...
    def get_system_resources(self) -> SystemResources: ...
    def get_optimal_worker_count(self, max_workers: int) -> int: ...


class ResourceMonitor(BaseService):
    """Monitor system resources and provide optimization recommendations."""
    
    def __init__(
        self,
        memory_threshold: float = 80.0,
        cpu_threshold: float = 90.0,
        disk_threshold: float = 85.0,
        monitoring_interval: float = 5.0
    ):
        super().__init__("ResourceMonitor")
        self.memory_threshold = memory_threshold
        self.cpu_threshold = cpu_threshold
        self.disk_threshold = disk_threshold
        self.monitoring_interval = monitoring_interval
        
        # Monitoring state
        self.is_monitoring = False
        self.monitor_task: Optional[asyncio.Task] = None
        self.callback: Optional[Callable[[SystemResources], None]] = None
        
        # Resource tracking
        self.memory_peak = 0.0
        self.cpu_peak = 0.0
        self.disk_peak = 0.0
        
        logger.info(f"ResourceMonitor initialized with thresholds: "
                   f"memory={memory_threshold}%, cpu={cpu_threshold}%, disk={disk_threshold}%")
    
    def get_system_resources(self) -> SystemResources:
        """Get current system resource usage."""
        try:
            memory = psutil.virtual_memory()
            cpu_percent = psutil.cpu_percent(interval=0.1)
            
            # Get disk usage for root directory
            disk = psutil.disk_usage('/')
            
            resources = SystemResources(
                memory_percent=memory.percent,
                cpu_percent=cpu_percent,
                disk_percent=disk.percent,
                memory_available_mb=memory.available / (1024 * 1024),
                cpu_count=psutil.cpu_count(logical=True)
            )
            
            # Update peaks
            self.memory_peak = max(self.memory_peak, memory.percent)
            self.cpu_peak = max(self.cpu_peak, cpu_percent)
            self.disk_peak = max(self.disk_peak, disk.percent)
            
            return resources
            
        except Exception as e:
            logger.error(f"Error getting system resources: {e}")
            # Return safe defaults
            return SystemResources(
                memory_percent=0.0,
                cpu_percent=0.0,
                disk_percent=0.0,
                memory_available_mb=1024.0,
                cpu_count=1
            )
    
    def is_memory_critical(self, memory_percent: float) -> bool:
        """Check if memory usage is critical."""
        return memory_percent > self.memory_threshold
    
    def is_cpu_critical(self, cpu_percent: float) -> bool:
        """Check if CPU usage is critical."""
        return cpu_percent > self.cpu_threshold
    
    def is_disk_critical(self, disk_percent: float) -> bool:
        """Check if disk usage is critical."""
        return disk_percent > self.disk_threshold
    
    def get_optimal_worker_count(self, max_workers: int) -> int:
        """Calculate optimal number of workers based on system resources."""
        resources = self.get_system_resources()
        
        # Base calculation on CPU cores
        optimal_workers = min(resources.cpu_count, max_workers)
        
        # Reduce workers if memory is critical
        if self.is_memory_critical(resources.memory_percent):
            optimal_workers = max(1, optimal_workers // 2)
            logger.warning(f"Memory usage critical ({resources.memory_percent}%), "
                          f"reducing workers to {optimal_workers}")
        
        # Reduce workers if CPU is critical
        if self.is_cpu_critical(resources.cpu_percent):
            optimal_workers = max(1, optimal_workers // 2)
            logger.warning(f"CPU usage critical ({resources.cpu_percent}%), "
                          f"reducing workers to {optimal_workers}")
        
        return optimal_workers
    
    def set_callback(self, callback: Callable[[SystemResources], None]) -> None:
        """Set callback function for resource monitoring."""
        self.callback = callback
    
    async def start_monitoring(self) -> None:
        """Start resource monitoring."""
        if self.is_monitoring:
            logger.warning("Resource monitoring already started")
            return
        
        self.is_monitoring = True
        self.monitor_task = asyncio.create_task(self._monitor_loop())
        logger.info("Resource monitoring started")
    
    async def stop_monitoring(self) -> None:
        """Stop resource monitoring."""
        if not self.is_monitoring:
            return
        
        self.is_monitoring = False
        
        if self.monitor_task:
            self.monitor_task.cancel()
            try:
                await self.monitor_task
            except asyncio.CancelledError:
                pass
            self.monitor_task = None
        
        logger.info("Resource monitoring stopped")
    
    async def _monitor_loop(self) -> None:
        """Main monitoring loop."""
        while self.is_monitoring:
            try:
                resources = self.get_system_resources()
                
                # Check for critical conditions
                if self.is_memory_critical(resources.memory_percent):
                    logger.warning(f"Critical memory usage: {resources.memory_percent}%")
                
                if self.is_cpu_critical(resources.cpu_percent):
                    logger.warning(f"Critical CPU usage: {resources.cpu_percent}%")
                
                if self.is_disk_critical(resources.disk_percent):
                    logger.warning(f"Critical disk usage: {resources.disk_percent}%")
                
                # Call callback if set
                if self.callback:
                    try:
                        self.callback(resources)
                    except Exception as e:
                        logger.error(f"Error in monitoring callback: {e}")
                
                await asyncio.sleep(self.monitoring_interval)
                
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Error in monitoring loop: {e}")
                await asyncio.sleep(self.monitoring_interval)
    
    def get_peak_usage(self) -> Dict[str, float]:
        """Get peak resource usage during monitoring."""
        return {
            "memory_peak_percent": self.memory_peak,
            "cpu_peak_percent": self.cpu_peak,
            "disk_peak_percent": self.disk_peak
        }
    
    async def _initialize_impl(self) -> None:
        """Initialize the resource monitor."""
        logger.info("ResourceMonitor initialized")
    
    async def cleanup(self) -> None:
        """Cleanup resources."""
        await self.stop_monitoring()
        logger.info("ResourceMonitor cleanup completed")


def create_resource_monitor(
    memory_threshold: float = 80.0,
    cpu_threshold: float = 90.0,
    disk_threshold: float = 85.0
) -> ResourceMonitor:
    """Create a resource monitor instance."""
    return ResourceMonitor(
        memory_threshold=memory_threshold,
        cpu_threshold=cpu_threshold,
        disk_threshold=disk_threshold
    )
