"""Diarization service for Trax platform.

This module provides speaker diarization functionality using Pyannote.audio
with support for parallel processing, speaker profiles, and memory optimization.
"""

import logging
import threading
import time
from pathlib import Path
from typing import Any, List, Optional

import torch
from pyannote.audio import Pipeline
from pyannote.audio.pipelines.utils.hook import ProgressHook

from ..base.services import BaseService
from ..config import config
from .diarization_types import (
    DiarizationConfig, DiarizationResult, SpeakerSegment,
    DiarizationServiceProtocol, ModelLoadingError, AudioProcessingError
)
from .diarization_utils import (
    determine_device, check_memory_before_loading, apply_memory_optimizations,
    convert_annotation_to_segments, calculate_confidence, cleanup_resources
)

logger = logging.getLogger(__name__)


class DiarizationManager(BaseService):
    """Manages speaker diarization using Pyannote.audio.
    
    Provides efficient diarization with model caching, memory optimization,
    and integration with the ModelManager singleton.
    """
    
    def __init__(self, config: Optional[DiarizationConfig] = None):
        """Initialize the DiarizationManager.
        
        Args:
            config: Configuration for diarization processing
        """
        super().__init__(name="DiarizationManager")
        self.config = config or DiarizationConfig()
        self._pipeline: Optional[Pipeline] = None
        self._lock = threading.Lock()
        self._initialized = False
        
        # Memory management
        self._memory_threshold_mb = 6000  # 6GB threshold
        self._device = determine_device()
        
        logger.info(f"DiarizationManager initialized with device: {self._device}")
    
    async def _initialize_impl(self) -> None:
        """Initialize the diarization pipeline."""
        try:
            # Load pipeline on initialization
            self._load_pipeline()
            self._initialized = True
            logger.info("DiarizationManager initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize DiarizationManager: {e}")
            raise
    
    def _load_pipeline(self) -> Pipeline:
        """Load the Pyannote.audio pipeline with error handling."""
        if self._pipeline is not None:
            return self._pipeline
        
        with self._lock:
            if self._pipeline is not None:
                return self._pipeline
            
            try:
                logger.info(f"Loading diarization model: {self.config.model_path}")
                
                # Check memory before loading
                check_memory_before_loading(self._device, self._memory_threshold_mb)
                
                # Load pipeline with configuration
                pipeline_kwargs = {
                    "model_path": self.config.model_path,
                    "device": self._device,
                }
                
                if self.config.use_auth_token and hasattr(config, 'HUGGINGFACE_TOKEN'):
                    pipeline_kwargs["use_auth_token"] = config.HUGGINGFACE_TOKEN
                
                self._pipeline = Pipeline.from_pretrained(**pipeline_kwargs)
                
                # Apply memory optimizations if enabled
                if self.config.memory_optimization:
                    apply_memory_optimizations(self._pipeline)
                
                self._initialized = True
                logger.info("Diarization pipeline loaded successfully")
                
                return self._pipeline
                
            except Exception as e:
                logger.error(f"Failed to load diarization model: {e}")
                raise ModelLoadingError(f"Failed to load diarization model: {e}") from e
    
    def process_audio(
        self, 
        audio_path: Path, 
        config: Optional[DiarizationConfig] = None
    ) -> DiarizationResult:
        """Process audio file for speaker diarization.
        
        Args:
            audio_path: Path to the audio file
            config: Optional configuration override
            
        Returns:
            DiarizationResult with speaker segments and metadata
            
        Raises:
            AudioProcessingError: If processing fails
            FileNotFoundError: If audio file doesn't exist
        """
        if not audio_path.exists():
            raise FileNotFoundError(f"Audio file not found: {audio_path}")
        
        # Use provided config or default
        processing_config = config or self.config
        
        try:
            logger.info(f"Starting diarization for: {audio_path.name}")
            
            # Load pipeline if needed
            pipeline = self._load_pipeline()
            
            # Prepare pipeline parameters
            pipeline_params = {
                "min_speakers": processing_config.min_speakers,
                "max_speakers": processing_config.max_speakers,
            }
            
            if processing_config.num_speakers is not None:
                pipeline_params["num_speakers"] = processing_config.num_speakers
            
            # Process audio with progress tracking
            start_time = time.time()
            
            with ProgressHook() as hook:
                diarization = pipeline(audio_path, hook=hook, **pipeline_params)
            
            processing_time = time.time() - start_time
            
            # Convert to our format
            segments = convert_annotation_to_segments(diarization, self.config.min_duration)
            
            # Calculate confidence and metadata
            confidence_score = calculate_confidence(segments)
            speaker_count = len(set(seg.speaker_id for seg in segments))
            audio_duration = max(seg.end for seg in segments) if segments else 0.0
            
            result = DiarizationResult(
                segments=segments,
                speaker_count=speaker_count,
                processing_time=processing_time,
                confidence_score=confidence_score,
                model_used=self.config.model_path,
                audio_duration=audio_duration
            )
            
            logger.info(f"Diarization completed: {speaker_count} speakers, "
                       f"{len(segments)} segments, {processing_time:.2f}s")
            
            return result
            
        except Exception as e:
            logger.error(f"Diarization failed for {audio_path}: {e}")
            raise AudioProcessingError(f"Diarization failed: {e}") from e
    
    def estimate_speaker_count(self, audio_path: Path) -> int:
        """Estimate the number of speakers in audio.
        
        This is a simplified estimation that can be overridden
        with more sophisticated methods.
        
        Args:
            audio_path: Path to the audio file
            
        Returns:
            Estimated number of speakers
        """
        try:
            # Use a quick pass with auto speaker detection
            quick_config = DiarizationConfig(
                model_path=self.config.model_path,
                min_speakers=1,
                max_speakers=10,  # Reasonable upper limit
                threshold=0.3  # Lower threshold for estimation
            )
            
            result = self.process_audio(audio_path, quick_config)
            return result.speaker_count
            
        except Exception as e:
            logger.warning(f"Speaker count estimation failed: {e}")
            return 2  # Default fallback
    
    def get_speaker_segments(
        self, 
        audio_path: Path, 
        speaker_id: str
    ) -> List[SpeakerSegment]:
        """Get all segments for a specific speaker.
        
        Args:
            audio_path: Path to the audio file
            speaker_id: ID of the speaker to extract
            
        Returns:
            List of segments for the specified speaker
        """
        result = self.process_audio(audio_path)
        return [seg for seg in result.segments if seg.speaker_id == speaker_id]
    
    def cleanup(self):
        """Clean up resources and free memory."""
        with self._lock:
            cleanup_resources(self._pipeline)
            self._pipeline = None
            self._initialized = False
            logger.info("DiarizationManager resources cleaned up")

