#!/usr/bin/env python3
"""Combined HTTP + WebSocket proxy for Qwen3-TTS on a single port.

Routes on one port (default 8091):
  POST /v1/audio/speech          → REST proxy to internal vllm-omni
  GET  /v1/audio/speech/stream   → WebSocket (PCM-paced, REST-sourced)
  GET  /health                   → health proxy to internal vllm-omni

Why REST-sourced WebSocket?
  The vllm-omni streaming codec drops the onset of short utterances.
  The batch REST endpoint (/v1/audio/speech) uses full decode and renders
  every word cleanly. This proxy keeps the client-facing WS protocol intact
  while sourcing audio sentence-by-sentence from REST, pacing the result
  into one continuous real-time PCM stream via pcm_pacer.

Why instructions is NOT forwarded:
  With a cloned ICL voice (ref_audio) the voice character comes from the
  reference; an extra style prompt makes the Talker emit EOS far too early
  and clips short words ("Hallo." → 0.16 s vs 0.40 s without).
"""
import argparse
import asyncio
import datetime as dt
import json
import math
import re

import aiohttp
import websockets
from aiohttp import web

# ---------------------------------------------------------------------------
SAMPLE_RATE      = 24_000
BYTES_PER_SAMPLE = 2
BYTES_PER_SEC    = SAMPLE_RATE * BYTES_PER_SAMPLE   # 48 000

SEND_CHUNK   = 960                                   # 20 ms at 48 kHz
FRAME_SECS   = SEND_CHUNK / BYTES_PER_SEC            # 0.020 s

LEADING_MS      = 300
LEADING_FRAMES  = math.ceil(LEADING_MS  / 1000 / FRAME_SECS)
TRAILING_MS     = 2_000
TRAILING_FRAMES = math.ceil(TRAILING_MS / 1000 / FRAME_SECS)

# Per-sentence synthesis timeout. A healthy sentence renders in <2 s.
SYNTH_TIMEOUT_S      = 30
# Direct REST proxy timeout — longer for potentially bigger inputs.
REST_PROXY_TIMEOUT_S = 120

SILENCE_CHUNK = bytes(SEND_CHUNK)

SPLIT_RE         = re.compile(r"(?<=[.!?])\s+|(?<=[。！？])")
MIN_SENTENCE_LEN = 2

# Fields forwarded to the internal REST synthesis request.
# `instructions` deliberately omitted — see module docstring.
REST_PASSTHROUGH = ("model", "voice", "task_type", "language",
                    "seed", "speaker_embedding")
# ---------------------------------------------------------------------------

_internal_base_url: str = ""   # e.g. http://localhost:8095
_internal_ws_url:   str = ""   # e.g. ws://localhost:8095/v1/audio/speech/stream


def now() -> str:
    return dt.datetime.now().isoformat(timespec="seconds")


class SentenceSplitter:
    """Incremental splitter: buffers text, yields complete sentences."""

    def __init__(self) -> None:
        self.buf = ""

    def add(self, text: str) -> list[str]:
        self.buf += text
        out: list[str] = []
        while True:
            m = SPLIT_RE.search(self.buf)
            if not m:
                break
            idx  = m.end()
            head = self.buf[:idx].strip()
            if len(head) < MIN_SENTENCE_LEN:
                break
            self.buf = self.buf[idx:]
            if head:
                out.append(head)
        return out

    def flush(self) -> list[str]:
        s = self.buf.strip()
        self.buf = ""
        return [s] if s else []


async def synth_rest(http: aiohttp.ClientSession, cfg: dict, text: str) -> bytes:
    """Synthesise one sentence via the internal batch REST endpoint → raw PCM."""
    payload = {"input": text, "response_format": "pcm"}
    for k in REST_PASSTHROUGH:
        if cfg.get(k) is not None:
            payload[k] = cfg[k]
    async with http.post(f"{_internal_base_url}/v1/audio/speech",
                         json=payload) as resp:
        resp.raise_for_status()
        return await resp.read()


# ── REST proxy ───────────────────────────────────────────────────────────────

async def handle_rest(request: web.Request) -> web.StreamResponse:
    """Forward POST /v1/audio/speech to the internal vllm-omni REST API."""
    body    = await request.read()
    ct      = request.headers.get("Content-Type", "application/json")
    timeout = aiohttp.ClientTimeout(total=REST_PROXY_TIMEOUT_S)
    async with aiohttp.ClientSession(timeout=timeout) as http:
        async with http.post(
            f"{_internal_base_url}/v1/audio/speech",
            data=body,
            headers={"Content-Type": ct},
        ) as resp:
            out_ct   = resp.headers.get("Content-Type", "application/octet-stream")
            response = web.StreamResponse(
                status=resp.status,
                headers={"Content-Type": out_ct},
            )
            await response.prepare(request)
            async for chunk in resp.content.iter_chunked(65536):
                await response.write(chunk)
            await response.write_eof()
            return response


async def handle_health(request: web.Request) -> web.Response:
    """Proxy GET /health to internal vllm-omni."""
    try:
        async with aiohttp.ClientSession() as http:
            async with http.get(f"{_internal_base_url}/health") as resp:
                return web.Response(status=resp.status, text=await resp.text())
    except Exception:
        return web.Response(status=503, text="unavailable")


# ── WebSocket PCM session ────────────────────────────────────────────────────

async def proxy_session_pcm(client_ws: web.WebSocketResponse,
                             conn_id: str, cfg: dict) -> None:
    """REST-sourced, paced, single continuous PCM stream per session."""
    pcm_queue:      asyncio.Queue = asyncio.Queue()
    sentence_queue: asyncio.Queue = asyncio.Queue()
    splitter = SentenceSplitter()
    started  = {"audio": False}

    def enqueue_pcm(data: bytes) -> None:
        for i in range(0, len(data), SEND_CHUNK):
            piece = data[i:i + SEND_CHUNK]
            if len(piece) < SEND_CHUNK:
                piece = piece + bytes(SEND_CHUNK - len(piece))
            pcm_queue.put_nowait(piece)

    async def pcm_pacer() -> bool:
        """Drain pcm_queue at 48 kHz. Returns True = fully drained, False = disconnected."""
        loop = asyncio.get_event_loop()
        t0   = loop.time()
        sent = 0
        try:
            while True:
                target = t0 + sent / BYTES_PER_SEC
                delay  = target - loop.time()
                if delay > 0.001:
                    await asyncio.sleep(delay)
                try:
                    chunk = pcm_queue.get_nowait()
                    if chunk is None:
                        break
                except asyncio.QueueEmpty:
                    chunk = SILENCE_CHUNK
                await client_ws.send_bytes(chunk)
                sent += len(chunk)
        except Exception as exc:
            print(f"[{now()}] {conn_id} PACER send failed after {sent} B "
                  f"({sent/BYTES_PER_SEC:.2f}s): {exc!r}", flush=True)
            return False
        print(f"[{now()}] {conn_id} PACER done, {sent} B "
              f"({sent/BYTES_PER_SEC:.2f}s)", flush=True)
        return True

    pacer_task = asyncio.create_task(pcm_pacer())

    async def synth_worker() -> None:
        timeout = aiohttp.ClientTimeout(total=SYNTH_TIMEOUT_S)
        async with aiohttp.ClientSession(timeout=timeout) as http:
            while True:
                text = await sentence_queue.get()
                if text is None:
                    break
                try:
                    pcm = await synth_rest(http, cfg, text)
                except Exception as exc:
                    print(f"[{now()}] {conn_id} REST synth failed for "
                          f"{text!r}: {exc!r}", flush=True)
                    continue
                print(f"[{now()}] {conn_id} REST {text!r} -> {len(pcm)} B "
                      f"({len(pcm)/BYTES_PER_SEC:.2f}s)", flush=True)
                if not started["audio"]:
                    started["audio"] = True
                    await client_ws.send_str(json.dumps(
                        {"type": "audio.start", "format": "pcm",
                         "sample_rate": SAMPLE_RATE}))
                    for _ in range(LEADING_FRAMES):
                        pcm_queue.put_nowait(SILENCE_CHUNK)
                enqueue_pcm(pcm)
        for _ in range(TRAILING_FRAMES):
            pcm_queue.put_nowait(SILENCE_CHUNK)
        pcm_queue.put_nowait(None)

    worker_task = asyncio.create_task(synth_worker())

    async def cleanup() -> None:
        for t in (worker_task, pacer_task):
            if not t.done():
                t.cancel()
        await asyncio.gather(worker_task, pacer_task, return_exceptions=True)

    try:
        async for msg in client_ws:
            if msg.type == aiohttp.WSMsgType.TEXT:
                try:
                    data = json.loads(msg.data)
                    typ  = data.get("type")
                except Exception:
                    continue
                if typ == "input.text":
                    text = data.get("text", "")
                    print(f"[{now()}] {conn_id} C text {text!r}", flush=True)
                    for sent in splitter.add(text):
                        sentence_queue.put_nowait(sent)
                elif typ == "input.done":
                    print(f"[{now()}] {conn_id} C input.done", flush=True)
                    for sent in splitter.flush():
                        sentence_queue.put_nowait(sent)
                    sentence_queue.put_nowait(None)
                    drained_ok = await pacer_task
                    if not worker_task.done():
                        worker_task.cancel()
                    if drained_ok:
                        try:
                            await client_ws.send_str(
                                json.dumps({"type": "session.done"}))
                        except Exception:
                            pass
                    return
            elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.ERROR):
                break
    finally:
        await cleanup()


async def proxy_session_passthrough(client_ws: web.WebSocketResponse,
                                    conn_id: str, first_msg_data: str) -> None:
    """Transparent WS passthrough to streaming upstream (non-PCM clients)."""
    async with websockets.connect(_internal_ws_url, max_size=None) as upstream:
        await upstream.send(first_msg_data)

        async def c2u() -> None:
            async for msg in client_ws:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    await upstream.send(msg.data)
                elif msg.type == aiohttp.WSMsgType.BINARY:
                    await upstream.send(msg.data)
                else:
                    break

        async def u2c() -> None:
            async for m in upstream:
                if isinstance(m, str):
                    await client_ws.send_str(m)
                else:
                    await client_ws.send_bytes(m)

        await asyncio.gather(c2u(), u2c())


async def handle_ws(request: web.Request) -> web.WebSocketResponse:
    conn_id   = f"ws-{id(request):x}"
    client_ws = web.WebSocketResponse(max_msg_size=0)
    await client_ws.prepare(request)
    print(f"[{now()}] {conn_id} open", flush=True)

    # First client message must be session.config.
    first = await client_ws.receive()
    if first.type != aiohttp.WSMsgType.TEXT:
        return client_ws
    try:
        cfg = json.loads(first.data)
    except Exception:
        cfg = {}

    if cfg.get("response_format", "pcm") == "pcm":
        log = dict(cfg)
        if log.get("speaker_embedding"):
            log["speaker_embedding"] = f"<{len(log['speaker_embedding'])} floats>"
        print(f"[{now()}] {conn_id} config "
              f"{json.dumps(log, ensure_ascii=False, sort_keys=True)}", flush=True)
        await proxy_session_pcm(client_ws, conn_id, cfg)
    else:
        print(f"[{now()}] {conn_id} non-PCM -> passthrough", flush=True)
        await proxy_session_passthrough(client_ws, conn_id, first.data)

    return client_ws


# ── main ─────────────────────────────────────────────────────────────────────

async def main() -> None:
    global _internal_base_url, _internal_ws_url

    parser = argparse.ArgumentParser()
    parser.add_argument("--host",     default="0.0.0.0")
    parser.add_argument("--port",     type=int, default=8091)
    parser.add_argument("--upstream", default="ws://localhost:8095/v1/audio/speech/stream")
    parser.add_argument("--rest",     default=None,
                        help="Internal REST base URL (default: derived from --upstream)")
    args = parser.parse_args()

    _internal_ws_url = args.upstream
    if args.rest:
        _internal_base_url = args.rest.rstrip("/")
    else:
        http = args.upstream.replace("ws://", "http://").replace("wss://", "https://")
        _internal_base_url = http.rsplit("/v1/", 1)[0]

    print(f"Proxy on {args.host}:{args.port}  "
          f"(REST: {_internal_base_url}/v1/audio/speech  "
          f"WS: {_internal_ws_url})", flush=True)

    app = web.Application()
    app.router.add_get("/v1/audio/speech/stream", handle_ws)
    app.router.add_post("/v1/audio/speech",        handle_rest)
    app.router.add_get("/health",                  handle_health)

    runner = web.AppRunner(app, access_log=None)
    await runner.setup()
    site = web.TCPSite(runner, args.host, args.port)
    await site.start()
    print(f"Listening on http+ws://{args.host}:{args.port}", flush=True)
    await asyncio.Future()


if __name__ == "__main__":
    asyncio.run(main())