"""Utility functions for diarization services."""

import gc
import json
import logging
import time
from pathlib import Path
from typing import Any, Dict, List

import numpy as np
import psutil
import torch
from sklearn.metrics.pairwise import cosine_similarity

from .diarization_types import (
    SpeakerSegment, ProfileMatch, SpeakerProfile, 
    ProfileValidationError, ProfileNotFoundError
)

logger = logging.getLogger(__name__)


def determine_device() -> str:
    """Determine the best device for processing."""
    if not torch.cuda.is_available():
        return "cpu"
    
    try:
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
        if gpu_memory < 4:  # Less than 4GB GPU memory
            return "cpu"
        return "cuda"
    except Exception:
        return "cpu"


def check_memory_before_loading(device: str, threshold_mb: int = 6000):
    """Check available memory before loading model."""
    if device == "cpu":
        memory = psutil.virtual_memory()
        available_mb = memory.available / 1024**2
        
        if available_mb < threshold_mb:
            logger.warning(f"Low memory available: {available_mb:.1f}MB")
            gc.collect()
            
            memory = psutil.virtual_memory()
            available_mb = memory.available / 1024**2
            
            if available_mb < threshold_mb:
                raise MemoryError(f"Insufficient memory: {available_mb:.1f}MB available")


def apply_memory_optimizations(pipeline: Any):
    """Apply memory optimization techniques."""
    if pipeline is None:
        return
    
    try:
        if hasattr(pipeline, 'model') and hasattr(pipeline.model, 'gradient_checkpointing_enable'):
            pipeline.model.gradient_checkpointing_enable()
        
        if hasattr(pipeline, 'model'):
            pipeline.model.eval()
        
        logger.debug("Applied memory optimizations to diarization pipeline")
        
    except Exception as e:
        logger.warning(f"Failed to apply memory optimizations: {e}")


def convert_annotation_to_segments(annotation: Any, min_duration: float = 0.5) -> List[SpeakerSegment]:
    """Convert Pyannote annotation to our segment format."""
    segments = []
    
    for segment, track, label in annotation.itertracks(yield_label=True):
        if segment.duration < min_duration:
            continue
        
        confidence = 0.8  # Default confidence
        if hasattr(annotation, 'confidence') and track in annotation.confidence:
            confidence = float(annotation.confidence[track])
        
        speaker_segment = SpeakerSegment(
            start=float(segment.start),
            end=float(segment.end),
            speaker_id=label,
            confidence=confidence
        )
        segments.append(speaker_segment)
    
    return segments


def calculate_confidence(segments: List[SpeakerSegment]) -> float:
    """Calculate overall confidence score for diarization."""
    if not segments:
        return 0.0
    
    total_duration = 0.0
    weighted_confidence = 0.0
    
    for segment in segments:
        duration = segment.end - segment.start
        total_duration += duration
        weighted_confidence += segment.confidence * duration
    
    return weighted_confidence / total_duration if total_duration > 0 else 0.0


def find_similar_speakers(
    query_embedding: np.ndarray,
    embeddings_cache: Dict[str, np.ndarray],
    profiles: Dict[str, SpeakerProfile],
    threshold: float = 0.7
) -> List[ProfileMatch]:
    """Find speakers with similar embeddings."""
    if not embeddings_cache:
        return []
    
    matches = []
    query_embedding_2d = query_embedding.reshape(1, -1)
    
    for speaker_id, profile_embedding in embeddings_cache.items():
        if profile_embedding is None:
            continue
        
        try:
            profile_embedding_2d = profile_embedding.reshape(1, -1)
            similarity = cosine_similarity(query_embedding_2d, profile_embedding_2d)[0][0]
            
            if similarity >= threshold:
                profile = profiles[speaker_id]
                match = ProfileMatch(
                    speaker_id=speaker_id,
                    similarity_score=similarity,
                    confidence=similarity,
                    profile=profile
                )
                matches.append(match)
                
        except Exception as e:
            logger.warning(f"Failed to calculate similarity for {speaker_id}: {e}")
    
    matches.sort(key=lambda x: x.similarity_score, reverse=True)
    return matches


def validate_speaker_profile(speaker_id: str, embedding: np.ndarray):
    """Validate speaker profile data."""
    if not speaker_id or not speaker_id.strip():
        raise ProfileValidationError("Speaker ID cannot be empty")
    
    if embedding is None or embedding.size == 0:
        raise ProfileValidationError("Embedding cannot be empty")


def save_profile_to_disk(profile: SpeakerProfile, storage_dir: Path):
    """Save a single profile to disk."""
    try:
        profile_file = storage_dir / f"{profile.speaker_id}.json"
        
        with open(profile_file, 'w') as f:
            json.dump(profile.to_dict(), f, indent=2)
            
    except Exception as e:
        logger.error(f"Failed to save profile {profile.speaker_id}: {e}")


def load_profile_from_disk(profile_file: Path) -> SpeakerProfile:
    """Load a profile from disk."""
    try:
        with open(profile_file, 'r') as f:
            data = json.load(f)
        
        return SpeakerProfile.from_dict(data)
        
    except Exception as e:
        logger.warning(f"Failed to load profile from {profile_file}: {e}")
        raise


def align_segments(
    diarization_segments: List[Any],
    transcription_segments: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
    """Align diarization segments with transcription segments."""
    merged_segments = []
    
    for trans_segment in transcription_segments:
        trans_start = trans_segment.get("start", 0.0)
        trans_end = trans_segment.get("end", 0.0)
        
        overlapping_speakers = []
        
        for diar_segment in diarization_segments:
            diar_start = diar_segment.start
            diar_end = diar_segment.end
            
            overlap_start = max(trans_start, diar_start)
            overlap_end = min(trans_end, diar_end)
            
            if overlap_end > overlap_start:
                overlap_duration = overlap_end - overlap_start
                segment_duration = trans_end - trans_start
                overlap_ratio = overlap_duration / segment_duration
                
                if overlap_ratio > 0.5:  # More than 50% overlap
                    overlapping_speakers.append({
                        "speaker_id": diar_segment.speaker_id,
                        "confidence": diar_segment.confidence,
                        "overlap_ratio": overlap_ratio
                    })
        
        primary_speaker = None
        if overlapping_speakers:
            overlapping_speakers.sort(
                key=lambda x: (x["overlap_ratio"], x["confidence"]),
                reverse=True
            )
            primary_speaker = overlapping_speakers[0]["speaker_id"]
        
        merged_segment = {
            "start": trans_start,
            "end": trans_end,
            "text": trans_segment.get("text", ""),
            "speaker_id": primary_speaker or "unknown",
            "confidence": trans_segment.get("confidence", 0.0),
            "overlapping_speakers": overlapping_speakers
        }
        
        merged_segments.append(merged_segment)
    
    return merged_segments


def cleanup_resources(pipeline: Any = None):
    """Clean up resources and free memory."""
    if pipeline is not None:
        del pipeline
    
    gc.collect()
    
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def measure_processing_time(func):
    """Decorator to measure processing time."""
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        processing_time = time.time() - start_time
        
        if hasattr(result, 'processing_time'):
            result.processing_time = processing_time
        
        return result
    return wrapper
