"""Qwen3-TTS (Aliyun Bailian realtime TTS) streaming handler.

Uses the dashscope QwenTtsRealtime WebSocket API (qwen3-tts-flash-realtime)
and streams the coalesced assistant text. The engine returns 24 kHz int16 PCM
inside response.audio.delta events (base64); the handler resamples it to the
16 kHz pipeline rate so the browser audio path and the LiveTalking /humanpcm
lip-sync bridge keep working unchanged. (speech-to-speech v0.2.11 pipeline.)
"""

from __future__ import annotations

import base64
import logging
import os
import queue
import threading
import time
from typing import Any, Iterator, Optional

import numpy as np

from speech_to_speech.baseHandler import BaseHandler
from speech_to_speech.pipeline.control import SESSION_END, is_control_message
from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END, EndOfResponse, TTSInput

logger = logging.getLogger(__name__)

PIPELINE_SR = 16000
ENGINE_SR = 24000
DEFAULT_MODEL = "qwen3-tts-flash-realtime"
DEFAULT_VOICE = "Cherry"
DEFAULT_BLOCKSIZE = 512
CONNECT_RETRIES = 4
CONNECT_BACKOFF_S = 2.0
EVENT_TIMEOUT_S = 60.0


class Qwen3CloudTTSHandler(BaseHandler[TTSIn, TTSOut]):
    def setup(
        self,
        should_listen: Any = None,  # accepted for setup_args=(should_listen,) compatibility
        api_key: str = "",
        model: str = DEFAULT_MODEL,
        voice: str = DEFAULT_VOICE,
        sample_rate: int = PIPELINE_SR,
        blocksize: int = DEFAULT_BLOCKSIZE,
        gen_kwargs: dict[str, Any] | None = None,
    ) -> None:
        import dashscope  # pip install dashscope>=1.20 (lazy: only when this backend is used)

        dashscope.api_key = api_key or os.environ.get("DASHSCOPE_API_KEY", "")
        if not dashscope.api_key:
            raise ValueError("DASHSCOPE_API_KEY is not set (pass --qwen3_cloud_api_key or set env).")
        self.model = model
        self.voice = voice
        self.sample_rate = int(sample_rate)
        self.blocksize = int(blocksize)
        logger.info(
            "Qwen3CloudTTSHandler ready (model=%s, voice=%s, engine_sr=%d -> pipeline_sr=%d, blocksize=%d)",
            self.model,
            self.voice,
            ENGINE_SR,
            self.sample_rate,
            self.blocksize,
        )

    def _coalesce_pending_tts_input(self, current_input: TTSInput) -> tuple[str, Optional[str]]:
        """Merge already-queued text chunks of the same response into one utterance."""
        if not hasattr(self.queue_in, "mutex") or not hasattr(self.queue_in, "queue"):
            return current_input.text, current_input.language_code

        parts = [current_input.text.strip()] if current_input.text and current_input.text.strip() else []
        language_code = current_input.language_code
        with self.queue_in.mutex:
            while self.queue_in.queue:
                next_item = self.queue_in.queue[0]
                if is_control_message(next_item, SESSION_END.kind):
                    break
                if isinstance(next_item, bytes) and next_item == PIPELINE_END:
                    break
                if isinstance(next_item, EndOfResponse):
                    break
                if not isinstance(next_item, TTSInput):
                    break
                if current_input.turn_id != next_item.turn_id or current_input.turn_revision != next_item.turn_revision:
                    break
                if (
                    language_code is not None
                    and next_item.language_code is not None
                    and next_item.language_code != language_code
                ):
                    break
                self.queue_in.queue.popleft()
                if next_item.text.strip():
                    parts.append(next_item.text.strip())
                if language_code is None:
                    language_code = next_item.language_code
        return " ".join(parts).strip(), language_code

    def _connect_with_retry(self, synthesizer: Any) -> None:
        last_exc: Exception | None = None
        for attempt in range(CONNECT_RETRIES):
            try:
                synthesizer.connect()
                return
            except Exception as e:  # noqa: BLE001
                last_exc = e
                logger.warning("Qwen3CloudTTS connect attempt %d failed: %s", attempt + 1, e)
                time.sleep(CONNECT_BACKOFF_S)
        raise RuntimeError(f"Qwen3CloudTTS could not connect after {CONNECT_RETRIES} attempts: {last_exc}")

    def _resample_24k_to_16k(self, audio: np.ndarray) -> np.ndarray:
        """int16 24k mono -> int16 16k mono (resample_poly up=2/down=3)."""
        if self.sample_rate == ENGINE_SR or self.sample_rate == 0:
            return audio
        from scipy.signal import resample_poly

        f32 = audio.astype(np.float32) / 32768.0
        out = resample_poly(f32, up=2, down=3)  # 24000 -> 16000
        return np.clip(out * 32768.0, -32768, 32767).astype(np.int16)

    def process(self, tts_input: TTSIn) -> Iterator[TTSOut]:
        if isinstance(tts_input, EndOfResponse):
            yield AUDIO_RESPONSE_DONE
            return

        text, _language = self._coalesce_pending_tts_input(tts_input)
        text = text or "Hello."
        logger.info("Qwen3CloudTTS synthesizing: %s", text)

        from dashscope.audio.qwen_tts_realtime import QwenTtsRealtime
        from dashscope.audio.tts_v2 import AudioFormat

        audio_q: "queue.Queue[tuple[str, Any]]" = queue.Queue()

        class Callback:
            def on_open(self) -> None:
                pass

            def on_close(self, *args: Any) -> None:
                audio_q.put(("close", None))

            def on_error(self, message: Any) -> None:
                audio_q.put(("error", str(message)))

            def on_data(self, data: bytes) -> None:
                audio_q.put(("data", data))

            def on_event(self, message: Any) -> None:
                event_type = message.get("type") if isinstance(message, dict) else str(message)
                if event_type == "response.audio.delta":
                    delta = (message.get("delta") or "") if isinstance(message, dict) else ""
                    if delta:
                        try:
                            audio_q.put(("pcm", base64.b64decode(delta)))
                        except Exception as e:  # noqa: BLE001
                            logger.warning("audio.delta decode failed: %s", e)
                elif event_type == "response.done":
                    audio_q.put(("done", None))

        synthesizer = QwenTtsRealtime(model=self.model, callback=Callback())
        self._connect_with_retry(synthesizer)
        try:
            synthesizer.update_session(
                voice=self.voice,
                response_format=AudioFormat.PCM_24000HZ_MONO_16BIT,
                mode="server_commit",
            )
            synthesizer.append_text(text)
            synthesizer.finish()

            leftover = b""
            out_leftover = np.array([], dtype=np.int16)
            block = self.blocksize * 2  # bytes per int16 block
            first_chunk = True
            while True:
                try:
                    kind, payload = audio_q.get(timeout=EVENT_TIMEOUT_S)
                except queue.Empty:
                    logger.error("Qwen3CloudTTS timed out (no events for %ds)", EVENT_TIMEOUT_S)
                    break
                if kind == "pcm":
                    if first_chunk:
                        logger.info("Qwen3CloudTTS first audio chunk received")
                        first_chunk = False
                    leftover += payload
                    # resample in ~0.5s batches, keep the resampled tail across
                    # batches so audio stays continuous
                    batch_bytes = ENGINE_SR // 2 * 2  # 0.5s @ 24k (int16 bytes)
                    while len(leftover) >= batch_bytes:
                        pcm24 = np.frombuffer(leftover[:batch_bytes], dtype=np.int16)
                        leftover = leftover[batch_bytes:]
                        pcm16 = self._resample_24k_to_16k(pcm24)
                        combined = np.concatenate([out_leftover, pcm16])
                        n = (len(combined) // self.blocksize) * self.blocksize
                        for i in range(0, n, self.blocksize):
                            yield combined[i : i + self.blocksize]
                        out_leftover = combined[n:]
                elif kind == "done":
                    break
                elif kind == "error":
                    logger.error("Qwen3CloudTTS error: %s", payload)
                    break
                elif kind == "close":
                    break

            if leftover:
                pcm24 = np.frombuffer(leftover, dtype=np.int16)
                pcm16 = self._resample_24k_to_16k(pcm24)
                out_leftover = np.concatenate([out_leftover, pcm16])
            if len(out_leftover) >= self.blocksize:
                n = (len(out_leftover) // self.blocksize) * self.blocksize
                for i in range(0, n, self.blocksize):
                    yield out_leftover[i : i + self.blocksize]
                out_leftover = out_leftover[n:]
            if len(out_leftover) > 0:
                yield out_leftover
        finally:
            try:
                synthesizer.close()
            except Exception:  # noqa: BLE001
                pass
