"""Utility functions for parallel processing operations."""

import concurrent.futures
import logging
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from .diarization_types import (
    DiarizationConfig, DiarizationResult, ProcessingResult,
    ParallelProcessingError, ResultMergingError
)
from .diarization_utils import align_segments

logger = logging.getLogger(__name__)


def initialize_services(diarization_manager, transcription_service):
    """Initialize diarization and transcription services."""
    if diarization_manager is None:
        from .diarization_service import DiarizationManager
        diarization_manager = DiarizationManager()
    
    if transcription_service is None:
        from .transcription_service import TranscriptionService
        transcription_service = TranscriptionService()
    
    return diarization_manager, transcription_service


def process_diarization_task(
    audio_path: Path,
    diarization_manager,
    config: Optional[DiarizationConfig] = None
) -> Optional[DiarizationResult]:
    """Process diarization in a separate thread."""
    try:
        return diarization_manager.process_audio(audio_path, config)
    except Exception as e:
        logger.error(f"Diarization failed for {audio_path}: {e}")
        return None


def process_transcription_task(
    audio_path: Path,
    transcription_service,
    config: Optional[Any] = None
) -> Optional[Any]:
    """Process transcription in a separate thread."""
    try:
        # Note: This assumes TranscriptionService has a transcribe_file method
        return transcription_service.transcribe_file(audio_path, config)
    except Exception as e:
        logger.error(f"Transcription failed for {audio_path}: {e}")
        return None


def merge_processing_results(
    diarization_result: Optional[DiarizationResult],
    transcription_result: Optional[Any]
) -> Dict[str, Any]:
    """Merge diarization and transcription results.
    
    Args:
        diarization_result: Result from diarization processing
        transcription_result: Result from transcription processing
        
    Returns:
        Merged result with speaker-labeled transcript segments
    """
    try:
        merged_result = {
            "audio_duration": 0.0,
            "speaker_count": 0,
            "segments": [],
            "processing_time": 0.0,
            "confidence_score": 0.0
        }
        
        if diarization_result:
            merged_result.update({
                "audio_duration": diarization_result.audio_duration,
                "speaker_count": diarization_result.speaker_count,
                "diarization_confidence": diarization_result.confidence_score
            })
        
        if transcription_result:
            merged_result.update({
                "transcription_confidence": transcription_result.accuracy_estimate,
                "word_count": transcription_result.word_count
            })
        
        # Merge segments if both results are available
        if diarization_result and transcription_result:
            merged_segments = align_segments(
                diarization_result.segments,
                transcription_result.segments
            )
            merged_result["segments"] = merged_segments
            merged_result["confidence_score"] = (
                diarization_result.confidence_score + transcription_result.accuracy_estimate
            ) / 2
        
        return merged_result
        
    except Exception as e:
        logger.error(f"Failed to merge results: {e}")
        raise ResultMergingError(f"Failed to merge results: {e}")


def execute_parallel_tasks(
    audio_path: Path,
    diarization_manager,
    transcription_service,
    diarization_config: Optional[DiarizationConfig] = None,
    transcription_config: Optional[Any] = None,
    timeout_seconds: int = 300
) -> tuple[Optional[DiarizationResult], Optional[Any]]:
    """Execute diarization and transcription tasks in parallel."""
    try:
        # Submit both tasks to thread pool
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            # Submit diarization task
            diarization_future = executor.submit(
                process_diarization_task,
                audio_path,
                diarization_manager,
                diarization_config
            )
            
            # Submit transcription task
            transcription_future = executor.submit(
                process_transcription_task,
                audio_path,
                transcription_service,
                transcription_config
            )
            
            # Wait for both tasks to complete
            try:
                diarization_result = diarization_future.result(timeout=timeout_seconds)
                transcription_result = transcription_future.result(timeout=timeout_seconds)
                
                return diarization_result, transcription_result
                
            except concurrent.futures.TimeoutError:
                raise ParallelProcessingError("Processing timeout exceeded")
            except Exception as e:
                raise ParallelProcessingError(f"Processing failed: {e}")
                
    except Exception as e:
        logger.error(f"Parallel task execution failed: {e}")
        raise


def process_batch_files(
    audio_paths: List[Path],
    parallel_processor,
    configs: Optional[Dict[str, Any]] = None,
    timeout_seconds: int = 300
) -> List[ProcessingResult]:
    """Process multiple files in batch using thread pool."""
    if not audio_paths:
        return []
    
    logger.info(f"Batch processing {len(audio_paths)} files")
    
    results = []
    
    # Process files in batches to control memory usage
    batch_size = min(len(audio_paths), parallel_processor.config.max_workers)
    
    for i in range(0, len(audio_paths), batch_size):
        batch = audio_paths[i:i + batch_size]
        
        # Submit batch to thread pool
        futures = []
        for audio_path in batch:
            future = parallel_processor.executor.submit(
                parallel_processor.process_file,
                audio_path,
                configs.get('diarization') if configs else None,
                configs.get('transcription') if configs else None
            )
            futures.append(future)
        
        # Collect results
        for future in concurrent.futures.as_completed(futures):
            try:
                result = future.result(timeout=timeout_seconds)
                results.append(result)
            except Exception as e:
                logger.error(f"Batch processing failed: {e}")
                # Create error result
                error_result = ProcessingResult(
                    task_id=f"error_{int(time.time() * 1000)}",
                    success=False,
                    error_message=str(e)
                )
                results.append(error_result)
    
    logger.info(f"Batch processing completed: {len(results)} results")
    return results


def update_processing_statistics(
    stats: Dict[str, Any],
    processing_time: float,
    success: bool,
    lock: Any
):
    """Update processing statistics thread-safely."""
    with lock:
        stats["total_files_processed"] += 1
        stats["total_processing_time"] += processing_time
        
        if success:
            stats["successful_processing"] += 1
        else:
            stats["failed_processing"] += 1
        
        # Calculate averages
        if stats["total_files_processed"] > 0:
            stats["average_processing_time"] = (
                stats["total_processing_time"] / stats["total_files_processed"]
            )


def calculate_success_rate(stats: Dict[str, Any]) -> float:
    """Calculate success rate from statistics."""
    if stats["total_files_processed"] > 0:
        return stats["successful_processing"] / stats["total_files_processed"]
    return 0.0


def estimate_processing_speedup(
    sequential_time: float,
    parallel_time: float,
    stats: Dict[str, Any],
    lock: Any
) -> float:
    """Estimate the speedup achieved by parallel processing."""
    if parallel_time <= 0:
        return 0.0
    
    speedup = sequential_time / parallel_time
    
    with lock:
        stats["parallel_speedup"] = speedup
    
    return speedup
