"""Merging service for aligning diarization and transcription results.

This module provides advanced merging capabilities for combining speaker diarization
with transcription results, including conflict resolution, confidence scoring,
and edge case handling.
"""

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from .diarization_types import (
    DiarizationResult, SpeakerSegment, MergingConfig,
    MergingError, SegmentAlignmentError
)

logger = logging.getLogger(__name__)


@dataclass
class MergedSegment:
    """Represents a merged segment with speaker and transcription data."""
    
    start: float
    end: float
    text: str
    speaker_id: str
    confidence: float
    transcription_confidence: float
    diarization_confidence: float
    overlap_ratio: float
    overlapping_speakers: List[Dict[str, Any]]
    segment_type: str = "speech"  # speech, silence, music, etc.


@dataclass
class MergingResult:
    """Result of merging diarization and transcription."""
    
    segments: List[MergedSegment]
    speaker_count: int
    total_duration: float
    confidence_score: float
    processing_time: float
    metadata: Dict[str, Any]


class MergingService:
    """Advanced service for merging diarization and transcription results.
    
    Provides sophisticated algorithms for aligning speaker diarization with
    transcription segments, handling overlapping speech, and resolving conflicts.
    """
    
    def __init__(self, config: Optional[MergingConfig] = None):
        """Initialize the merging service.
        
        Args:
            config: Configuration for merging behavior
        """
        self.config = config or MergingConfig()
        logger.info(f"MergingService initialized with config: {self.config}")
    
    def merge_results(
        self,
        diarization_result: DiarizationResult,
        transcription_result: Dict[str, Any]
    ) -> MergingResult:
        """Merge diarization and transcription results.
        
        Args:
            diarization_result: Speaker diarization results
            transcription_result: Transcription results with segments
            
        Returns:
            MergedResult with aligned speaker-labeled segments
            
        Raises:
            MergingError: If merging fails
        """
        try:
            logger.info("Starting diarization-transcription merging")
            
            # Extract segments from transcription result
            transcription_segments = transcription_result.get("segments", [])
            if not transcription_segments:
                raise MergingError("No transcription segments found")
            
            # Align segments with advanced conflict resolution
            merged_segments = self._align_segments_advanced(
                diarization_result.segments,
                transcription_segments
            )
            
            # Post-process merged segments
            processed_segments = self._post_process_segments(merged_segments)
            
            # Calculate overall confidence and metadata
            confidence_score = self._calculate_overall_confidence(processed_segments)
            metadata = self._generate_metadata(
                diarization_result, transcription_result, processed_segments
            )
            
            result = MergingResult(
                segments=processed_segments,
                speaker_count=diarization_result.speaker_count,
                total_duration=diarization_result.audio_duration,
                confidence_score=confidence_score,
                processing_time=0.0,  # Will be set by caller
                metadata=metadata
            )
            
            logger.info(f"Merging completed: {len(processed_segments)} segments, "
                       f"{diarization_result.speaker_count} speakers, "
                       f"confidence: {confidence_score:.3f}")
            
            return result
            
        except Exception as e:
            logger.error(f"Failed to merge results: {e}")
            raise MergingError(f"Failed to merge results: {e}") from e
    
    def _align_segments_advanced(
        self,
        diarization_segments: List[SpeakerSegment],
        transcription_segments: List[Dict[str, Any]]
    ) -> List[MergedSegment]:
        """Advanced segment alignment with conflict resolution.
        
        Args:
            diarization_segments: Speaker diarization segments
            transcription_segments: Transcription segments
            
        Returns:
            List of merged segments with speaker labels
        """
        merged_segments = []
        
        for trans_segment in transcription_segments:
            trans_start = trans_segment.get("start", 0.0)
            trans_end = trans_segment.get("end", 0.0)
            trans_text = trans_segment.get("text", "").strip()
            trans_confidence = trans_segment.get("confidence", 0.0)
            
            # Skip empty segments
            if not trans_text:
                continue
            
            # Find overlapping speaker segments
            overlapping_speakers = self._find_overlapping_speakers(
                trans_start, trans_end, diarization_segments
            )
            
            # Resolve speaker assignment with confidence weighting
            primary_speaker, speaker_confidence, overlap_ratio = self._resolve_speaker_assignment(
                overlapping_speakers, trans_start, trans_end
            )
            
            # Create merged segment
            merged_segment = MergedSegment(
                start=trans_start,
                end=trans_end,
                text=trans_text,
                speaker_id=primary_speaker or "unknown",
                confidence=speaker_confidence,
                transcription_confidence=trans_confidence,
                diarization_confidence=overlap_ratio,
                overlap_ratio=overlap_ratio,
                overlapping_speakers=overlapping_speakers
            )
            
            merged_segments.append(merged_segment)
        
        return merged_segments
    
    def _find_overlapping_speakers(
        self,
        segment_start: float,
        segment_end: float,
        diarization_segments: List[SpeakerSegment]
    ) -> List[Dict[str, Any]]:
        """Find speaker segments that overlap with the given time range.
        
        Args:
            segment_start: Start time of the segment
            segment_end: End time of the segment
            diarization_segments: List of speaker diarization segments
            
        Returns:
            List of overlapping speakers with metadata
        """
        overlapping_speakers = []
        
        for diar_segment in diarization_segments:
            # Calculate overlap
            overlap_start = max(segment_start, diar_segment.start)
            overlap_end = min(segment_end, diar_segment.end)
            
            if overlap_end > overlap_start:
                overlap_duration = overlap_end - overlap_start
                segment_duration = segment_end - segment_start
                overlap_ratio = overlap_duration / segment_duration
                
                # Only include if overlap meets threshold
                if overlap_ratio >= self.config.min_overlap_ratio:
                    overlapping_speakers.append({
                        "speaker_id": diar_segment.speaker_id,
                        "confidence": diar_segment.confidence,
                        "overlap_ratio": overlap_ratio,
                        "overlap_duration": overlap_duration,
                        "segment_start": diar_segment.start,
                        "segment_end": diar_segment.end
                    })
        
        return overlapping_speakers
    
    def _resolve_speaker_assignment(
        self,
        overlapping_speakers: List[Dict[str, Any]],
        segment_start: float,
        segment_end: float
    ) -> Tuple[Optional[str], float, float]:
        """Resolve speaker assignment from overlapping speakers.
        
        Args:
            overlapping_speakers: List of overlapping speaker candidates
            segment_start: Start time of the segment
            segment_end: End time of the segment
            
        Returns:
            Tuple of (primary_speaker_id, confidence, overlap_ratio)
        """
        if not overlapping_speakers:
            return None, 0.0, 0.0
        
        # Sort by weighted score (overlap ratio * confidence)
        for speaker in overlapping_speakers:
            speaker["weighted_score"] = (
                speaker["overlap_ratio"] * speaker["confidence"]
            )
        
        overlapping_speakers.sort(
            key=lambda x: x["weighted_score"], reverse=True
        )
        
        primary_speaker = overlapping_speakers[0]
        
        # Check for conflicts (multiple speakers with similar scores)
        if len(overlapping_speakers) > 1:
            top_score = primary_speaker["weighted_score"]
            second_score = overlapping_speakers[1]["weighted_score"]
            
            # If scores are too close, mark as conflict
            if (top_score - second_score) / top_score < self.config.conflict_threshold:
                logger.warning(f"Speaker conflict detected: {primary_speaker['speaker_id']} "
                             f"vs {overlapping_speakers[1]['speaker_id']}")
                # Use the speaker with higher overlap ratio as tiebreaker
                overlapping_speakers.sort(
                    key=lambda x: x["overlap_ratio"], reverse=True
                )
                primary_speaker = overlapping_speakers[0]
        
        return (
            primary_speaker["speaker_id"],
            primary_speaker["confidence"],
            primary_speaker["overlap_ratio"]
        )
    
    def _post_process_segments(self, segments: List[MergedSegment]) -> List[MergedSegment]:
        """Post-process merged segments for consistency and quality.
        
        Args:
            segments: List of merged segments
            
        Returns:
            Post-processed segments
        """
        if not segments:
            return segments
        
        processed_segments = []
        i = 0
        
        while i < len(segments):
            segment = segments[i]
            segment_duration = segment.end - segment.start
            
            # Handle very short segments
            if segment_duration < self.config.min_segment_duration:
                logger.debug(f"Found short segment: {segment.start}-{segment.end} ({segment_duration}s) "
                           f"speaker: {segment.speaker_id}")
                
                # Try to merge with next segment if it has the same speaker
                if i < len(segments) - 1:
                    next_segment = segments[i + 1]
                    logger.debug(f"Next speaker: {next_segment.speaker_id}")
                    
                    if next_segment.speaker_id == segment.speaker_id:
                        # Merge with next
                        logger.debug(f"Merging short segment with next")
                        next_segment.text = segment.text + " " + next_segment.text
                        next_segment.start = segment.start
                        next_segment.confidence = (
                            segment.confidence + next_segment.confidence
                        ) / 2
                        i += 1  # Skip the current segment
                        continue
                
                # Try to merge with previous segment if it has the same speaker
                if i > 0:
                    prev_segment = segments[i - 1]
                    logger.debug(f"Prev speaker: {prev_segment.speaker_id}")
                    
                    if prev_segment.speaker_id == segment.speaker_id:
                        # Merge with previous
                        logger.debug(f"Merging short segment with previous")
                        prev_segment.text += " " + segment.text
                        prev_segment.end = segment.end
                        prev_segment.confidence = (
                            prev_segment.confidence + segment.confidence
                        ) / 2
                        i += 1  # Skip the current segment
                        continue
                
                logger.debug(f"No suitable merge target found for short segment")
            
            # Handle low-confidence segments
            if segment.confidence < self.config.min_confidence_threshold:
                segment.speaker_id = "unknown"
                segment.confidence = 0.0
            
            processed_segments.append(segment)
            i += 1
        
        return processed_segments
    
    def _calculate_overall_confidence(self, segments: List[MergedSegment]) -> float:
        """Calculate overall confidence score for the merged result.
        
        Args:
            segments: List of merged segments
            
        Returns:
            Overall confidence score (0.0 to 1.0)
        """
        if not segments:
            return 0.0
        
        # Weighted average of segment confidences
        total_weight = 0.0
        weighted_sum = 0.0
        
        for segment in segments:
            # Weight by segment duration
            duration = segment.end - segment.start
            weight = duration
            
            weighted_sum += segment.confidence * weight
            total_weight += weight
        
        if total_weight == 0.0:
            return 0.0
        
        return weighted_sum / total_weight
    
    def _generate_metadata(
        self,
        diarization_result: DiarizationResult,
        transcription_result: Dict[str, Any],
        merged_segments: List[MergedSegment]
    ) -> Dict[str, Any]:
        """Generate metadata for the merged result.
        
        Args:
            diarization_result: Original diarization result
            transcription_result: Original transcription result
            merged_segments: Final merged segments
            
        Returns:
            Metadata dictionary
        """
        # Count speakers in merged result
        speaker_ids = set(seg.speaker_id for seg in merged_segments)
        speaker_ids.discard("unknown")
        
        # Calculate statistics
        total_words = sum(len(seg.text.split()) for seg in merged_segments)
        avg_segment_duration = sum(
            seg.end - seg.start for seg in merged_segments
        ) / len(merged_segments) if merged_segments else 0.0
        
        return {
            "speaker_count_merged": len(speaker_ids),
            "speaker_count_original": diarization_result.speaker_count,
            "total_words": total_words,
            "total_segments": len(merged_segments),
            "average_segment_duration": avg_segment_duration,
            "unknown_speaker_segments": sum(
                1 for seg in merged_segments if seg.speaker_id == "unknown"
            ),
            "diarization_confidence": diarization_result.confidence_score,
            "transcription_confidence": transcription_result.get("accuracy_estimate", 0.0),
            "merging_config": {
                "min_overlap_ratio": self.config.min_overlap_ratio,
                "min_confidence_threshold": self.config.min_confidence_threshold,
                "min_segment_duration": self.config.min_segment_duration,
                "conflict_threshold": self.config.conflict_threshold
            }
        }
