#!/usr/bin/env python3
"""Logging WebSocket proxy for Qwen3-TTS streaming calls."""
import argparse
import asyncio
import datetime as dt
import json
import websockets


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


def shorten(value, limit=500):
    text = value if isinstance(value, str) else repr(value)
    return text if len(text) <= limit else text[:limit] + "..."


async def proxy(client_ws, upstream_url):
    conn_id = f"ws-{id(client_ws):x}"
    print(f"[{now()}] {conn_id} open -> {upstream_url}", flush=True)
    async with websockets.connect(upstream_url, max_size=None) as upstream_ws:
        async def client_to_upstream():
            async for msg in client_ws:
                if isinstance(msg, str):
                    try:
                        data = json.loads(msg)
                        typ = data.get("type")
                        if typ == "session.config":
                            log_data = dict(data)
                            if "speaker_embedding" in log_data:
                                emb = log_data["speaker_embedding"] or []
                                log_data["speaker_embedding"] = f"<{len(emb)} floats>"
                            print(f"[{now()}] {conn_id} C->S config {json.dumps(log_data, ensure_ascii=False, sort_keys=True)}", flush=True)
                        elif typ == "input.text":
                            print(f"[{now()}] {conn_id} C->S text {data.get('text')!r}", flush=True)
                        else:
                            print(f"[{now()}] {conn_id} C->S {shorten(msg)}", flush=True)
                    except Exception:
                        print(f"[{now()}] {conn_id} C->S {shorten(msg)}", flush=True)
                else:
                    print(f"[{now()}] {conn_id} C->S binary {len(msg)} bytes", flush=True)
                await upstream_ws.send(msg)

        async def upstream_to_client():
            async for msg in upstream_ws:
                if isinstance(msg, str):
                    try:
                        data = json.loads(msg)
                        print(f"[{now()}] {conn_id} S->C {data}", flush=True)
                    except Exception:
                        print(f"[{now()}] {conn_id} S->C {shorten(msg)}", flush=True)
                else:
                    print(f"[{now()}] {conn_id} S->C binary {len(msg)} bytes", flush=True)
                await client_ws.send(msg)

        await asyncio.gather(client_to_upstream(), upstream_to_client())


async def handler(client_ws):
    try:
        await proxy(client_ws, handler.upstream_url)
    except Exception as exc:
        print(f"[{now()}] proxy error: {exc}", flush=True)
        try:
            await client_ws.close()
        except Exception:
            pass


async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8094)
    parser.add_argument("--upstream", default="ws://localhost:8091/v1/audio/speech/stream")
    args = parser.parse_args()
    handler.upstream_url = args.upstream
    print(f"WS log proxy: ws://{args.host}:{args.port}/v1/audio/speech/stream -> {args.upstream}", flush=True)
    async with websockets.serve(handler, args.host, args.port, max_size=None):
        await asyncio.Future()


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