#!/usr/bin/env python3
"""
Dictate streaming speech-to-text - Python example.

Authenticates, requests a WebSocket token, streams 16 kHz mono s16le PCM and
prints every text update.

Usage:
    pip install httpx websockets
    DICTATE_EMAIL=you@example.com DICTATE_PASSWORD=... \\
        python transcribe.py speech.raw

``speech.raw`` is 16 kHz, mono, signed 16-bit little-endian PCM. Convert with:
    ffmpeg -i input.wav -ar 16000 -ac 1 -f s16le -acodec pcm_s16le speech.raw
"""

import asyncio
import gzip
import json
import os
import struct
import sys

import httpx
import websockets

BASE_URL = os.environ.get("DICTATE_BASE_URL", "http://localhost:7932")
FRAME_MS = 200
BYTES_PER_MS = 16000 * 2 // 1000  # 16 kHz * 16 bit * mono = 32 bytes/ms


def parse_frame(data: bytes) -> dict:
    """Decode one binary frame: header(4) + sequence(4) + size(4) + gzip JSON."""
    if len(data) < 8:
        return {}
    header_size = (data[0] & 0x0F) * 4 or 4
    message_type = (data[1] >> 4) & 0x0F
    flags = data[1] & 0x0F
    compression = data[2] & 0x0F

    offset = header_size
    if flags in (0b0001, 0b0011):  # a sequence number is present
        offset += 4
    size = struct.unpack(">I", data[offset : offset + 4])[0]
    offset += 4
    payload = data[offset : offset + size]

    if compression == 0b0001:
        try:
            payload = gzip.decompress(payload)
        except OSError:
            pass

    try:
        body = json.loads(payload)
    except (ValueError, TypeError):
        body = {}

    text = ""
    result = body.get("result")
    if isinstance(result, list):
        text = "".join(
            (item.get("payload_msg") or {}).get("text", "") for item in result
        )
    elif isinstance(result, dict):
        text = result.get("text", "") or body.get("text", "")
    else:
        text = body.get("text", "")

    return {
        "message_type": message_type,
        "is_final": flags in (0b0010, 0b0011),
        "text": text,
    }


async def main() -> None:
    if len(sys.argv) < 2:
        print("usage: python transcribe.py <speech.raw>", file=sys.stderr)
        raise SystemExit(1)

    async with httpx.AsyncClient(base_url=BASE_URL, timeout=30) as http:
        login = await http.post(
            "/api/v1/auth/login",
            json={
                "email": os.environ.get("DICTATE_EMAIL"),
                "password": os.environ.get("DICTATE_PASSWORD"),
            },
        )
        login.raise_for_status()
        token = login.json()["access_token"]

        auth = {"Authorization": f"Bearer {token}"}

        grant_resp = await http.post("/api/v1/ws-token", headers=auth)
        grant_resp.raise_for_status()
        grant = grant_resp.json()

    print(f"remaining quota: {grant['remaining_ms']} ms")

    host = BASE_URL.split("://", 1)[1].split(":")[0]
    port = grant.get("port") or 7932
    url = f"ws://{host}:{port}{grant['path']}?token={grant['token']}"

    with open(sys.argv[1], "rb") as handle:
        pcm = handle.read()

    step = BYTES_PER_MS * FRAME_MS
    frames = [pcm[i : i + step] for i in range(0, len(pcm), step)]

    async with websockets.connect(url, max_size=16 * 1024 * 1024) as socket:
        for frame in frames:
            await socket.send(frame)
        print(f"sent {len(frames)} frames")

        # Closing is how the server learns the session ended; it then flushes the
        # final utterance. Do not keep the socket open waiting for one.
        async def closer():
            await asyncio.sleep(2)
            await socket.close()

        asyncio.create_task(closer())

        try:
            async for message in socket:
                if isinstance(message, str):
                    continue  # JSON control message
                frame = parse_frame(message)
                if frame.get("text"):
                    # Each frame carries the whole transcript so far.
                    print(frame["text"])
                if frame.get("is_final"):
                    print("\n--- final ---")
                    break
        except websockets.ConnectionClosed:
            pass


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nstopped")