"""DashScope (Aliyun Bailian) streaming ASR STT handler — Omni realtime API.

Uses dashscope.audio.qwen_omni.OmniRealtimeConversation in ASR-only mode
(output_modalities=[TEXT] + input audio transcription) with the
qwen3-asr-flash-realtime model. Each VAD segment is uploaded (streamed in
chunks) and committed; partial/final transcriptions are yielded to the
pipeline. VAD stays local, exactly as faster-whisper used to behave.
"""

from __future__ import annotations

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

import numpy as np

from speech_to_speech.STT.base_stt_handler import BaseSTTHandler
from speech_to_speech.pipeline.handler_types import STTIn, STTOut
from speech_to_speech.pipeline.messages import PartialTranscription, Transcription

logger = logging.getLogger(__name__)

DEFAULT_MODEL = "qwen3-asr-flash-realtime"
DEFAULT_SAMPLE_RATE = 16000
CONNECT_RETRIES = 4
CONNECT_BACKOFF_S = 2.0
EVENT_TIMEOUT_S = 15.0
APPEND_CHUNK_SAMPLES = 16000 * 2  # 2s chunks


class DashScopeASRSTTHandler(BaseSTTHandler):
    """Streaming Chinese ASR via DashScope Omni realtime conversation (ASR-only)."""

    def setup(
        self,
        api_key: str = "",
        model: str = DEFAULT_MODEL,
        language: str = "zh",
        sample_rate: int = DEFAULT_SAMPLE_RATE,
        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 --dashscope_asr_api_key or set env).")
        self.model = model
        self.language = language
        self.sample_rate = int(sample_rate)
        logger.info(
            "DashScopeASRSTTHandler ready (model=%s, lang=%s, sr=%d)",
            self.model,
            self.language,
            self.sample_rate,
        )

    @staticmethod
    def _float32_to_pcm16(audio: np.ndarray) -> bytes:
        if audio.dtype == np.int16:
            return audio.tobytes()
        return (np.clip(audio, -1.0, 1.0) * 32767).astype(np.int16).tobytes()

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

    def process(self, vad_audio: STTIn) -> Iterator[STTOut]:
        # v1: cloud ASR opens one websocket session per utterance. Progressive
        # (live) segments are skipped — transcribing every 0.5s revision would
        # re-open sessions serially and stall the pipeline. Launch with
        # --enable_live_transcription false so VAD emits one final segment per
        # utterance; keep this guard as defense-in-depth.
        if vad_audio.mode == "progressive":
            return

        from dashscope.audio.qwen_omni import MultiModality, OmniRealtimeConversation
        from dashscope.audio.qwen_omni.omni_realtime import TranscriptionParams

        pcm = self._float32_to_pcm16(vad_audio.audio)

        events: "queue.Queue[tuple[str, str]]" = queue.Queue()
        partial_text = ""
        final_text = ""

        class Callback:
            def on_open(self) -> None:
                events.put(("open", ""))

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

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

            def on_event(self, message: Any) -> None:
                event_type = message.get("type") if isinstance(message, dict) else str(message)
                if event_type == "conversation.item.input_audio_transcription.text":
                    stash = message.get("stash", "") if isinstance(message, dict) else ""
                    events.put(("partial", str(stash)))
                elif event_type == "conversation.item.input_audio_transcription.completed":
                    transcript = message.get("transcript", "") if isinstance(message, dict) else ""
                    events.put(("final", str(transcript)))

        conversation = OmniRealtimeConversation(
            model=self.model,
            url="wss://dashscope.aliyuncs.com/api-ws/v1/realtime",
            callback=Callback(),
        )
        self._connect_with_retry(conversation)
        try:
            conversation.update_session(
                output_modalities=[MultiModality.TEXT],
                enable_input_audio_transcription=True,
                # Disable the server-side VAD: it races our explicit commit()
                # (auto-consuming the audio buffer can make commit fail with
                # "input messages do not contain elements with the role of user").
                enable_turn_detection=False,
                transcription_params=TranscriptionParams(
                    language=self.language,
                    sample_rate=self.sample_rate,
                    input_audio_format="pcm",
                ),
            )
            chunk_bytes = APPEND_CHUNK_SAMPLES * 2
            for i in range(0, len(pcm), chunk_bytes):
                conversation.append_audio(base64.b64encode(pcm[i : i + chunk_bytes]).decode())
                time.sleep(0.05)
            time.sleep(1.0)  # let the server ingest the trailing audio before commit
            conversation.commit()

            deadline = time.monotonic() + EVENT_TIMEOUT_S
            while time.monotonic() < deadline:
                try:
                    kind, payload = events.get(timeout=1.0)
                except queue.Empty:
                    continue
                if kind == "partial":
                    partial_text = payload
                    if vad_audio.mode == "progressive":
                        yield PartialTranscription(
                            text=payload,
                            turn_id=vad_audio.turn_id,
                            turn_revision=vad_audio.turn_revision,
                        )
                elif kind == "final":
                    final_text = payload
                    if not vad_audio.mode == "progressive":
                        yield Transcription(
                            text=payload,
                            language_code=self.language,
                            turn_id=vad_audio.turn_id,
                            turn_revision=vad_audio.turn_revision,
                            speech_stopped_at_s=vad_audio.created_at_s,
                        )
                        return
                elif kind == "error":
                    logger.error("DashScope ASR error: %s", payload)
                    return
                elif kind == "close":
                    break
        finally:
            try:
                conversation.close()
            except Exception:  # noqa: BLE001
                pass

        # Fallback: progressive mode never got a final event.
        if vad_audio.mode != "progressive" and final_text:
            yield Transcription(
                text=final_text,
                language_code=self.language,
                turn_id=vad_audio.turn_id,
                turn_revision=vad_audio.turn_revision,
                speech_stopped_at_s=vad_audio.created_at_s,
            )
