"""Configuration manager for diarization services with advanced optimization features."""

import logging
import psutil
import torch
import numpy as np
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, List
from dataclasses import dataclass
from datetime import datetime, timezone

from .diarization_types import DiarizationConfig, MergingConfig
from .memory_optimization import MemoryConfig, MemoryOptimizer

logger = logging.getLogger(__name__)


@dataclass
class SystemResources:
    """System resource information."""
    
    total_memory_gb: float
    available_memory_gb: float
    cpu_count: int
    gpu_available: bool
    gpu_memory_gb: Optional[float] = None
    gpu_name: Optional[str] = None


@dataclass
class OptimizationRecommendations:
    """Optimization recommendations based on system resources."""
    
    recommended_batch_size: int
    recommended_chunk_duration: int
    enable_quantization: bool
    enable_offloading: bool
    enable_chunking: bool
    target_sample_rate: int
    memory_optimizations: List[str]


class DiarizationConfigManager:
    """Manages diarization configuration with automatic optimization and resource management."""
    
    def __init__(self, base_config: Optional[DiarizationConfig] = None):
        """Initialize the configuration manager.
        
        Args:
            base_config: Base configuration to start with
        """
        self.base_config = base_config or DiarizationConfig()
        self.system_resources = self._analyze_system_resources()
        self.memory_optimizer = MemoryOptimizer()
        
        logger.info(f"Config manager initialized with {self.system_resources.total_memory_gb:.1f}GB total memory")
    
    def _analyze_system_resources(self) -> SystemResources:
        """Analyze available system resources."""
        # Get system memory
        memory = psutil.virtual_memory()
        total_memory_gb = memory.total / (1024**3)
        available_memory_gb = memory.available / (1024**3)
        
        # Get CPU info
        cpu_count = psutil.cpu_count()
        
        # Check GPU availability
        gpu_available = torch.cuda.is_available()
        gpu_memory_gb = None
        gpu_name = None
        
        if gpu_available:
            gpu_memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
            gpu_name = torch.cuda.get_device_name(0)
        
        return SystemResources(
            total_memory_gb=total_memory_gb,
            available_memory_gb=available_memory_gb,
            cpu_count=cpu_count,
            gpu_available=gpu_available,
            gpu_memory_gb=gpu_memory_gb,
            gpu_name=gpu_name
        )
    
    def get_optimization_recommendations(self) -> OptimizationRecommendations:
        """Get optimization recommendations based on system resources."""
        available_memory_gb = self.system_resources.available_memory_gb
        
        # Calculate recommended batch size
        if available_memory_gb >= 16:
            recommended_batch_size = 4
        elif available_memory_gb >= 8:
            recommended_batch_size = 2
        else:
            recommended_batch_size = 1
        
        # Calculate recommended chunk duration
        if available_memory_gb >= 12:
            recommended_chunk_duration = 900  # 15 minutes
        elif available_memory_gb >= 6:
            recommended_chunk_duration = 600  # 10 minutes
        else:
            recommended_chunk_duration = 300  # 5 minutes
        
        # Determine optimization strategies
        enable_quantization = available_memory_gb < 12
        enable_offloading = available_memory_gb < 8
        enable_chunking = available_memory_gb < 16
        
        # Determine target sample rate
        if available_memory_gb >= 8:
            target_sample_rate = 16000
        else:
            target_sample_rate = 8000  # Lower quality but more memory efficient
        
        # List memory optimizations
        memory_optimizations = []
        if enable_quantization:
            memory_optimizations.append("quantization")
        if enable_offloading:
            memory_optimizations.append("model_offloading")
        if enable_chunking:
            memory_optimizations.append("audio_chunking")
        if target_sample_rate < 16000:
            memory_optimizations.append("downsampling")
        
        return OptimizationRecommendations(
            recommended_batch_size=recommended_batch_size,
            recommended_chunk_duration=recommended_chunk_duration,
            enable_quantization=enable_quantization,
            enable_offloading=enable_offloading,
            enable_chunking=enable_chunking,
            target_sample_rate=target_sample_rate,
            memory_optimizations=memory_optimizations
        )
    
    def create_optimized_config(self, audio_duration_seconds: Optional[float] = None) -> DiarizationConfig:
        """Create an optimized configuration based on system resources and audio characteristics.
        
        Args:
            audio_duration_seconds: Duration of audio file in seconds (for chunking decisions)
            
        Returns:
            Optimized diarization configuration
        """
        recommendations = self.get_optimization_recommendations()
        
        # Start with base configuration
        config = DiarizationConfig()
        
        # Apply system-based optimizations
        config.batch_size = recommendations.recommended_batch_size
        config.enable_quantization = recommendations.enable_quantization
        config.enable_model_offloading = recommendations.enable_offloading
        config.enable_chunking = recommendations.enable_chunking
        config.target_sample_rate = recommendations.target_sample_rate
        config.chunk_duration_seconds = recommendations.recommended_chunk_duration
        
        # Apply memory limits based on system
        config.max_memory_gb = min(
            self.base_config.max_memory_gb,
            self.system_resources.available_memory_gb * 0.8  # Use 80% of available memory
        )
        
        # Adjust chunking based on audio duration
        if audio_duration_seconds and audio_duration_seconds > config.chunk_duration_seconds * 2:
            config.enable_chunking = True
        elif audio_duration_seconds and audio_duration_seconds <= config.chunk_duration_seconds:
            config.enable_chunking = False
        
        # Set device based on GPU availability
        if self.system_resources.gpu_available and self.system_resources.gpu_memory_gb >= 4:
            config.device = "cuda"
        else:
            config.device = "cpu"
        
        logger.info(f"Created optimized config: batch_size={config.batch_size}, "
                   f"chunking={config.enable_chunking}, device={config.device}")
        
        return config
    
    def estimate_speaker_count(self, audio_path: Path, config: DiarizationConfig) -> Optional[int]:
        """Estimate the number of speakers in an audio file.
        
        Args:
            audio_path: Path to the audio file
            config: Diarization configuration
            
        Returns:
            Estimated number of speakers, or None if estimation fails
        """
        if not config.enable_speaker_estimation:
            return config.num_speakers
        
        try:
            # This is a simplified estimation - in practice, you might use
            # a lightweight model or heuristic based on audio characteristics
            import librosa
            
            # Load audio for analysis
            y, sr = librosa.load(audio_path, sr=config.target_sample_rate, duration=60)  # Analyze first minute
            
            # Simple heuristic: more complex audio patterns suggest more speakers
            # This is a basic implementation - real estimation would be more sophisticated
            
            # Calculate spectral features
            spectral_centroids = librosa.feature.spectral_centroid(y=y, sr=sr)
            spectral_rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)
            
            # Calculate complexity metrics
            spectral_complexity = np.std(spectral_centroids) + np.std(spectral_rolloff)
            
            # Estimate speakers based on complexity (simplified heuristic)
            if spectral_complexity < 0.1:
                estimated_speakers = 1
            elif spectral_complexity < 0.2:
                estimated_speakers = 2
            elif spectral_complexity < 0.3:
                estimated_speakers = 3
            else:
                estimated_speakers = 4
            
            # Apply constraints
            if config.min_speakers:
                estimated_speakers = max(estimated_speakers, config.min_speakers)
            if config.max_speakers:
                estimated_speakers = min(estimated_speakers, config.max_speakers)
            
            logger.info(f"Estimated {estimated_speakers} speakers based on audio complexity")
            return estimated_speakers
            
        except Exception as e:
            logger.warning(f"Failed to estimate speaker count: {e}")
            return config.num_speakers
    
    def validate_config(self, config: DiarizationConfig) -> Tuple[bool, List[str]]:
        """Validate configuration against system resources.
        
        Args:
            config: Configuration to validate
            
        Returns:
            Tuple of (is_valid, list_of_warnings)
        """
        warnings = []
        
        # Check memory requirements
        required_memory_gb = config.max_memory_gb
        if required_memory_gb > self.system_resources.available_memory_gb:
            warnings.append(f"Required memory ({required_memory_gb:.1f}GB) exceeds available memory "
                          f"({self.system_resources.available_memory_gb:.1f}GB)")
        
        # Check batch size
        if config.batch_size > 4:
            warnings.append("Large batch size may cause memory issues")
        
        # Check chunk duration
        if config.chunk_duration_seconds > 1800:  # 30 minutes
            warnings.append("Very long chunk duration may cause memory issues")
        
        # Check device compatibility
        if config.device == "cuda" and not self.system_resources.gpu_available:
            warnings.append("CUDA device requested but GPU not available")
        
        # Config is valid unless there are critical errors (memory issues)
        critical_errors = [w for w in warnings if "memory" in w.lower() and "exceeds" in w.lower()]
        is_valid = len(critical_errors) == 0
        return is_valid, warnings
    
    def get_memory_usage_estimate(self, config: DiarizationConfig, audio_duration_seconds: float) -> Dict[str, float]:
        """Estimate memory usage for processing.
        
        Args:
            config: Diarization configuration
            audio_duration_seconds: Duration of audio file
            
        Returns:
            Dictionary with memory usage estimates
        """
        # Base model memory (Pyannote.audio model)
        model_memory_gb = 2.0
        
        # Audio data memory (depends on sample rate and duration)
        audio_memory_gb = (config.target_sample_rate * audio_duration_seconds * 4) / (1024**3)  # 4 bytes per sample
        
        # Processing overhead
        processing_overhead_gb = 1.0
        
        # Chunking overhead
        if config.enable_chunking:
            chunk_overhead_gb = 0.5
        else:
            chunk_overhead_gb = 0.0
        
        # Quantization savings
        if config.enable_quantization:
            model_memory_gb *= 0.5  # 50% reduction
        
        total_memory_gb = model_memory_gb + audio_memory_gb + processing_overhead_gb + chunk_overhead_gb
        
        return {
            "model_memory_gb": model_memory_gb,
            "audio_memory_gb": audio_memory_gb,
            "processing_overhead_gb": processing_overhead_gb,
            "chunk_overhead_gb": chunk_overhead_gb,
            "total_memory_gb": total_memory_gb,
            "available_memory_gb": self.system_resources.available_memory_gb
        }
    
    def create_merging_config(self, diarization_config: DiarizationConfig) -> MergingConfig:
        """Create optimized merging configuration based on diarization config.
        
        Args:
            diarization_config: Diarization configuration
            
        Returns:
            Optimized merging configuration
        """
        # Adjust merging parameters based on diarization quality
        if diarization_config.quality_threshold > 0.8:
            # High quality diarization - can use stricter thresholds
            min_overlap_ratio = 0.6
            min_confidence_threshold = 0.5
        elif diarization_config.quality_threshold > 0.6:
            # Medium quality - balanced thresholds
            min_overlap_ratio = 0.5
            min_confidence_threshold = 0.4
        else:
            # Lower quality - more lenient thresholds
            min_overlap_ratio = 0.4
            min_confidence_threshold = 0.3
        
        return MergingConfig(
            min_overlap_ratio=min_overlap_ratio,
            min_confidence_threshold=min_confidence_threshold,
            min_segment_duration=diarization_config.min_duration,
            conflict_threshold=0.1,
            enable_post_processing=True,
            merge_short_segments=True
        )
