"""
Home Assistant Faster Whisper Integration.

This module provides integration with the Home Assistant Faster Whisper component
for speech recognition.
"""

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

from fmusvid.ai.speech import BaseSpeechRecognizer

logger = logging.getLogger(__name__)

# Try to import Home Assistant Faster Whisper component
try:
    # Import the necessary Home Assistant components
    from homeassistant.components import stt
    from wyoming_faster_whisper.faster_whisper import WhisperModel
    HOME_ASSISTANT_AVAILABLE = True
except ImportError:
    logger.warning("Home Assistant Faster Whisper component not found. "
                   "To use this integration, install the Home Assistant and Wyoming Faster Whisper packages.")
    HOME_ASSISTANT_AVAILABLE = False


class HomeAssistantWhisperIntegration(BaseSpeechRecognizer):
    """
    Integration with Home Assistant's Faster Whisper component.

    This class implements the BaseSpeechRecognizer interface and uses the
    Home Assistant Faster Whisper component for speech recognition.
    """

    def __init__(
        self,
        model: str = "medium",
        language: Optional[str] = None,
        translate: bool = False,
        models_dir: Optional[str] = None,
        buffer_size: int = 5,
        streaming_buffer_size: float = 5.0,
        chunk_duration: float = 1.0,
        sample_rate: int = 16000,
        cpu_threads: int = 0,
        beam_size: int = 5,
        **kwargs
    ):
        """
        Initialize the Home Assistant Whisper Integration.

        Args:
            model: Whisper model size ('tiny', 'base', 'small', 'medium', 'large')
            language: Language code (None for auto-detection)
            translate: Whether to translate to English
            models_dir: Directory to store models (default is None, which uses the default location)
            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)
            cpu_threads: Number of CPU threads to use (0 for auto)
            beam_size: Beam size for beam search
        """
        if not HOME_ASSISTANT_AVAILABLE:
            raise ImportError("Home Assistant Faster Whisper component is required for this integration")

        self.model_name = model
        self.language = language or "en"
        self.translate = translate
        self.models_dir = models_dir
        self.buffer_size = buffer_size
        self.streaming_buffer_size = streaming_buffer_size
        self.chunk_duration = chunk_duration
        self.sample_rate = sample_rate
        self.cpu_threads = cpu_threads
        self.beam_size = beam_size

        # 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}'...")

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

            # Initialize the model
            model_path = self.model_name
            if self.models_dir:
                if not os.path.exists(self.models_dir):
                    os.makedirs(self.models_dir, exist_ok=True)
                model_path = os.path.join(self.models_dir, self.model_name)

            self._model = WhisperModel(
                model_path,
                cpu_threads=self.cpu_threads
            )

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

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

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

        Args:
            audio: Audio file path, numpy array or bytes
            language: Language code (None for auto-detection)
            translate: Whether to translate to English
            **kwargs: Additional arguments to pass to the transcriber

        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
        beam_size = kwargs.get("beam_size", self.beam_size)

        try:
            start_time = time.time()

            # Handle different audio input types
            audio_data = None
            if isinstance(audio, str):
                # File path - keep as is
                audio_data = audio
            elif isinstance(audio, np.ndarray):
                # Numpy array
                import soundfile as sf
                sf.write(self._temp_audio_file, audio, self.sample_rate)
                audio_data = self._temp_audio_file
            elif isinstance(audio, bytes):
                # Bytes
                with open(self._temp_audio_file, "wb") as f:
                    f.write(audio)
                audio_data = self._temp_audio_file
            else:
                raise ValueError("Audio must be a file path, numpy array, or bytes")

            # Transcribe audio
            task = "translate" if translate else "transcribe"
            segments, info = self._model.transcribe(
                audio_data,
                beam_size=beam_size,
                language=language,
                task=task
            )

            # Extract text and segments
            segments_list = []
            text_parts = []

            for i, segment in enumerate(segments):
                text_parts.append(segment.text)
                segments_list.append({
                    "id": i,
                    "text": segment.text,
                    "start": segment.start,
                    "end": segment.end
                })

            # Combine all text
            full_text = " ".join(text_parts)

            # Format response similar to Whisper models
            result = {
                "text": full_text,
                "segments": segments_list,
                "language": info.language,
                "language_probability": info.language_probability
            }

            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()
