"""
Whisper-based speech recognition implementations.

This module provides classes for transcribing speech using OpenAI's Whisper models,
including both the standard Whisper implementation and the faster-whisper variant.
"""

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

from fmusvid.ai.speech import BaseSpeechRecognizer

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. Using standard Whisper implementation instead.")
    FASTER_WHISPER_AVAILABLE = False

# Try to import standard whisper, but don't fail if it's not available
try:
    import whisper
    WHISPER_AVAILABLE = True
except ImportError:
    logger.warning("whisper not installed. Some speech recognition features may be limited.")
    WHISPER_AVAILABLE = False


class WhisperTranscriber(BaseSpeechRecognizer):
    """
    Transcribe audio using OpenAI's Whisper models.

    This class provides both batch and streaming transcription of audio using the
    standard Whisper implementation.
    """

    def __init__(
        self,
        model: str = "medium",
        device: str = "cuda",
        compute_type: str = "float16",
        download_root: Optional[str] = None,
        vad_filter: bool = False,  # Standard Whisper doesn't have VAD
        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 (not available in standard Whisper)
            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
        """
        if not WHISPER_AVAILABLE:
            raise ImportError("whisper is required for WhisperTranscriber")

        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

        # 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 = whisper.load_model(
                name=self.model_name,
                device=self.device,
                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,
        **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
            **kwargs: Additional arguments to pass to whisper.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
                import soundfile as sf
                audio_data, _ = sf.read(audio)
            elif isinstance(audio, np.ndarray):
                # Numpy array
                audio_data = audio
            else:
                raise ValueError("Audio must be a file path or numpy array")

            # Transcribe the audio
            options = {
                "language": language,
                "task": "translate" if translate else "transcribe",
            }
            options.update(kwargs)

            result = self._model.transcribe(audio_data, **options)

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

            return result

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

    def start_streaming(self) -> bool:
        """
        Start streaming mode for real-time transcription.

        Returns:
            True if streaming started successfully, False otherwise
        """
        if self._running:
            logger.warning("Streaming already running")
            return True

        try:
            # Clear queues
            while not self._audio_queue.empty():
                self._audio_queue.get_nowait()

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

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

            logger.info("Streaming started")
            return True

        except Exception as e:
            logger.error(f"Error starting streaming: {e}")
            self._running = False
            return False

    def stop_streaming(self) -> None:
        """Stop streaming mode."""
        if not self._running:
            return

        self._running = False

        if self._processing_thread:
            self._processing_thread.join(timeout=2.0)
            self._processing_thread = None

        logger.info("Streaming stopped")

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

        Args:
            audio: Audio chunk as numpy array
        """
        if not self._running:
            logger.warning("Streaming not running, starting now")
            self.start_streaming()

        try:
            self._audio_queue.put(audio)
        except Exception as e:
            logger.error(f"Error adding audio chunk to queue: {e}")

    def _process_audio_queue(self) -> None:
        """Process audio chunks from the queue in a background thread."""
        buffer = []
        buffer_duration = 0.0

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

                # Add to buffer
                buffer.append(audio_chunk)
                chunk_duration = len(audio_chunk) / self.sample_rate
                buffer_duration += chunk_duration

                # Process if buffer is full
                if buffer_duration >= self.streaming_buffer_size:
                    # Combine chunks
                    audio_data = np.concatenate(buffer)

                    # Transcribe
                    result = self.transcribe(
                        audio_data,
                        language=self.language,
                        translate=self.translate
                    )

                    # Add to transcript queue if there's text
                    if result and result.get("text", "").strip():
                        self._transcript_queue.put(result)

                    # Reset buffer
                    buffer = []
                    buffer_duration = 0.0

            except Exception as e:
                logger.error(f"Error processing audio chunk: {e}")
                time.sleep(0.1)  # Prevent CPU spinning on errors

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

        Args:
            timeout: Time to wait for a transcript (0 for non-blocking)

        Returns:
            Transcript dict, or None if no transcript is available
        """
        try:
            return self._transcript_queue.get(block=timeout > 0, timeout=timeout)
        except queue.Empty:
            return None

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

        Args:
            audio_chunks: Generator yielding audio chunks

        Yields:
            Transcribed text as it becomes available
        """
        # Start streaming if not already running
        if not self._running:
            self.start_streaming()

        try:
            # Process each chunk
            for chunk in audio_chunks:
                self.add_audio_chunk(chunk)

                # Check for transcripts
                transcript = self.get_transcript()
                if transcript:
                    yield transcript.get("text", "")

        finally:
            # Stop streaming when done
            self.stop_streaming()

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

        Args:
            audio: Audio chunk as numpy array

        Returns:
            Transcribed text, or None if no speech was detected
        """
        # Run transcription in a thread pool to avoid blocking
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            None,
            lambda: self.transcribe(
                audio,
                language=self.language,
                translate=self.translate
            )
        )

        # Extract text from result
        if result and result.get("text", "").strip():
            return result.get("text", "").strip()

        return None

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

        Args:
            audio_chunks: Async generator yielding audio chunks

        Yields:
            Transcribed text as it becomes available
        """
        buffer = []
        buffer_duration = 0.0

        async for chunk in audio_chunks:
            # Add to buffer
            buffer.append(chunk)
            chunk_duration = len(chunk) / self.sample_rate
            buffer_duration += chunk_duration

            # Process if buffer is full
            if buffer_duration >= self.streaming_buffer_size:
                # Combine chunks
                audio_data = np.concatenate(buffer)

                # Transcribe
                text = await self.transcribe_chunk(audio_data)

                # Yield text if not empty
                if text:
                    yield text

                # Reset buffer
                buffer = []
                buffer_duration = 0.0

        # Process any remaining audio
        if buffer:
            audio_data = np.concatenate(buffer)
            text = await self.transcribe_chunk(audio_data)
            if text:
                yield text

    def __del__(self):
        """Clean up resources."""
        self.stop_streaming()

        if self._temp_dir:
            self._temp_dir.cleanup()


class FasterWhisperTranscriber(BaseSpeechRecognizer):
    """
    Transcribe audio using Faster Whisper models.

    This class provides both batch and streaming transcription of audio using the
    faster-whisper implementation, which offers better performance than the standard
    Whisper implementation.
    """

    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 FasterWhisperTranscriber.

        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
        """
        if not FASTER_WHISPER_AVAILABLE:
            raise ImportError("faster-whisper is required for FasterWhisperTranscriber")

        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

        # Initialize the model
        self._initialize_model()

    def _initialize_model(self):
        """Initialize the Whisper model."""
        try:
            logger.info(f"Loading Faster 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"Faster 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 Faster 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("Faster 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"Error transcribing audio: {e}")
            raise

    def start_streaming(self) -> bool:
        """
        Start streaming mode for real-time transcription.

        Returns:
            True if streaming started successfully, False otherwise
        """
        if self._running:
            logger.warning("Streaming already running")
            return True

        try:
            # Clear queues
            while not self._audio_queue.empty():
                self._audio_queue.get_nowait()

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

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

            logger.info("Streaming started")
            return True

        except Exception as e:
            logger.error(f"Error starting streaming: {e}")
            self._running = False
            return False

    def stop_streaming(self) -> None:
        """Stop streaming mode."""
        if not self._running:
            return

        self._running = False

        if self._processing_thread:
            self._processing_thread.join(timeout=2.0)
            self._processing_thread = None

        logger.info("Streaming stopped")

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

        Args:
            audio: Audio chunk as numpy array
        """
        if not self._running:
            logger.warning("Streaming not running, starting now")
            self.start_streaming()

        try:
            self._audio_queue.put(audio)
        except Exception as e:
            logger.error(f"Error adding audio chunk to queue: {e}")

    def _process_audio_queue(self) -> None:
        """Process audio chunks from the queue in a background thread."""
        buffer = []
        buffer_duration = 0.0

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

                # Add to buffer
                buffer.append(audio_chunk)
                chunk_duration = len(audio_chunk) / self.sample_rate
                buffer_duration += chunk_duration

                # Process if buffer is full
                if buffer_duration >= self.streaming_buffer_size:
                    # Combine chunks
                    audio_data = np.concatenate(buffer)

                    # Transcribe
                    result = self.transcribe(
                        audio_data,
                        language=self.language,
                        translate=self.translate
                    )

                    # Add to transcript queue if there's text
                    if result and result.get("text", "").strip():
                        self._transcript_queue.put(result)

                    # Reset buffer
                    buffer = []
                    buffer_duration = 0.0

            except Exception as e:
                logger.error(f"Error processing audio chunk: {e}")
                time.sleep(0.1)  # Prevent CPU spinning on errors

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

        Args:
            timeout: Time to wait for a transcript (0 for non-blocking)

        Returns:
            Transcript dict, or None if no transcript is available
        """
        try:
            return self._transcript_queue.get(block=timeout > 0, timeout=timeout)
        except queue.Empty:
            return None

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

        Args:
            audio_chunks: Generator yielding audio chunks

        Yields:
            Transcribed text as it becomes available
        """
        # Start streaming if not already running
        if not self._running:
            self.start_streaming()

        try:
            # Process each chunk
            for chunk in audio_chunks:
                self.add_audio_chunk(chunk)

                # Check for transcripts
                transcript = self.get_transcript()
                if transcript:
                    yield transcript.get("text", "")

        finally:
            # Stop streaming when done
            self.stop_streaming()

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

        Args:
            audio: Audio chunk as numpy array

        Returns:
            Transcribed text, or None if no speech was detected
        """
        # Run transcription in a thread pool to avoid blocking
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            None,
            lambda: self.transcribe(
                audio,
                language=self.language,
                translate=self.translate
            )
        )

        # Extract text from result
        if result and result.get("text", "").strip():
            return result.get("text", "").strip()

        return None

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

        Args:
            audio_chunks: Async generator yielding audio chunks

        Yields:
            Transcribed text as it becomes available
        """
        buffer = []
        buffer_duration = 0.0

        async for chunk in audio_chunks:
            # Add to buffer
            buffer.append(chunk)
            chunk_duration = len(chunk) / self.sample_rate
            buffer_duration += chunk_duration

            # Process if buffer is full
            if buffer_duration >= self.streaming_buffer_size:
                # Combine chunks
                audio_data = np.concatenate(buffer)

                # Transcribe
                text = await self.transcribe_chunk(audio_data)

                # Yield text if not empty
                if text:
                    yield text

                # Reset buffer
                buffer = []
                buffer_duration = 0.0

        # Process any remaining audio
        if buffer:
            audio_data = np.concatenate(buffer)
            text = await self.transcribe_chunk(audio_data)
            if text:
                yield text

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

        Args:
            model_size: New model size/name
        """
        if self.model_name == model_size:
            return

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

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

        # Update model name
        self.model_name = model_size

        # Re-initialize model
        self._initialize_model()

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

    def __del__(self):
        """Clean up resources."""
        self.stop_streaming()

        if self._temp_dir:
            self._temp_dir.cleanup()
