import asyncio
import json
import wave

import websockets


GATEWAY_URL = "ws://127.0.0.1:8000/ws/audio"

AUDIO_FILE = (
    "/home/ubuntu/voice-assistant/audios/chafik_normalize.wav"
)

CHUNK_FRAMES = 1600


async def main() -> None:
    print("========================================")
    print("STT GATEWAY SMOKE TEST")
    print("========================================")
    print(f"Gateway : {GATEWAY_URL}")
    print(f"Audio   : {AUDIO_FILE}")
    print(f"Chunk   : {CHUNK_FRAMES} frames")
    print()

    async with websockets.connect(GATEWAY_URL) as websocket:

        # -------------------------------------------------------------
        # 1. session.start
        # -------------------------------------------------------------
        await websocket.send(
            json.dumps(
                {
                    "type": "session.start",
                    "audio": {
                        "encoding": "pcm_s16le",
                        "sample_rate": 16000,
                        "channels": 1,
                    },
                }
            )
        )

        response = json.loads(
            await websocket.recv()
        )

        print(
            f"[RECV] {response['type']}"
        )

        if response["type"] != "session.ready":
            raise RuntimeError(
                f"Expected session.ready, got {response}"
            )

        session_id = response["session_id"]

        # -------------------------------------------------------------
        # 2. audio.start
        # -------------------------------------------------------------
        await websocket.send(
            json.dumps(
                {
                    "type": "audio.start",
                }
            )
        )

        response = json.loads(
            await websocket.recv()
        )

        print(
            f"[RECV] {response['type']}"
        )

        if response["type"] != "turn.started":
            raise RuntimeError(
                f"Expected turn.started, got {response}"
            )

        turn_id = response["turn_id"]

        print(
            f"Session : {session_id}"
        )
        print(
            f"Turn    : {turn_id}"
        )
        print()

        # -------------------------------------------------------------
        # 3. Start a receiver task.
        #
        #    This is necessary because partial/final events can arrive
        #    while we are still sending audio.
        # -------------------------------------------------------------
        received_events: list[dict] = []

        async def receive_events() -> None:
            while True:
                message = await websocket.recv()
                event = json.loads(message)

                received_events.append(event)

                event_type = event.get("type")

                if event_type == "stt.partial":
                    print(
                        f"[PARTIAL] {event.get('text', '')}"
                    )

                elif event_type == "stt.final":
                    print(
                        f"[FINAL  ] {event.get('text', '')}"
                    )

                elif event_type == "turn.ended":
                    print(
                        "[RECV] turn.ended"
                    )
                    break

                elif event_type == "error":
                    print(
                        "[ERROR]",
                        event,
                    )

                    break

                else:
                    print(
                        f"[RECV] {event_type}"
                    )

        receiver_task = asyncio.create_task(
            receive_events()
        )

        # -------------------------------------------------------------
        # 4. Stream WAV as:
        #
        #       audio.chunk JSON
        #       binary PCM
        #
        #    The sequence number belongs to the logical chunk.
        # -------------------------------------------------------------
        with wave.open(AUDIO_FILE, "rb") as wav:
            sequence_number = 0

            while True:
                audio = wav.readframes(CHUNK_FRAMES)

                if not audio:
                    break

                await websocket.send(
                    json.dumps(
                        {
                            "type": "audio.chunk",
                            "turn_id": turn_id,
                            "sequence_number": sequence_number,
                        }
                    )
                )

                await websocket.send(audio)

                sequence_number += 1

        print()
        print(
            f"Audio sent: {sequence_number} chunks"
        )

        # -------------------------------------------------------------
        # 5. audio.stop
        # -------------------------------------------------------------
        await websocket.send(
            json.dumps(
                {
                    "type": "audio.stop",
                    "turn_id": turn_id,
                }
            )
        )

        # -------------------------------------------------------------
        # 6. Wait for:
        #
        #       stt.final
        #       turn.ended
        #
        #    The receiver task handles both.
        # -------------------------------------------------------------
        await receiver_task

        print()
        print("========================================")
        print("RESULT")
        print("========================================")

        partials = [
            event
            for event in received_events
            if event.get("type") == "stt.partial"
        ]

        finals = [
            event
            for event in received_events
            if event.get("type") == "stt.final"
        ]

        turn_ended = [
            event
            for event in received_events
            if event.get("type") == "turn.ended"
        ]

        print(
            f"Partial events : {len(partials)}"
        )
        print(
            f"Final events   : {len(finals)}"
        )
        print(
            f"Turn ended     : {len(turn_ended)}"
        )

        if finals:
            print()
            print("Final transcript:")
            print(
                finals[-1].get("text", "")
            )

        print()
        print("STT Gateway smoke test finished.")


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