"""
Speech recognition module for FMUS-VID.

This module provides classes for transcribing speech in videos using
Whisper and other speech recognition models.
"""

import os
import time
import logging
import tempfile
import asyncio
import threading
import queue
from typing import List, Dict, Optional, Union, Generator, AsyncGenerator, Tuple, Any
import numpy as np

logger = logging.getLogger(__name__)

# Try to import faster-whisper, but don't fail if it's not available
try:
    from faster_whisper import WhisperModel
    FASTER_WHISPER_AVAILABLE = True
except ImportError:
    logger.warning("faster-whisper not installed. Speech recognition features will be limited.")
    FASTER_WHISPER_AVAILABLE = False


class WhisperTranscriber:
    """
    Transcribe audio using Whisper models.

    This class provides both batch and streaming transcription of audio.
    """

    def __init__(
        self,
        model: str = "medium",
        device: str = "cuda",
        compute_type: str = "float16",
        download_root: Optional[str] = None,
        vad_filter: bool = True,
        buffer_size: int = 5,
        streaming_buffer_size: float = 5.0,
        chunk_duration: float = 1.0,
        sample_rate: int = 16000,
        language: Optional[str] = None,
        translate: bool = False
    ):
        """
        Initialize the WhisperTranscriber.

        Args:
            model: Whisper model size ('tiny', 'base', 'small', 'medium', 'large')
            device: Device to run on ('cuda', 'cpu', 'cuda:0', etc.)
            compute_type: Compute type ('float16', 'float32', 'int8')
            download_root: Directory to download models to (None for default)
            vad_filter: Whether to use voice activity detection
            buffer_size: Number of audio chunks to buffer
            streaming_buffer_size: Size of streaming buffer in seconds
            chunk_duration: Duration of each audio chunk in seconds
            sample_rate: Audio sample rate (16000 for Whisper)
            language: Language code (None for auto-detection)
            translate: Whether to translate to English
        """
        self.model_name = model
        self.device = device
        self.compute_type = compute_type
        self.download_root = download_root
        self.vad_filter = vad_filter
        self.buffer_size = buffer_size
        self.streaming_buffer_size = streaming_buffer_size
        self.chunk_duration = chunk_duration
        self.sample_rate = sample_rate
        self.language = language
        self.translate = translate

        # Internal state
        self._model = None
        self._audio_buffer = []
        self._buffer_lock = threading.Lock()
        self._audio_queue = queue.Queue()
        self._transcript_queue = queue.Queue()
        self._processing_thread = None
        self._running = False
        self._last_timestamp = 0.0
        self._temp_dir = None
        self._temp_audio_file = None

        if not FASTER_WHISPER_AVAILABLE:
            raise ImportError("faster-whisper is required for WhisperTranscriber")

        # Initialize the model
        self._initialize_model()

    def _initialize_model(self):
        """Initialize the Whisper model."""
        try:
            logger.info(f"Loading Whisper model '{self.model_name}' on {self.device}...")

            self._model = WhisperModel(
                model_size_or_path=self.model_name,
                device=self.device,
                compute_type=self.compute_type,
                download_root=self.download_root
            )

            logger.info(f"Whisper model loaded successfully")

            # Create temporary directory for audio files
            self._temp_dir = tempfile.TemporaryDirectory()
            self._temp_audio_file = os.path.join(self._temp_dir.name, "temp_audio.wav")

        except Exception as e:
            logger.error(f"Failed to initialize Whisper model: {e}")
            raise

    def transcribe(
        self,
        audio: Union[str, np.ndarray],
        language: Optional[str] = None,
        translate: Optional[bool] = None,
        return_timestamps: bool = False,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Transcribe audio in batch mode.

        Args:
            audio: Audio file path or numpy array
            language: Language code (None for auto-detection)
            translate: Whether to translate to English
            return_timestamps: Whether to return word-level timestamps
            **kwargs: Additional arguments to pass to WhisperModel.transcribe

        Returns:
            Dict containing the transcription results
        """
        if self._model is None:
            raise RuntimeError("Whisper model not initialized")

        # Set up parameters
        language = language or self.language
        translate = translate if translate is not None else self.translate

        try:
            start_time = time.time()

            # Handle different audio input types
            if isinstance(audio, str):
                # File path
                audio_path = audio
            elif isinstance(audio, np.ndarray):
                # Numpy array
                import soundfile as sf
                sf.write(self._temp_audio_file, audio, self.sample_rate)
                audio_path = self._temp_audio_file
            else:
                raise ValueError("Audio must be a file path or numpy array")

            # Transcribe the audio
            segments, info = self._model.transcribe(
                audio=audio_path,
                language=language,
                task="translate" if translate else "transcribe",
                vad_filter=self.vad_filter,
                word_timestamps=return_timestamps,
                **kwargs
            )

            # Collect segments
            result = {
                "text": "",
                "segments": [],
                "language": info.language,
                "language_probability": info.language_probability
            }

            for segment in segments:
                result["text"] += segment.text + " "
                segment_dict = {
                    "id": segment.id,
                    "start": segment.start,
                    "end": segment.end,
                    "text": segment.text
                }

                if return_timestamps and hasattr(segment, "words"):
                    segment_dict["words"] = [
                        {"text": word.word, "start": word.start, "end": word.end}
                        for word in segment.words
                    ]

                result["segments"].append(segment_dict)

            result["text"] = result["text"].strip()

            elapsed = time.time() - start_time
            logger.info(f"Transcription completed in {elapsed:.2f}s")

            return result

        except Exception as e:
            logger.error(f"Transcription error: {e}")
            raise

    def start_streaming(self):
        """Start streaming transcription mode."""
        if self._running:
            return

        self._running = True

        # Clear buffers
        with self._buffer_lock:
            self._audio_buffer = []

        while not self._audio_queue.empty():
            self._audio_queue.get()

        while not self._transcript_queue.empty():
            self._transcript_queue.get()

        # Start processing thread
        self._processing_thread = threading.Thread(target=self._process_audio_queue)
        self._processing_thread.daemon = True
        self._processing_thread.start()

        logger.info("Started streaming transcription")

    def stop_streaming(self):
        """Stop streaming transcription mode."""
        if not self._running:
            return

        self._running = False

        # Wait for processing thread to finish
        if self._processing_thread:
            self._processing_thread.join(timeout=2.0)
            self._processing_thread = None

        logger.info("Stopped streaming transcription")

    def add_audio_chunk(self, audio: np.ndarray):
        """
        Add an audio chunk to the streaming buffer.

        Args:
            audio: Audio samples as numpy array
        """
        if not self._running:
            self.start_streaming()

        # Add to buffer
        with self._buffer_lock:
            self._audio_buffer.append(audio)

            # Keep buffer at desired size
            buffer_duration = 0.0
            for chunk in self._audio_buffer:
                buffer_duration += len(chunk) / self.sample_rate

            # Trim buffer if needed
            while buffer_duration > self.streaming_buffer_size and len(self._audio_buffer) > 1:
                removed = self._audio_buffer.pop(0)
                buffer_duration -= len(removed) / self.sample_rate

        # Add to processing queue
        self._audio_queue.put(audio)

    def _process_audio_queue(self):
        """Background thread for processing audio chunks."""
        try:
            logger.info("Audio processing thread started")

            while self._running:
                # Get audio chunk from queue
                try:
                    audio = self._audio_queue.get(timeout=0.1)
                except queue.Empty:
                    continue

                # Skip empty chunks
                if len(audio) == 0:
                    self._audio_queue.task_done()
                    continue

                # Process audio chunk
                try:
                    # Get current buffer for context
                    with self._buffer_lock:
                        # Concatenate recent chunks for context
                        buffer_audio = np.concatenate(self._audio_buffer[-self.buffer_size:])

                    # Write to temporary file
                    import soundfile as sf
                    sf.write(self._temp_audio_file, buffer_audio, self.sample_rate)

                    # Transcribe
                    result = self.transcribe(
                        self._temp_audio_file,
                        return_timestamps=True
                    )

                    # Filter segments by recency
                    buffer_duration = len(buffer_audio) / self.sample_rate
                    recent_segments = []

                    for segment in result["segments"]:
                        # Only include segments from the recent portion
                        if segment["end"] >= buffer_duration - self.chunk_duration * 1.5:
                            # Adjust timestamps relative to current time
                            segment["start"] = max(0, segment["start"] - (buffer_duration - self.chunk_duration))
                            segment["end"] = max(0, segment["end"] - (buffer_duration - self.chunk_duration))
                            recent_segments.append(segment)

                    # Add to transcript queue if we have new segments
                    if recent_segments:
                        self._transcript_queue.put(recent_segments)

                except Exception as e:
                    logger.error(f"Error processing audio chunk: {e}")

                finally:
                    self._audio_queue.task_done()

        except Exception as e:
            logger.error(f"Audio processing thread error: {e}")

        finally:
            logger.info("Audio processing thread stopped")

    def get_transcript(self, timeout: float = 0.0) -> Optional[List[Dict]]:
        """
        Get the latest transcript segments.

        Args:
            timeout: Timeout in seconds (0 for non-blocking)

        Returns:
            List of transcript segments or None if no new transcripts
        """
        try:
            return self._transcript_queue.get(timeout=timeout)
        except queue.Empty:
            return None

    def transcribe_stream(self, audio_chunks: Generator[np.ndarray, None, None]) -> Generator[str, None, None]:
        """
        Transcribe an audio stream.

        Args:
            audio_chunks: Generator yielding audio chunks

        Yields:
            Transcribed text segments
        """
        self.start_streaming()

        try:
            for chunk in audio_chunks:
                self.add_audio_chunk(chunk)

                # Check for new transcripts
                transcript = self.get_transcript(timeout=0.0)
                if transcript:
                    for segment in transcript:
                        yield segment["text"]

                # Small sleep to avoid busy waiting
                time.sleep(0.01)

        finally:
            self.stop_streaming()

    async def transcribe_chunk(self, audio: np.ndarray) -> Optional[str]:
        """
        Transcribe a single audio chunk asynchronously.

        This is optimized for real-time processing of small chunks.

        Args:
            audio: Audio samples as numpy array

        Returns:
            Transcribed text or None if no speech detected
        """
        if not self._running:
            self.start_streaming()

        self.add_audio_chunk(audio)

        # Use asyncio to wait for transcript without blocking
        for _ in range(10):  # Try for a short time
            transcript = self.get_transcript(timeout=0.0)
            if transcript:
                # Return the text of all segments concatenated
                return " ".join(segment["text"] for segment in transcript)

            # Wait a bit and try again
            await asyncio.sleep(0.05)

        return None

    async def aiter_transcribe(self, audio_chunks: AsyncGenerator[np.ndarray, None]) -> AsyncGenerator[str, None]:
        """
        Asynchronously transcribe an audio stream.

        Args:
            audio_chunks: Async generator yielding audio chunks

        Yields:
            Transcribed text segments
        """
        self.start_streaming()

        try:
            async for chunk in audio_chunks:
                self.add_audio_chunk(chunk)

                # Check for new transcripts
                transcript = self.get_transcript(timeout=0.0)
                if transcript:
                    for segment in transcript:
                        yield segment["text"]

                # Small sleep to avoid busy waiting
                await asyncio.sleep(0.01)

        finally:
            self.stop_streaming()

    def set_model_size(self, model_size: str):
        """
        Change the model size on the fly.

        Args:
            model_size: New model size ('tiny', 'base', 'small', 'medium', 'large')
        """
        if model_size == self.model_name:
            return

        logger.info(f"Changing model size from {self.model_name} to {model_size}")

        # Stop streaming if running
        was_running = self._running
        if was_running:
            self.stop_streaming()

        # Update model
        self.model_name = model_size
        self._initialize_model()

        # Restart streaming if it was running
        if was_running:
            self.start_streaming()

    def __del__(self):
        """Clean up resources when the object is garbage collected."""
        self.stop_streaming()

        # Clean up temporary directory
        if self._temp_dir:
            self._temp_dir.cleanup()
            self._temp_dir = None
