from __future__ import annotations

import argparse
import asyncio
import json

from aiortc import RTCPeerConnection, RTCSessionDescription
from aiortc.contrib.media import MediaPlayer
import httpx


async def run(server_url: str, audio_path: str, timeout: float) -> None:
    peer = RTCPeerConnection()
    channel = peer.createDataChannel("transcripts")
    player = MediaPlayer(audio_path)
    if player.audio is None:
        raise RuntimeError("Input file does not contain an audio track")
    peer.addTrack(player.audio)
    final_received = asyncio.Event()
    session_id = None

    @channel.on("message")
    def on_message(message):
        event = json.loads(message)
        print(json.dumps(event, ensure_ascii=False), flush=True)
        if event.get("type") == "transcript.final":
            final_received.set()

    try:
        offer = await peer.createOffer()
        await peer.setLocalDescription(offer)
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(
                f"{server_url}/api/webrtc/offer",
                json={
                    "sdp": peer.localDescription.sdp,
                    "type": peer.localDescription.type,
                },
            )
            response.raise_for_status()
            answer = response.json()
            session_id = answer["sessionId"]
            await peer.setRemoteDescription(
                RTCSessionDescription(
                    sdp=answer["sdp"],
                    type=answer["type"],
                )
            )
            await asyncio.wait_for(final_received.wait(), timeout=timeout)
            await client.post(
                f"{server_url}/api/webrtc/session/{session_id}/close"
            )
    finally:
        await peer.close()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("audio")
    parser.add_argument(
        "--server",
        default="http://127.0.0.1:8090",
    )
    parser.add_argument("--timeout", type=float, default=60.0)
    args = parser.parse_args()
    asyncio.run(run(args.server.rstrip("/"), args.audio, args.timeout))


if __name__ == "__main__":
    main()