from __future__ import annotations

import logging
from threading import Event
from time import perf_counter
from typing import Any, Iterator

import numpy as np
from rich.console import Console

from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.cancel_scope import CancelScope
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse
from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker

logger = logging.getLogger(__name__)
console = Console()


class PocketTTSHandler(BaseHandler[TTSIn, TTSOut]):
    """
    Handler for Pocket TTS model from Kyutai Labs.
    Supports streaming audio generation with voice cloning.
    """

    def setup(
        self,
        should_listen: Event,
        device: str = "cpu",
        voice: str = "alba",  # Default voice from catalog
        sample_rate: int = 16000,  # Match the pipeline's audio output (LocalAudioStreamer uses 16kHz)
        blocksize: int = 512,
        max_tokens: int = 50,
        gen_kwargs: dict[str, Any] | None = None,  # For compatibility with pipeline, not used
        cancel_scope: CancelScope | None = None,
        speculative_turns: SpeculativeTurnTracker | None = None,
    ) -> None:
        """
        Initialize Pocket TTS handler.

        Args:
            should_listen: Event to control when to start listening again
            device: Device to run model on ('cpu', 'cuda', 'mps')
            voice: Voice to use. Can be:
                - A preset name: 'alba', 'marius', 'javert', 'jean', 'fantine', 'cosette', 'eponine', 'azelma'
                - A local audio file path
                - A Hugging Face path like "hf://kyutai/tts-voices/..."
            sample_rate: Output sample rate (pocket-tts generates at 24kHz and will be resampled to this rate). Default 16kHz matches the pipeline's audio output.
            blocksize: Size of audio blocks to yield
            max_tokens: Maximum tokens to generate
        """
        self.should_listen = should_listen
        self.cancel_scope = cancel_scope
        self.speculative_turns = speculative_turns
        self.device = device
        self.voice = voice
        self.sample_rate = sample_rate
        self.blocksize = blocksize
        self.max_tokens = max_tokens

        # Suppress verbose logging from pocket_tts library
        logging.getLogger("pocket_tts").setLevel(logging.WARNING)
        logging.getLogger("pocket_tts.models.tts_model").setLevel(logging.WARNING)
        logging.getLogger("pocket_tts.utils.utils").setLevel(logging.WARNING)

        # Import and load model
        from pocket_tts import TTSModel

        logger.info("Loading Pocket TTS model")
        self.model = TTSModel.load_model()

        # Move model to specified device
        if device == "cuda":
            self.model = self.model.cuda()
        elif device == "mps":
            self.model = self.model.to("mps")
        elif device != "cpu":
            self.model = self.model.to(device)

        logger.info(f"Pocket TTS model moved to {device}")

        # Load voice state
        logger.info(f"Loading voice: {voice}")
        self.voice_state = self.model.get_state_for_audio_prompt(voice)

        logger.info(f"Pocket TTS model sample rate: {self.model.sample_rate}")

    @property
    def min_time_to_debug(self) -> float:
        """
        Override to suppress logging for individual audio chunks.
        Pocket TTS yields many small chunks (~10-20ms each), which would flood logs.
        Only log if a chunk takes unusually long (>100ms).
        """
        return 0.1  # 100ms threshold

    def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
        speculative_turns = getattr(self, "speculative_turns", None)
        if isinstance(tts_input, EndOfResponse):
            if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
                tts_input.turn_id,
                tts_input.turn_revision,
            ):
                return
            yield AUDIO_RESPONSE_DONE
            return

        if speculative_turns and not speculative_turns.is_latest_after_reopen_grace(
            tts_input.turn_id,
            tts_input.turn_revision,
        ):
            logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision)
            return
        if speculative_turns:
            speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision)

        gen = self.cancel_scope.generation if self.cancel_scope else None
        language_code = tts_input.language_code
        text = tts_input.text
        logger.debug(f"Received language code: {language_code}")

        console.print(f"[green]ASSISTANT: {text}")

        # Generate audio stream
        logger.debug(f"Generating audio for: {text[:50]}...")

        pipeline_start = perf_counter()
        first_chunk = True

        # Calculate target chunk size in original sample rate
        # We need enough samples so that after resampling we get blocksize samples
        needs_resampling = self.model.sample_rate != self.sample_rate
        if needs_resampling:
            from scipy.signal import resample_poly

            # Target chunk size in original sample rate (before resampling)
            resample_ratio = self.model.sample_rate / self.sample_rate
            target_chunk_size = int(self.blocksize * resample_ratio)
            # Calculate up/down factors for polyphase resampling
            # For 24kHz -> 16kHz: 16000/24000 = 2/3
            from math import gcd

            g = gcd(self.sample_rate, self.model.sample_rate)
            self._resample_up = self.sample_rate // g
            self._resample_down = self.model.sample_rate // g
        else:
            target_chunk_size = self.blocksize

        # Buffer to accumulate audio until we have enough for one block
        audio_buffer = []
        buffer_size = 0

        for audio_chunk in self.model.generate_audio_stream(
            self.voice_state,
            text,
            max_tokens=self.max_tokens,
            copy_state=True,  # Don't modify the original voice state
        ):
            if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen):
                logger.info("TTS generation cancelled (interruption)")
                return
            if first_chunk:
                logger.debug(f"Time to first audio: {perf_counter() - pipeline_start:.3f}s")
                first_chunk = False

            # Convert from torch tensor to numpy and add to buffer
            audio_np = audio_chunk.cpu().numpy()
            audio_buffer.append(audio_np)
            buffer_size += len(audio_np)

            # Process buffer when we have enough samples for at least one block
            while buffer_size >= target_chunk_size:
                # Concatenate buffer
                concatenated = np.concatenate(audio_buffer)

                # Extract enough samples for one block
                chunk_to_process = concatenated[:target_chunk_size]
                remainder = concatenated[target_chunk_size:]

                # Resample this chunk if needed
                if needs_resampling:
                    chunk_resampled = resample_poly(
                        chunk_to_process,
                        up=self._resample_up,
                        down=self._resample_down,
                    )
                else:
                    chunk_resampled = chunk_to_process

                # Convert to int16 format expected by audio output
                audio_int16 = (chunk_resampled * 32768).astype(np.int16)

                # Ensure exact blocksize (pad or trim if resampling caused slight size differences)
                if len(audio_int16) < self.blocksize:
                    audio_int16 = np.pad(audio_int16, (0, self.blocksize - len(audio_int16)))
                elif len(audio_int16) > self.blocksize:
                    audio_int16 = audio_int16[: self.blocksize]

                yield audio_int16

                # Update buffer with remainder
                audio_buffer = [remainder] if len(remainder) > 0 else []
                buffer_size = len(remainder)

        # Process any remaining audio in buffer
        if audio_buffer and buffer_size > 0:
            concatenated = np.concatenate(audio_buffer)

            # Resample the remainder
            if needs_resampling:
                concatenated = resample_poly(
                    concatenated,
                    up=self._resample_up,
                    down=self._resample_down,
                )

            # Convert to int16
            audio_int16 = (concatenated * 32768).astype(np.int16)

            # Yield in blocks
            for i in range(0, len(audio_int16), self.blocksize):
                chunk = audio_int16[i : i + self.blocksize]
                # Pad last chunk if needed
                if len(chunk) < self.blocksize:
                    chunk = np.pad(chunk, (0, self.blocksize - len(chunk)))
                yield chunk
