import asyncio
import websockets
import json

AUDIO_FILE = "./data/khabar_1_5.wav"
WS_URL = "ws://127.0.0.1:3002/asr"

async def send_audio():
    async with websockets.connect(WS_URL, max_size=None) as ws:
        print("Connected to ASR WebSocket")

        # --- Send audio in binary chunks ---
        with open(AUDIO_FILE, "rb") as f:
            chunk_size = 100000
            while chunk := f.read(chunk_size):
                await ws.send(chunk)

        # --- Tell server that audio stream is finished ---
        await ws.send(json.dumps({"event": "end"}))

        # --- Receive partial + final transcripts ---
        try:
            while True:
                msg = await ws.recv()

                # Triton code may send JSON or text
                try:
                    data = json.loads(msg)
                except Exception:
                    print("Non-JSON frame:", msg)
                    continue

                print("Received:", data)

                # Optional: stop when server sends is_final
                if data.get("is_final"):
                    print("Final:", data["text"])

        except websockets.exceptions.ConnectionClosed:
            print("WebSocket closed by server")

asyncio.run(send_audio())