"""Simple Python WebSocket client for the Voice TTS server.

Streams text in chunks, receives base64 PCM audio and plays it via sounddevice.
Install dependencies:
    pip install websockets sounddevice
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import json
import sys
from typing import Any

import numpy as np
import sounddevice as sd
import websockets


class VoiceTTSClient:
    """WebSocket client that speaks LLM/agent output in real time."""

    def __init__(
        self,
        uri: str = "ws://localhost:8765/ws",
        voice_ref: str | None = None,
        voice_refs: dict[str, str] | None = None,
        language: str = "ru",
        speed: float = 1.0,
        emotion: str = "neutral",
        sample_rate: int = 24_000,
        block_size: int = 2048,
    ):
        self.uri = uri
        self.voice_ref = voice_ref
        self.voice_refs = voice_refs or {}
        self.language = language
        self.speed = speed
        self.emotion = emotion
        self.sample_rate = sample_rate
        self.block_size = block_size

        self._ws: websockets.WebSocketClientProtocol | None = None
        self._seq = 0
        self._audio_buffer: bytearray = bytearray()
        self._current_segment_seq: int | None = None
        self._stream: sd.RawOutputStream | None = None
        self._lock = asyncio.Lock()

    def _next_seq(self) -> int:
        self._seq += 1
        return self._seq

    def _send_dict(self, payload: dict[str, Any]) -> None:
        if self._ws is None:
            raise RuntimeError("WebSocket is not connected")
        asyncio.create_task(self._ws.send(json.dumps(payload, ensure_ascii=False)))

    async def connect(self) -> None:
        self._ws = await websockets.connect(self.uri)

        # Start audio output stream. sounddevice resamples if necessary.
        self._stream = sd.RawOutputStream(
            samplerate=self.sample_rate,
            channels=1,
            dtype="int16",
            blocksize=self.block_size,
            callback=self._audio_callback,
        )
        self._stream.start()

        init_msg = {
            "type": "init",
            "seq": self._next_seq(),
            "session_id": "python-client",
            "language": self.language,
            "speed": self.speed,
            "emotion": self.emotion,
        }
        if self.voice_ref:
            init_msg["voice_ref"] = self.voice_ref
        if self.voice_refs:
            init_msg["voice_refs"] = self.voice_refs

        await self._ws.send(json.dumps(init_msg, ensure_ascii=False))

    def _audio_callback(self, outdata: np.ndarray, frames: int, _time, _status) -> None:
        """Pull audio bytes from the buffer into the sounddevice stream."""
        needed = frames * 2  # int16 = 2 bytes
        available = len(self._audio_buffer)
        if available >= needed:
            chunk = bytes(self._audio_buffer[:needed])
            self._audio_buffer = self._audio_buffer[needed:]
        else:
            chunk = bytes(self._audio_buffer) + b"\x00" * (needed - available)
            self._audio_buffer = bytearray()
        outdata[:] = np.frombuffer(chunk, dtype=np.int16).reshape(-1, 1)

    async def speak_text(self, text: str, chunk_delay: float = 0.15) -> None:
        """Simulate streaming text by sending it word-by-word."""
        if self._ws is None:
            raise RuntimeError("Call connect() first")

        words = text.split()
        for i, word in enumerate(words):
            payload = word + (" " if i < len(words) - 1 else "")
            await self._ws.send(
                json.dumps(
                    {"type": "text", "payload": payload, "seq": self._next_seq()},
                    ensure_ascii=False,
                )
            )
            await asyncio.sleep(chunk_delay)

        await self._ws.send(json.dumps({"type": "flush", "seq": self._next_seq()}))

    async def stop(self, reason: str = "interrupt") -> None:
        if self._ws is None:
            return
        await self._ws.send(
            json.dumps({"type": "stop", "reason": reason, "seq": self._next_seq()})
        )
        async with self._lock:
            self._audio_buffer = bytearray()

    async def run(self, text: str) -> None:
        await self.connect()
        assert self._ws is not None
        try:
            receive_task = asyncio.create_task(self._receive_loop())
            await self.speak_text(text)
            await receive_task
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            await self.close()

    async def _receive_loop(self) -> None:
        assert self._ws is not None
        finished_events = {"stopped", "finished"}
        while True:
            try:
                raw = await self._ws.recv()
            except websockets.exceptions.ConnectionClosed:
                break

            msg = json.loads(raw)
            msg_type = msg.get("type")

            if msg_type == "audio":
                pcm = base64.b64decode(msg["data"])
                async with self._lock:
                    self._audio_buffer.extend(pcm)

            elif msg_type == "status":
                event = msg.get("event")
                if event in finished_events:
                    # Wait for the audio buffer to drain before exiting.
                    await self._drain()
                    break
                print(f"[status] {event} seq={msg.get('seq')}")

            elif msg_type == "error":
                print(f"[error] {msg.get('message')}", file=sys.stderr)

    async def _drain(self) -> None:
        """Wait until the local audio buffer has been played."""
        while True:
            async with self._lock:
                if len(self._audio_buffer) == 0:
                    break
            await asyncio.sleep(0.05)
        # Give sounddevice a little extra time to finish its current block.
        await asyncio.sleep(0.2)

    async def close(self) -> None:
        if self._stream is not None:
            self._stream.stop()
            self._stream.close()
            self._stream = None
        if self._ws is not None:
            await self._ws.close()
            self._ws = None


def main() -> None:
    parser = argparse.ArgumentParser(description="Voice TTS WebSocket client")
    parser.add_argument("--uri", default="ws://localhost:8765/ws")
    parser.add_argument("--voice-ref", default=None)
    parser.add_argument("--language", default="ru")
    parser.add_argument("--speed", type=float, default=1.0)
    parser.add_argument("--emotion", default="neutral")
    parser.add_argument("--sample-rate", type=int, default=24_000)
    parser.add_argument("text", nargs="*", default=["Привет. Это тестовая фраза."])
    args = parser.parse_args()

    text = " ".join(args.text)
    client = VoiceTTSClient(
        uri=args.uri,
        voice_ref=args.voice_ref,
        language=args.language,
        speed=args.speed,
        emotion=args.emotion,
        sample_rate=args.sample_rate,
    )

    try:
        asyncio.run(client.run(text))
    except KeyboardInterrupt:
        print("Interrupted")


if __name__ == "__main__":
    main()