"""Qwen-Audio Realtime speech-to-speech service for Pipecat.

The adapter intentionally depends only on Pipecat's public frame/service APIs.
Provider-specific WebSocket events stay in this module so the shared pipeline
does not need to know about Qwen-Audio protocol details.
"""

from __future__ import annotations

import asyncio
import base64
import inspect
import json
from collections.abc import Awaitable, Callable
from typing import Any
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from uuid import uuid4

from loguru import logger
from pipecat.frames.frames import (
    CancelFrame,
    EndFrame,
    Frame,
    InputAudioRawFrame,
    InterruptionFrame,
    LLMMessagesAppendFrame,
    OutputTransportMessageUrgentFrame,
    StartFrame,
    TTSAudioRawFrame,
)
from pipecat.processors.frame_processor import FrameDirection
from pipecat.services.ai_service import AIService
from pipecat.services.settings import ServiceSettings
from pipecat.utils.time import time_now_iso8601
from websockets.asyncio.client import connect as websocket_connect
from websockets.protocol import State


DEFAULT_QWEN_AUDIO_REALTIME_MODEL = "qwen-audio-3.0-realtime-flash"
DEFAULT_QWEN_AUDIO_REALTIME_VOICE = "longanqian"
QWEN_INPUT_SAMPLE_RATE = 16_000
QWEN_OUTPUT_SAMPLE_RATE = 24_000
SUPPORTED_TURN_DETECTION_MODES = frozenset({"server_vad", "smart_turn"})

ExtraEventHandler = Callable[[dict[str, Any]], Awaitable[None] | None]


class QwenAudioRealtimeService(AIService):
    """Translate Pipecat frames to Qwen-Audio Realtime WebSocket events.

    ``extra_event_handlers`` is deliberately small: future features such as
    Function Calling can subscribe to provider events without changing the
    audio, transcript, or interruption paths implemented here.
    """

    def __init__(
        self,
        *,
        api_key: str,
        model: str = DEFAULT_QWEN_AUDIO_REALTIME_MODEL,
        base_url: str,
        instructions: str = "",
        voice: str = DEFAULT_QWEN_AUDIO_REALTIME_VOICE,
        input_sample_rate: int = QWEN_INPUT_SAMPLE_RATE,
        output_sample_rate: int = QWEN_OUTPUT_SAMPLE_RATE,
        turn_detection_mode: str = "server_vad",
        vad_threshold: float = 0.5,
        silence_duration_ms: int = 800,
        max_history_turns: int = 20,
        extra_event_handlers: dict[str, ExtraEventHandler] | None = None,
        **kwargs,
    ) -> None:
        if turn_detection_mode not in SUPPORTED_TURN_DETECTION_MODES:
            supported = ", ".join(sorted(SUPPORTED_TURN_DETECTION_MODES))
            raise ValueError(
                f"Unsupported Qwen turn detection mode {turn_detection_mode!r}; "
                f"expected one of: {supported}"
            )
        if not -1.0 <= vad_threshold <= 1.0:
            raise ValueError("Qwen VAD threshold must be between -1.0 and 1.0")
        if not 200 <= silence_duration_ms <= 6_000:
            raise ValueError(
                "Qwen VAD silence duration must be between 200 and 6000 ms"
            )
        if not 1 <= max_history_turns <= 50:
            raise ValueError("Qwen max history turns must be between 1 and 50")

        super().__init__(settings=ServiceSettings(model=model), **kwargs)
        self._api_key = api_key
        self._model = model
        self._base_url = base_url
        self._instructions = instructions
        self._voice = voice
        self._input_sample_rate = input_sample_rate
        self._output_sample_rate = output_sample_rate
        self._turn_detection_mode = turn_detection_mode
        self._vad_threshold = vad_threshold
        self._silence_duration_ms = silence_duration_ms
        self._max_history_turns = max_history_turns
        self._extra_event_handlers = extra_event_handlers or {}

        self._websocket = None
        self._receive_task: asyncio.Task | None = None
        self._session_ready = asyncio.Event()
        self._pending_events: list[dict[str, Any]] = []
        self._warned_input_sample_rate = False

        self._response_active = False
        self._audio_suppressed = False
        self._assistant_turn_id: str | None = None
        self._assistant_text = ""
        self._assistant_timestamp = ""
        self._greeting_request_item_id: str | None = None

    async def start(self, frame: StartFrame) -> None:
        await super().start(frame)
        if not self._api_key or not self._model or not self._base_url:
            await self.push_error(
                "Qwen-Audio Realtime requires api_key, model, and base_url",
                fatal=True,
            )
            return
        await self._connect()

    async def stop(self, frame: EndFrame) -> None:
        await self._disconnect()
        await super().stop(frame)

    async def cancel(self, frame: CancelFrame) -> None:
        await self._disconnect()
        await super().cancel(frame)

    async def cleanup(self) -> None:
        await self._disconnect()
        await super().cleanup()

    async def process_frame(self, frame: Frame, direction: FrameDirection) -> None:
        await super().process_frame(frame, direction)

        if isinstance(frame, InputAudioRawFrame):
            if (
                frame.sample_rate != self._input_sample_rate
                and not self._warned_input_sample_rate
            ):
                self._warned_input_sample_rate = True
                logger.warning(
                    "Qwen-Audio Realtime expected {} Hz input, received {} Hz",
                    self._input_sample_rate,
                    frame.sample_rate,
                )
            await self._send_event(
                {
                    "type": "input_audio_buffer.append",
                    "audio": base64.b64encode(frame.audio).decode("ascii"),
                }
            )
            return

        if isinstance(frame, LLMMessagesAppendFrame):
            for message in frame.messages:
                text = self._message_text(message)
                if text:
                    await self.send_text(
                        text,
                        run_immediately=frame.run_llm is not False,
                    )
            return

        if isinstance(frame, InterruptionFrame):
            await self._cancel_active_response()

        await self.push_frame(frame, direction)

    async def send_text(self, text: str, *, run_immediately: bool = True) -> None:
        """Append text to Qwen's conversation and optionally start a response."""
        if not text:
            return
        await self._send_event(
            {
                "type": "conversation.item.create",
                "item": {
                    "type": "message",
                    "role": "user",
                    "content": [{"type": "input_text", "text": text}],
                },
            }
        )
        if run_immediately:
            await self._send_event({"type": "response.create"})

    async def interrupt(self) -> None:
        """Cancel current model output and notify the rest of the pipeline."""
        await self._cancel_active_response()
        await self.broadcast_interruption()

    async def speak(self, text: str) -> None:
        """Ask Qwen to speak a fixed greeting, then remove the hidden request.

        Qwen's documented ``response.create`` payload has no per-response
        instruction field. A temporary user item keeps this behavior within
        the supported protocol; it is deleted after the response completes.
        """
        if not text:
            return
        item_id = f"item_{uuid4().hex}"
        self._greeting_request_item_id = item_id
        await self._send_event(
            {
                "type": "conversation.item.create",
                "item": {
                    "id": item_id,
                    "type": "message",
                    "role": "user",
                    "content": [
                        {
                            "type": "input_text",
                            "text": (
                                "请直接朗读下面的开场白，不要增删、解释或添加前后缀：\n"
                                f"{text}"
                            ),
                        }
                    ],
                },
            }
        )
        await self._send_event({"type": "response.create"})

    async def update_instructions(self, instructions: str) -> None:
        """Update only instructions after startup.

        Qwen accepts voice and turn detection only in the first session update,
        so resending the full initial configuration would break live dynamic
        variable updates.
        """
        self._instructions = instructions
        if self._session_ready.is_set():
            await self._send_event(
                {
                    "type": "session.update",
                    "session": {"instructions": instructions},
                },
                wait_until_ready=False,
            )

    def _connection_url(self) -> str:
        parts = urlsplit(self._base_url)
        query = dict(parse_qsl(parts.query))
        query["model"] = self._model
        return urlunsplit(
            (parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment)
        )

    def _initial_session_config(self) -> dict[str, Any]:
        return {
            "modalities": ["text", "audio"],
            "instructions": self._instructions,
            "voice": self._voice,
            "input_audio_format": "pcm",
            "output_audio_format": "pcm",
            "turn_detection": self._turn_detection_config(),
            "max_history_turns": self._max_history_turns,
        }

    def _turn_detection_config(self) -> dict[str, Any]:
        if self._turn_detection_mode == "smart_turn":
            return {"type": "smart_turn"}
        return {
            "type": "server_vad",
            "threshold": self._vad_threshold,
            "silence_duration_ms": self._silence_duration_ms,
        }

    async def _connect(self) -> None:
        if self._websocket and self._websocket.state is State.OPEN:
            return
        try:
            self._websocket = await websocket_connect(
                self._connection_url(),
                additional_headers={
                    "Authorization": f"Bearer {self._api_key}",
                    "x-dashscope-dataInspection": "disable",
                },
                max_size=None,
                open_timeout=10,
            )
            self._receive_task = self.create_task(
                self._receive_messages(), name="qwen_audio_realtime_receive"
            )
        except Exception as exc:
            self._websocket = None
            await self.push_error(
                f"Qwen-Audio Realtime connection failed: {exc}",
                exception=exc,
                fatal=True,
            )

    async def _disconnect(self) -> None:
        current_task = asyncio.current_task()
        task = self._receive_task
        self._receive_task = None
        if task and task is not current_task:
            await self.cancel_task(task)

        websocket = self._websocket
        self._websocket = None
        self._session_ready.clear()
        self._pending_events.clear()
        self._response_active = False
        if websocket and websocket.state is State.OPEN:
            try:
                await websocket.close()
            except Exception:
                pass

    async def _receive_messages(self) -> None:
        websocket = self._websocket
        if not websocket:
            return
        try:
            async for raw_message in websocket:
                await self._handle_server_event(json.loads(raw_message))
        except Exception as exc:
            if self._websocket is websocket:
                await self.push_error(
                    f"Qwen-Audio Realtime receive failed: {exc}", exception=exc
                )
        finally:
            if self._websocket is websocket:
                self._websocket = None
                self._session_ready.clear()
            if self._receive_task is asyncio.current_task():
                self._receive_task = None

    async def _handle_server_event(self, event: dict[str, Any]) -> None:
        event_type = str(event.get("type") or "")
        if event_type == "session.created":
            await self._send_event(
                {
                    "type": "session.update",
                    "session": self._initial_session_config(),
                },
                wait_until_ready=False,
            )
        elif event_type == "session.updated":
            self._session_ready.set()
            pending, self._pending_events = self._pending_events, []
            for payload in pending:
                await self._send_event(payload, wait_until_ready=False)
        elif event_type == "response.created":
            self._response_active = True
            self._audio_suppressed = False
        elif event_type == "response.audio.delta":
            audio = event.get("delta")
            if audio and not self._audio_suppressed:
                await self.push_frame(
                    TTSAudioRawFrame(
                        base64.b64decode(audio),
                        self._output_sample_rate,
                        1,
                    )
                )
        elif event_type in {"response.audio_transcript.delta", "response.text.delta"}:
            if not self._audio_suppressed:
                await self._append_assistant_text(str(event.get("delta") or ""))
        elif event_type in {"response.audio_transcript.done", "response.text.done"}:
            transcript = str(event.get("transcript") or event.get("text") or "")
            if transcript and not self._audio_suppressed:
                if self._assistant_turn_id:
                    self._assistant_text = transcript
                else:
                    await self._append_assistant_text(transcript)
        elif event_type == "conversation.item.input_audio_transcription.completed":
            await self._send_transcript("user", str(event.get("transcript") or ""))
        elif event_type == "input_audio_buffer.speech_started":
            await self._cancel_active_response()
            await self.broadcast_interruption()
        elif event_type == "response.done":
            response = event.get("response")
            status = response.get("status") if isinstance(response, dict) else None
            interrupted = status in {"cancelled", "incomplete", "interrupted", "failed"}
            self._response_active = False
            await self._finish_assistant_text(interrupted=interrupted)
            await self._delete_greeting_request()
        elif event_type == "error":
            error = event.get("error")
            message = error.get("message") if isinstance(error, dict) else str(error)
            if "cancel" not in str(message).lower():
                await self.push_error(f"Qwen-Audio Realtime error: {message}")

        handler = self._extra_event_handlers.get(event_type)
        if handler:
            result = handler(event)
            if inspect.isawaitable(result):
                await result

    async def _cancel_active_response(self) -> None:
        if not self._response_active and not self._assistant_turn_id:
            return
        self._audio_suppressed = True
        if self._response_active:
            await self._send_event(
                {"type": "response.cancel"}, wait_until_ready=False
            )
        self._response_active = False
        await self._finish_assistant_text(interrupted=True)

    async def _delete_greeting_request(self) -> None:
        item_id = self._greeting_request_item_id
        self._greeting_request_item_id = None
        if item_id:
            await self._send_event(
                {"type": "conversation.item.delete", "item_id": item_id},
                wait_until_ready=False,
            )

    async def _send_event(
        self, payload: dict[str, Any], *, wait_until_ready: bool = True
    ) -> None:
        if wait_until_ready and not self._session_ready.is_set():
            self._pending_events.append(payload)
            return
        if not self._websocket or self._websocket.state is not State.OPEN:
            return
        event = {"event_id": f"event_{uuid4().hex}", **payload}
        await self._websocket.send(json.dumps(event, ensure_ascii=False))

    async def _append_assistant_text(self, delta: str) -> None:
        if not delta:
            return
        if not self._assistant_turn_id:
            self._assistant_turn_id = uuid4().hex
            self._assistant_timestamp = time_now_iso8601()
            await self._send_transport_message(
                {
                    "type": "assistant-text-start",
                    "turn_id": self._assistant_turn_id,
                    "timestamp": self._assistant_timestamp,
                }
            )
        self._assistant_text += delta
        await self._send_transport_message(
            {
                "type": "assistant-text-delta",
                "turn_id": self._assistant_turn_id,
                "delta": delta,
            }
        )

    async def _finish_assistant_text(self, *, interrupted: bool) -> None:
        if not self._assistant_turn_id:
            return
        await self._send_transport_message(
            {
                "type": "assistant-text-end",
                "turn_id": self._assistant_turn_id,
                "content": self._assistant_text,
                "interrupted": interrupted,
            }
        )
        self._assistant_turn_id = None
        self._assistant_text = ""
        self._assistant_timestamp = ""

    async def _send_transcript(self, role: str, content: str) -> None:
        if content:
            await self._send_transport_message(
                {
                    "type": "transcript",
                    "role": role,
                    "content": content,
                    "timestamp": time_now_iso8601(),
                }
            )

    async def _send_transport_message(self, message: dict[str, Any]) -> None:
        await self.push_frame(OutputTransportMessageUrgentFrame(message=message))

    @staticmethod
    def _message_text(message: Any) -> str:
        if not isinstance(message, dict):
            return ""
        content = message.get("content")
        if isinstance(content, str):
            return content.strip()
        if isinstance(content, list):
            return "\n".join(
                str(part.get("text") or "")
                for part in content
                if isinstance(part, dict) and part.get("type") == "text"
            ).strip()
        return ""