import asyncio
import json
import logging
import uuid

import websockets

from protocols import EventType, MsgType, full_client_request, receive_message

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def main() -> None:
    headers = {
        "X-Api-Key": "your_api_key",
        "X-Api-Resource-Id": "seed-tts-2.0",
        "X-Api-Request-Id": str(uuid.uuid4()),
        "X-Control-Require-Usage-Tokens-Return": "*"
    }

    body = {
        "req_params": {
            "speaker": "zh_female_gaolengyujie_uranus_bigtts",
            "text": "你好",
            "audio_params": {
                "format": "mp3",
                "sample_rate": 24000,
            }
        }
    }

    websocket = await websockets.connect(
        "wss://openspeech.bytedance.com/api/v3/tts/unidirectional/stream",
        additional_headers=headers,
        max_size=10 * 1024 * 1024,
    )
    logger.info(
        f"Connected to WebSocket server, Logid: {websocket.response.headers.get('x-tt-logid', '')}"
    )

    try:
        await full_client_request(websocket, json.dumps(body).encode())

        audio_data = bytearray()
        while True:
            msg = await receive_message(websocket)

            if msg.type == MsgType.FullServerResponse and msg.event == EventType.SessionFinished:
                break
            if msg.type == MsgType.AudioOnlyServer and msg.payload:
                audio_data.extend(msg.payload)
            elif msg.type == MsgType.Error:
                raise RuntimeError(f"TTS conversion failed: {msg}")

        if not audio_data:
            raise RuntimeError("No audio data received")

        with open("unidirectional_stream.mp3", "wb") as f:
            f.write(audio_data)
        logger.info(f"Audio received: {len(audio_data)}, saved to unidirectional_stream.mp3")
    finally:
        await websocket.close()
        logger.info("Connection closed")

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

{
  "Logid": "20260616211035105B13415E264F69E61D",
  "Received": [
    {
      "MsgType": "FullServerResponse",
      "EventType": "TTSSentenceStart",
      "SessionId": "76ad2906-71c4-4b0b-8655-a77d2718b078",
      "Payload": {
        "phonemes": [],
        "text": "",
        "words": []
      }
    },
    {
      "MsgType": "AudioOnlyServer",
      "EventType": "TTSResponse",
      "SessionId": "76ad2906-71c4-4b0b-8655-a77d2718b078",
      "PayloadSize": 2541
    },
    {
      "MsgType": "FullServerResponse",
      "EventType": "TTSSentenceEnd",
      "SessionId": "76ad2906-71c4-4b0b-8655-a77d2718b078",
      "Payload": {
        "phonemes": [],
        "text": "你好",
        "words": []
      }
    },
    {
      "MsgType": "FullServerResponse",
      "EventType": "TTSSubtitle",
      "SessionId": "76ad2906-71c4-4b0b-8655-a77d2718b078",
      "Payload": {
        "phonemes": [],
        "text": "你好",
        "words": [
          {
            "confidence": 0.6805309,
            "endTime": 0.285,
            "startTime": 0.155,
            "word": "你"
          },
          {
            "confidence": 0.9384748,
            "endTime": 0.735,
            "startTime": 0.285,
            "word": "好"
          }
        ]
      }
    },
    {
      "MsgType": "FullServerResponse",
      "EventType": "SessionFinished",
      "SessionId": "76ad2906-71c4-4b0b-8655-a77d2718b078",
      "Payload": {
        "usage": {
          "text_words": 2
        }
      }
    }
  ]
}