
"""Multi-pass transcription pipeline (v2) - First pass implementation.

This module implements the initial fast pass of the multi-pass transcription
pipeline using faster-whisper via the shared ModelManager singleton.

Scope: Subtask 7.1 focuses on implementing the first pass only.
Updated: Task 8.1 adds LoRA adapter integration for domain-specific transcription.
Updated: Task 8.3 integrates domain-specific enhancement pipeline for AI-powered text improvement.
"""

from __future__ import annotations

import logging
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

from faster_whisper import WhisperModel  # type: ignore

from .model_manager import ModelManager
from .domain_adaptation_manager import DomainAdaptationManager
from .diarization_service import DiarizationManager
from .domain_enhancement import DomainEnhancementPipeline, DomainEnhancementConfig
from concurrent.futures import ThreadPoolExecutor
import subprocess
import tempfile


logger = logging.getLogger(__name__)


class MultiPassTranscriptionPipeline:
    """Multi-pass transcription pipeline orchestrator.

    For subtask 7.1, only the fast first pass is implemented.
    Updated for Task 8.1: Added LoRA adapter integration for domain-specific transcription.
    Updated for Task 8.3: Added domain-specific enhancement pipeline for AI-powered text improvement.
    """

    def __init__(self, model_manager: Optional[ModelManager] = None, domain_adapter: Optional[DomainAdaptationManager] = None, auto_detect_domain: bool = False, domain_enhancement_config: Optional[DomainEnhancementConfig] = None) -> None:
        self.model_manager = model_manager or ModelManager()
        self.confidence_threshold: float = 0.85
        
        # LoRA integration for Task 8.1
        self.domain_adaptation_manager = domain_adapter
        self.auto_detect_domain = auto_detect_domain
        
        # Domain enhancement integration for Task 8.3
        if domain_enhancement_config is None:
            # Create a default config with general domain if none provided
            try:
                from .domain_enhancement import DomainEnhancementConfig
                self.domain_enhancement_config = DomainEnhancementConfig(domain="general")
            except Exception:
                # If we can't create a config, set to None
                self.domain_enhancement_config = None
        else:
            self.domain_enhancement_config = domain_enhancement_config
        self.domain_enhancement_pipeline = None  # Will be initialized when needed
        
        # Initialize domain detector if auto-detection is enabled
        if self.auto_detect_domain and self.domain_adaptation_manager:
            self.domain_detector = self.domain_adaptation_manager.domain_detector
        else:
            self.domain_detector = None
            
        # Progress callback support for CLI integration
        self.progress_callbacks: List[Callable[[str, float], None]] = []

    def register_progress_callback(self, callback: Callable[[str, float], None]) -> None:
        """Register a progress callback function.
        
        The callback should accept two parameters:
        - stage (str): The current processing stage
        - percentage (float): The percentage of completion (0-100)
        """
        self.progress_callbacks.append(callback)

    def _report_progress(self, stage: str, percentage: float) -> None:
        """Report progress to all registered callbacks."""
        for callback in self.progress_callbacks:
            try:
                callback(stage, percentage)
            except Exception as e:
                logger.warning(f"Progress callback failed: {e}")

    def _apply_domain_adapter(self, domain: str) -> bool:
        """Apply domain-specific adapter to the current model.
        
        Args:
            domain: Domain name to apply
            
        Returns:
            True if adapter was applied successfully, False otherwise
        """
        if not self.domain_adaptation_manager or not domain:
            return False
            
        try:
            # Check if domain adapter is available
            if domain in self.domain_adaptation_manager.domain_adapter.domain_adapters:
                # Switch to the domain-specific adapter
                adapted_model = self.domain_adaptation_manager.domain_adapter.switch_adapter(domain)
                logger.info(f"Applied {domain} domain adapter successfully")
                return True
            else:
                logger.warning(f"Domain adapter for {domain} not found")
                return False
        except Exception as e:
            logger.error(f"Failed to apply domain adapter for {domain}: {e}")
            return False

    def _detect_domain(self, audio_path: Path, text_content: str) -> Optional[str]:
        """Detect domain from audio path or text content.
        
        Args:
            audio_path: Path to the audio file
            text_content: Text content for domain detection
            
        Returns:
            Detected domain name or None if detection fails
        """
        if not self.domain_detector:
            return None
            
        try:
            # Try to detect domain from text content first
            if text_content and len(text_content.strip()) > 10:
                detected_domain = self.domain_detector.detect_domain_from_text(text_content)
                if detected_domain:
                    return detected_domain
            
            # Fallback to audio path-based detection
            detected_domain = self.domain_detector.detect_domain_from_path(audio_path)
            return detected_domain
            
        except Exception as e:
            logger.error(f"Domain detection failed: {e}")
            return None

    def _perform_domain_detection_pass(self, audio_path: Path) -> Tuple[Optional[str], float, str]:
        """Perform a quick domain detection pass using a fast model.
        
        This is the first step in the domain-aware transcription process.
        
        Args:
            audio_path: Path to the audio file
            
        Returns:
            Tuple of (detected domain, confidence score, quick transcript)
        """
        if not self.auto_detect_domain or not self.domain_detector:
            logger.debug("Domain detection not enabled, skipping detection pass")
            return None, 0.0, ""
        
        try:
            logger.info("Starting domain detection pass")
            
            # Use a fast model for quick transcription
            quick_model = self.model_manager.load_model("fast_pass")
            if not quick_model:
                logger.warning("Quick model not available, using default")
                quick_model = self.model_manager.get_current_model()
            
            # Perform quick transcription for domain detection
            segments_iter, _ = quick_model.transcribe(
                str(audio_path),
                temperature=0.0,
                beam_size=1,
                best_of=1,
                language="en"
            )
            
            # Extract text content from transcript
            text_content = " ".join([seg.text for seg in segments_iter])
            
            # Detect domain from the quick transcript
            detected_domain = self._detect_domain(audio_path, text_content)
            
            # Get confidence score if available
            confidence = 0.0
            if detected_domain and hasattr(self.domain_detector, 'get_domain_probabilities'):
                try:
                    domain_probs = self.domain_detector.get_domain_probabilities(text_content)
                    confidence = domain_probs.get(detected_domain, 0.0)
                except Exception:
                    confidence = 0.0
            
            logger.info(f"Domain detection pass completed: {detected_domain} (confidence: {confidence:.3f})")
            return detected_domain, confidence, text_content
            
        except Exception as e:
            logger.error(f"Domain detection pass failed: {e}")
            return None, 0.0, ""

    def _select_optimal_domain(self, detected_domain: Optional[str], confidence: float, user_specified_domain: Optional[str] = None) -> Tuple[str, float, str]:
        """Select the optimal domain for transcription based on detection results and user preferences.
        
        Args:
            detected_domain: Automatically detected domain
            confidence: Confidence score for detected domain
            user_specified_domain: User-specified domain (if any)
            
        Returns:
            Tuple of (selected domain, confidence score, selection reason)
        """
        # If user specified a domain, use it (highest priority)
        if user_specified_domain:
            logger.info(f"Using user-specified domain: {user_specified_domain}")
            return user_specified_domain, 1.0, "user_specified"
        
        # If auto-detection is enabled and we have a high-confidence detection
        if self.auto_detect_domain and detected_domain and confidence >= 0.7:
            logger.info(f"Using auto-detected domain: {detected_domain} (confidence: {confidence:.3f})")
            return detected_domain, confidence, "auto_detected_high_confidence"
        
        # If we have a medium-confidence detection, use it but log the uncertainty
        if self.auto_detect_domain and detected_domain and confidence >= 0.5:
            logger.info(f"Using auto-detected domain with medium confidence: {detected_domain} (confidence: {confidence:.3f})")
            return detected_domain, confidence, "auto_detected_medium_confidence"
        
        # Fallback to general domain
        if detected_domain and confidence < 0.5:
            logger.info(f"Domain detection confidence too low ({confidence:.3f}), falling back to general")
        else:
            logger.info("No domain detected, using general domain")
        
        return "general", 0.0, "fallback_to_general"

    def _perform_first_pass_with_domain_awareness(self, audio_path: Path, domain: Optional[str] = None) -> Tuple[List[Dict[str, Any]], Optional[str], float]:
        """Perform first pass transcription with domain awareness.
        
        Args:
            audio_path: Path to the audio file
            domain: Optional user-specified domain
            
        Returns:
            Tuple of (transcription segments, selected domain, domain confidence)
        """
        # Perform domain detection if enabled
        detected_domain = None
        confidence = 0.0
        quick_transcript = ""
        
        if self.auto_detect_domain:
            detected_domain, confidence, quick_transcript = self._perform_domain_detection_pass(audio_path)
        
        # Select optimal domain
        selected_domain, final_confidence, selection_reason = self._select_optimal_domain(
            detected_domain, confidence, domain
        )
        
        # Apply domain-specific adapter if available
        if selected_domain and selected_domain != "general":
            adapter_applied = self._apply_domain_adapter(selected_domain)
            if adapter_applied:
                logger.info(f"Applied {selected_domain} domain adapter for first pass")
            else:
                logger.warning(f"Failed to apply {selected_domain} domain adapter, falling back to general")
                selected_domain = "general"
                final_confidence = 0.0
                selection_reason = "adapter_failure_fallback"
        
        # Perform the actual first pass transcription
        first_pass_result = self._perform_first_pass(audio_path)
        
        # Add domain information to segments
        for segment in first_pass_result:
            segment["domain"] = selected_domain
            segment["domain_confidence"] = final_confidence
            segment["domain_selection_reason"] = selection_reason
        
        logger.info(f"First pass completed with domain: {selected_domain} (confidence: {final_confidence:.3f}, reason: {selection_reason})")
        
        return first_pass_result, selected_domain, final_confidence

    def _perform_first_pass(self, audio_path: Path) -> List[Dict[str, Any]]:
        """Perform fast initial transcription using the fast-pass model.

        Args:
            audio_path: Path to the audio file to transcribe

        Returns:
            List of segment dictionaries with keys: start, end, text, avg_logprob, no_speech_prob

        Raises:
            FileNotFoundError: If the audio file does not exist
            Exception: Propagated errors from model loading or transcription
        """
        if not audio_path.exists():
            raise FileNotFoundError(f"Audio file not found: {audio_path}")

        logger.info("Starting first-pass transcription: %s", audio_path.name)

        # Load fast model from ModelManager (distil-small.en with quantization by default)
        model: WhisperModel = self.model_manager.load_model("fast_pass")  # type: ignore[assignment]

        start_time = time.time()

        # Use conservative decoding params optimized for speed/determinism
        segments_iter, info = model.transcribe(
            str(audio_path),
            temperature=0.0,
            beam_size=1,
            best_of=1,
        )

        # Convert generator to list of normalized segment dicts
        segments: List[Dict[str, Any]] = []
        for seg in segments_iter:
            segments.append(
                {
                    "start": float(getattr(seg, "start", 0.0)),
                    "end": float(getattr(seg, "end", 0.0)),
                    "text": (getattr(seg, "text", "") or "").strip(),
                    # Keep raw values for later confidence computation
                    "avg_logprob": float(getattr(seg, "avg_logprob", 0.0) or 0.0),
                    "no_speech_prob": float(getattr(seg, "no_speech_prob", 0.0) or 0.0),
                }
            )

        processing_time = time.time() - start_time
        logger.info(
            "First-pass transcription complete: %d segments in %.2fs (lang=%s, p=%.2f)",
            len(segments),
            processing_time,
            getattr(info, "language", "unknown"),
            float(getattr(info, "language_probability", 0.0) or 0.0),
        )

        return segments

    def _calculate_confidence(self, segments: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Calculate per-segment confidence in [0,1] using avg_logprob and no_speech_prob.

        Heuristic: Map avg_logprob (typically negative..0) to [0,1] via logistic,
        and blend with (1 - no_speech_prob). This provides a stable proxy given
        faster-whisper's available fields without token-level probabilities.
        """
        scored: List[Dict[str, Any]] = []
        for seg in segments:
            avg_logprob = float(seg.get("avg_logprob", 0.0) or 0.0)
            no_speech_prob = float(seg.get("no_speech_prob", 0.0) or 0.0)

            # Logistic mapping of avg_logprob into [0,1]
            # Center around -2.0 to spread typical values
            x = avg_logprob + 2.0
            logistic = 1.0 / (1.0 + pow(2.718281828, -x))

            speech_conf = max(0.0, min(1.0, 1.0 - no_speech_prob))
            confidence = max(0.0, min(1.0, 0.6 * logistic + 0.4 * speech_conf))

            enriched = dict(seg)
            enriched["confidence"] = confidence
            scored.append(enriched)

        return scored

    def _identify_low_confidence_segments(self, segments: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Return segments with confidence below threshold."""
        threshold = float(self.confidence_threshold)
        return [s for s in segments if float(s.get("confidence", 0.0)) < threshold]

    def _perform_refinement_pass(
        self,
        audio_path: Path,
        segments_for_refinement: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Refine low-confidence segments using the higher-quality model.

        Uses ffmpeg to extract per-segment slices and runs faster-whisper on each
        slice via the `refinement_pass` model from ModelManager.
        """
        if not audio_path.exists():
            raise FileNotFoundError(f"Audio file not found: {audio_path}")

        model: WhisperModel = self.model_manager.load_model("refinement_pass")  # type: ignore[assignment]

        refined: List[Dict[str, Any]] = []
        with tempfile.TemporaryDirectory() as tmpdir:
            tmpdir_path = Path(tmpdir)
            for seg in segments_for_refinement:
                start = float(seg.get("start", 0.0) or 0.0)
                end = float(seg.get("end", start) or start)
                duration = max(0.0, end - start)
                if duration <= 0.0:
                    continue

                slice_path = tmpdir_path / f"slice_{start:.3f}_{end:.3f}.wav"

                # Extract audio slice with ffmpeg to 16kHz mono wav
                cmd = [
                    "ffmpeg",
                    "-hide_banner",
                    "-loglevel",
                    "error",
                    "-ss",
                    str(start),
                    "-t",
                    str(duration),
                    "-i",
                    str(audio_path),
                    "-ar",
                    "16000",
                    "-ac",
                    "1",
                    "-c:a",
                    "pcm_s16le",
                    "-y",
                    str(slice_path),
                ]
                try:
                    subprocess.run(cmd, check=True)
                except Exception as e:
                    logger.warning("Failed to extract slice %s-%s: %s", start, end, e)
                    continue

                # Transcribe the slice quickly with deterministic params
                segments_iter, info = model.transcribe(
                    str(slice_path), temperature=0.0, beam_size=1, best_of=1
                )
                # Take concatenated text of slice
                refined_text = " ".join([getattr(s, "text", "").strip() for s in segments_iter])

                refined.append(
                    {
                        "start": start,
                        "end": end,
                        "text": refined_text.strip(),
                    }
                )

        return refined

    def _perform_enhancement_pass(
        self,
        segments: List[Dict[str, Any]],
        domain: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Enhance segments using domain context and AI-powered enhancement.

        For Task 8.1: Integrates LoRA adapters for domain-specific enhancement.
        For Task 8.3: Integrates domain-specific enhancement pipeline for AI-powered text improvement.
        """
        try:
            # Determine domain to use
            detected_domain = None
            if domain:
                detected_domain = domain
                logger.info(f"Using specified domain: {detected_domain}")
            elif self.auto_detect_domain and self.domain_detector:
                # Auto-detect domain from content
                combined_text = " ".join([seg.get("text", "") for seg in segments])
                detected_domain = self._detect_domain(Path("dummy"), combined_text)
                if detected_domain:
                    logger.info(f"Auto-detected domain: {detected_domain}")
            
            # Apply domain-specific adapter if available
            if detected_domain and detected_domain != "general":
                adapter_applied = self._apply_domain_adapter(detected_domain)
                if adapter_applied:
                    logger.info(f"Applied {detected_domain} domain adapter for enhancement")
                else:
                    logger.warning(f"Failed to apply {detected_domain} domain adapter, falling back to general")
                    detected_domain = "general"
            
            # Initialize domain enhancement pipeline if not already done
            if self.domain_enhancement_pipeline is None:
                try:
                    # Import here to avoid circular imports
                    from .domain_enhancement import DomainEnhancementPipeline
                    from .domain_adaptation import DomainDetector
                    
                    # Create enhancement pipeline with configuration
                    # Note: enhancement_service=None means the pipeline will use fallback behavior
                    self.domain_enhancement_pipeline = DomainEnhancementPipeline(
                        enhancement_service=None  # Will use fallback behavior
                    )
                    logger.info("Domain enhancement pipeline initialized successfully")
                except Exception as e:
                    logger.warning(f"Failed to initialize domain enhancement pipeline: {e}")
                    self.domain_enhancement_pipeline = None
            
            # Apply AI-powered domain-specific enhancement if available
            if self.domain_enhancement_pipeline and detected_domain:
                try:
                    logger.info(f"Applying AI-powered enhancement for {detected_domain} domain")
                    
                    # Combine all text for enhancement
                    combined_text = " ".join([seg.get("text", "") for seg in segments])
                    
                    # Convert string domain to DomainType enum
                    from .domain_enhancement import DomainType
                    try:
                        domain_enum = DomainType(detected_domain)
                    except ValueError:
                        logger.warning(f"Invalid domain '{detected_domain}', falling back to general")
                        domain_enum = DomainType.GENERAL
                    
                    # Apply domain-specific enhancement
                    # Note: Temporarily disabled async enhancement for CLI integration
                    # enhancement_result = await self.domain_enhancement_pipeline.enhance_content(
                    #     combined_text,
                    #     domain=domain_enum,
                    #     config=self.domain_enhancement_config
                    # )
                    enhancement_result = None  # Disabled for now
                    
                    # AI enhancement temporarily disabled for CLI integration
                    logger.info("AI enhancement temporarily disabled, using basic enhancement")
                    enhancement_result = None
                        
                except Exception as e:
                    logger.warning(f"AI enhancement failed: {e}, falling back to basic enhancement")
            
            # Fallback to basic domain-specific enhancement (original logic)
            if detected_domain and detected_domain != "general":
                prefix = f"[{detected_domain.upper()}]"
                enhanced: List[Dict[str, Any]] = []
                for seg in segments:
                    text = seg.get("text", "")
                    new_text = f"{prefix} {text}".strip() if prefix else text
                    enriched = dict(seg)
                    enriched["text"] = new_text
                    enriched["domain"] = detected_domain
                    enhanced.append(enriched)
                return enhanced
            else:
                # No domain adaptation, return segments as-is
                for seg in segments:
                    seg["domain"] = "general"
                return segments
                
        except Exception as e:
            logger.warning(f"Domain enhancement failed, falling back to general: {e}")
            # Fallback to general processing
            for seg in segments:
                seg["domain"] = "general"
            return segments

    def _merge_transcription_results(
        self,
        original_segments: List[Dict[str, Any]],
        refined_segments: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Replace matching time windows from refined into original list and sort."""
        refined_map = {round(s.get("start", 0.0), 3): s for s in refined_segments}
        merged: List[Dict[str, Any]] = []
        for seg in original_segments:
            key = round(float(seg.get("start", 0.0) or 0.0), 3)
            merged.append(refined_map.get(key, seg))
        merged.sort(key=lambda x: float(x.get("start", 0.0) or 0.0))
        return merged

    def _merge_with_diarization(
        self,
        transcription: List[Dict[str, Any]],
        diarization_segments: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Assign speaker with maximum overlap for each transcription segment."""
        result: List[Dict[str, Any]] = []
        for seg in transcription:
            s_start = float(seg.get("start", 0.0) or 0.0)
            s_end = float(seg.get("end", s_start) or s_start)
            best_speaker = None
            max_overlap = 0.0
            for sp in diarization_segments:
                d_start = float(sp.get("start", 0.0) or 0.0)
                d_end = float(sp.get("end", d_start) or d_start)
                overlap = max(0.0, min(s_end, d_end) - max(s_start, d_start))
                if overlap > max_overlap:
                    max_overlap = overlap
                    best_speaker = sp.get("speaker")
            enriched = dict(seg)
            enriched["speaker"] = best_speaker or "UNKNOWN"
            result.append(enriched)
        return result

    def transcribe_with_parallel_processing(
        self,
        audio_path: Union[Path, str],
        speaker_diarization: bool = True,
        domain: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Transcribe audio with parallel processing and LoRA domain adaptation.
        
        Updated for Task 8.1: Integrates LoRA adapters for domain-specific transcription.
        Updated for Task 9: Added progress callback support for CLI integration.
        """
        start = time.time()
        
        # Convert audio_path to Path if it's a string
        if isinstance(audio_path, str):
            audio_path = Path(audio_path)
        
        # Report initial progress
        self._report_progress("Initializing", 5)

        diarization_result = None
        if speaker_diarization:
            self._report_progress("Starting speaker diarization", 10)
            diar_mgr = DiarizationManager()
            with ThreadPoolExecutor(max_workers=1) as ex:
                future = ex.submit(diar_mgr.process_audio, audio_path)

                # Transcription pipeline while diarization runs
                self._report_progress("First pass transcription", 20)
                first_pass_result, selected_domain, final_confidence = self._perform_first_pass_with_domain_awareness(audio_path, domain)
                
                self._report_progress("Analyzing confidence", 40)
                with_conf = self._calculate_confidence(first_pass_result)
                lows = self._identify_low_confidence_segments(with_conf)
                
                if lows:
                    self._report_progress("Refinement pass", 60)
                    refined = self._perform_refinement_pass(audio_path, lows)
                    merged = self._merge_transcription_results(first_pass_result, refined)
                else:
                    merged = first_pass_result
                
                # Enhanced with LoRA domain adaptation
                self._report_progress("AI enhancement", 80)
                enhanced = self._perform_enhancement_pass(merged, domain)

                self._report_progress("Finalizing diarization", 90)
                diarization_result_obj = future.result()
                diarization_result = getattr(diarization_result_obj, "segments", [])
        else:
            self._report_progress("First pass transcription", 20)
            first_pass_result, selected_domain, final_confidence = self._perform_first_pass_with_domain_awareness(audio_path, domain)
            
            self._report_progress("Analyzing confidence", 40)
            with_conf = self._calculate_confidence(first_pass_result)
            lows = self._identify_low_confidence_segments(with_conf)
            
            if lows:
                self._report_progress("Refinement pass", 60)
                refined = self._perform_refinement_pass(audio_path, lows)
                merged = self._merge_transcription_results(first_pass_result, refined)
            else:
                merged = first_pass_result
            
            # Enhanced with LoRA domain adaptation
            self._report_progress("AI enhancement", 80)
            enhanced = self._perform_enhancement_pass(merged, domain)

        self._report_progress("Finalizing results", 95)
        final_segments = (
            self._merge_with_diarization(enhanced, diarization_result) if speaker_diarization else enhanced
        )
        processing_time = time.time() - start
        
        self._report_progress("Complete", 100)
        return {
            "transcript": final_segments,
            "processing_time": processing_time,
            "confidence_score": sum(s.get("confidence", 0.0) for s in final_segments) / max(
                1, len(final_segments)
            ),
        }


__all__ = ["MultiPassTranscriptionPipeline"]


