"""Shared types and data structures for diarization services."""

import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Protocol, runtime_checkable

import numpy as np

logger = logging.getLogger(__name__)


class DiarizationError(Exception):
    """Base exception for diarization errors."""
    pass


class ModelLoadingError(DiarizationError):
    """Exception raised when diarization model fails to load."""
    pass


class AudioProcessingError(DiarizationError):
    """Exception raised when audio processing fails."""
    pass


class SpeakerProfileError(DiarizationError):
    """Base exception for speaker profile errors."""
    pass


class ProfileNotFoundError(SpeakerProfileError):
    """Exception raised when a speaker profile is not found."""
    pass


class ProfileValidationError(SpeakerProfileError):
    """Exception raised when a speaker profile is invalid."""
    pass


class ParallelProcessingError(DiarizationError):
    """Base exception for parallel processing errors."""
    pass


class WorkerPoolError(ParallelProcessingError):
    """Exception raised when worker pool operations fail."""
    pass


class ResultMergingError(ParallelProcessingError):
    """Exception raised when merging parallel results fails."""
    pass


class MergingError(DiarizationError):
    """Exception raised when merging diarization and transcription fails."""
    pass


class SegmentAlignmentError(MergingError):
    """Exception raised when segment alignment fails."""
    pass


@dataclass
class SpeakerSegment:
    """Represents a speaker segment with timing and confidence."""
    
    start: float
    end: float
    speaker_id: str
    confidence: float
    speaker_label: Optional[str] = None


@dataclass
class DiarizationResult:
    """Result of diarization processing."""
    
    segments: List[SpeakerSegment]
    speaker_count: int
    processing_time: float
    confidence_score: float
    model_used: str
    audio_duration: float


@dataclass
class DiarizationConfig:
    """Configuration for diarization processing with comprehensive optimization options."""
    
    # Model configuration
    model_path: str = "pyannote/speaker-diarization-3.0"
    use_auth_token: bool = True
    device: str = "auto"
    
    # Speaker count configuration
    num_speakers: Optional[int] = None  # Fixed speaker count (overrides auto-detection)
    min_speakers: Optional[int] = None  # Minimum speaker count for auto-detection
    max_speakers: Optional[int] = None  # Maximum speaker count for auto-detection
    enable_speaker_estimation: bool = True  # Enable automatic speaker count estimation
    speaker_estimation_confidence: float = 0.8  # Confidence threshold for speaker estimation
    
    # Quality thresholds
    threshold: float = 0.5  # Diarization threshold
    min_duration: float = 0.5  # Minimum segment duration
    quality_threshold: float = 0.7  # Overall quality threshold
    confidence_threshold: float = 0.6  # Minimum confidence for speaker assignment
    
    # Memory optimization settings
    memory_optimization: bool = True  # Enable memory optimizations
    max_memory_gb: float = 8.0  # Maximum memory usage in GB
    memory_safety_margin: float = 0.2  # Safety margin for memory usage
    enable_gradient_checkpointing: bool = True  # Enable gradient checkpointing
    enable_quantization: bool = True  # Enable model quantization
    enable_model_offloading: bool = True  # Enable model offloading to CPU
    enable_precision_optimization: bool = True  # Enable precision optimization
    
    # Audio processing configuration
    batch_size: int = 1  # Processing batch size
    enable_audio_downsampling: bool = True  # Enable audio downsampling for memory efficiency
    target_sample_rate: int = 16000  # Target sample rate for processing
    enable_chunking: bool = True  # Enable audio chunking for large files
    chunk_duration_seconds: int = 600  # Chunk duration in seconds (10 minutes)
    chunk_overlap_seconds: int = 30  # Overlap between chunks in seconds
    
    # Resource management
    enable_resource_cleanup: bool = True  # Enable automatic resource cleanup
    cleanup_interval_seconds: int = 60  # Cleanup interval in seconds
    max_processing_time_seconds: int = 3600  # Maximum processing time (1 hour)
    
    # Integration settings
    integrate_with_model_manager: bool = True  # Integrate with ModelManager singleton
    enable_caching: bool = True  # Enable result caching
    cache_ttl_seconds: int = 3600  # Cache TTL in seconds


@dataclass
class MergingConfig:
    """Configuration for merging diarization and transcription results."""
    
    min_overlap_ratio: float = 0.5  # Minimum overlap ratio to consider speaker assignment
    min_confidence_threshold: float = 0.3  # Minimum confidence for speaker assignment
    min_segment_duration: float = 0.5  # Minimum segment duration in seconds
    conflict_threshold: float = 0.1  # Threshold for detecting speaker conflicts
    enable_post_processing: bool = True  # Enable post-processing of merged segments
    merge_short_segments: bool = True  # Merge very short segments with adjacent ones
    unknown_speaker_label: str = "unknown"  # Label for unknown speakers


@dataclass
class SpeakerProfile:
    """Represents a speaker profile with embeddings and metadata."""
    
    speaker_id: str
    name: Optional[str] = None
    embedding: Optional[np.ndarray] = None
    segments: List[Dict[str, Any]] = None
    confidence_scores: List[float] = None
    created_at: datetime = None
    updated_at: datetime = None
    version: str = "1.0"
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        """Initialize default values."""
        if self.segments is None:
            self.segments = []
        if self.confidence_scores is None:
            self.confidence_scores = []
        if self.created_at is None:
            self.created_at = datetime.now(timezone.utc)
        if self.updated_at is None:
            self.updated_at = self.created_at
        if self.metadata is None:
            self.metadata = {}
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert profile to dictionary for serialization."""
        from dataclasses import asdict
        data = asdict(self)
        if self.embedding is not None:
            data['embedding'] = self.embedding.tolist()
        data['created_at'] = self.created_at.isoformat()
        data['updated_at'] = self.updated_at.isoformat()
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'SpeakerProfile':
        """Create profile from dictionary."""
        if 'embedding' in data and data['embedding'] is not None:
            data['embedding'] = np.array(data['embedding'])
        if 'created_at' in data:
            data['created_at'] = datetime.fromisoformat(data['created_at'])
        if 'updated_at' in data:
            data['updated_at'] = datetime.fromisoformat(data['updated_at'])
        return cls(**data)


@dataclass
class ProfileMatch:
    """Represents a match between a query and a speaker profile."""
    
    speaker_id: str
    similarity_score: float
    confidence: float
    profile: SpeakerProfile


@dataclass
class ProcessingTask:
    """Represents a processing task with metadata."""
    
    task_id: str
    audio_path: Path
    task_type: str  # 'diarization' or 'transcription'
    config: Any
    priority: int = 0
    created_at: float = None
    
    def __post_init__(self):
        import time
        if self.created_at is None:
            self.created_at = time.time()


@dataclass
class ProcessingResult:
    """Result of parallel processing."""
    
    task_id: str
    diarization_result: Optional[DiarizationResult] = None
    transcription_result: Optional[Any] = None
    merged_result: Optional[Dict[str, Any]] = None
    processing_time: float = 0.0
    success: bool = True
    error_message: Optional[str] = None


@dataclass
class ParallelProcessingConfig:
    """Configuration for parallel processing."""
    
    max_workers: int = 2
    timeout_seconds: int = 300  # 5 minutes
    enable_progress_tracking: bool = True
    memory_limit_mb: int = 6000  # 6GB
    retry_failed_tasks: bool = True
    max_retries: int = 2
    worker_cleanup_interval: int = 60  # seconds


@runtime_checkable
class DiarizationServiceProtocol(Protocol):
    """Protocol for diarization services."""
    
    def process_audio(
        self, 
        audio_path: Path, 
        config: Optional[DiarizationConfig] = None
    ) -> DiarizationResult:
        """Process audio file for speaker diarization."""
        ...
    
    def estimate_speaker_count(self, audio_path: Path) -> int:
        """Estimate the number of speakers in audio."""
        ...
    
    def get_speaker_segments(
        self, 
        audio_path: Path, 
        speaker_id: str
    ) -> List[SpeakerSegment]:
        """Get all segments for a specific speaker."""
        ...


@runtime_checkable
class SpeakerProfileManagerProtocol(Protocol):
    """Protocol for speaker profile managers."""
    
    def add_speaker(self, speaker_id: str, embedding: np.ndarray, **kwargs) -> SpeakerProfile:
        """Add a new speaker profile."""
        ...
    
    def get_speaker(self, speaker_id: str) -> Optional[SpeakerProfile]:
        """Get a speaker profile by ID."""
        ...
    
    def find_similar_speakers(self, embedding: np.ndarray, threshold: float = 0.7) -> List[ProfileMatch]:
        """Find speakers with similar embeddings."""
        ...
    
    def update_speaker(self, speaker_id: str, embedding: np.ndarray, **kwargs) -> SpeakerProfile:
        """Update an existing speaker profile."""
        ...
    
    def remove_speaker(self, speaker_id: str) -> bool:
        """Remove a speaker profile."""
        ...
    
    def save_profiles(self, file_path: Path) -> bool:
        """Save all profiles to disk."""
        ...
    
    def load_profiles(self, file_path: Path) -> bool:
        """Load profiles from disk."""
        ...


@runtime_checkable
class ParallelProcessorProtocol(Protocol):
    """Protocol for parallel processors."""
    
    def process_file(
        self, 
        audio_path: Path,
        diarization_config: Optional[DiarizationConfig] = None,
        transcription_config: Optional[Any] = None
    ) -> ProcessingResult:
        """Process a single file with parallel diarization and transcription."""
        ...
    
    def process_batch(
        self, 
        audio_paths: List[Path],
        configs: Optional[Dict[str, Any]] = None
    ) -> List[ProcessingResult]:
        """Process multiple files in parallel."""
        ...
    
    def get_processing_stats(self) -> Dict[str, Any]:
        """Get statistics about processing performance."""
        ...
