import asyncio
import json
import wave
from pathlib import Path

import websockets


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

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

OUTPUT_FILE = Path(
    "/tmp/voice_pipeline_response.pcm"
)

CHUNK_FRAMES = 1600


async def receive_messages(
    websocket,
    state: dict,
) -> None:
    while True:
        message = await websocket.recv()

        if isinstance(message, bytes):
            state["tts_chunks"] += 1
            state["tts_bytes"] += len(message)
            state["tts_audio"].extend(message)

            print(
                f"[TTS AUDIO] "
                f"chunk={state['tts_chunks']} "
                f"bytes={len(message)}"
            )

            continue

        event = json.loads(message)
        event_type = event.get("type")

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

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

        elif event_type == "agent.response":
            state["agent_response"] = event.get("text")
            print(
                f"[AGENT  ] "
                f"{state['agent_response']}"
            )

        elif event_type == "turn.ended":
            state["turn_ended"] = True
            print("[RECV] turn.ended")
            return

        elif event_type == "error":
            raise RuntimeError(
                f"Gateway error: {event}"
            )

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


async def send_audio(
    websocket,
    turn_id: int,
) -> int:
    sent_chunks = 0

    with wave.open(str(AUDIO_FILE), "rb") as audio:
        sample_rate = audio.getframerate()
        channels = audio.getnchannels()
        sample_width = audio.getsampwidth()

        if sample_rate != 16000:
            raise RuntimeError(
                f"Expected 16 kHz audio, got {sample_rate} Hz."
            )

        if channels != 1:
            raise RuntimeError(
                f"Expected mono audio, got {channels} channels."
            )

        if sample_width != 2:
            raise RuntimeError(
                f"Expected 16-bit PCM audio, got "
                f"{sample_width * 8}-bit."
            )

        while True:
            pcm = audio.readframes(CHUNK_FRAMES)

            if not pcm:
                break

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

            await websocket.send(pcm)

            sent_chunks += 1

    return sent_chunks


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

    state = {
        "partial_count": 0,
        "final_text": None,
        "agent_response": None,
        "tts_chunks": 0,
        "tts_bytes": 0,
        "tts_audio": bytearray(),
        "turn_ended": False,
    }

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

        message = await websocket.recv()

        if isinstance(message, bytes):
            raise RuntimeError(
                "Expected JSON session.ready."
            )

        session_ready = json.loads(message)

        print(
            f"[RECV] "
            f"{session_ready.get('type')}"
        )

        if session_ready.get("type") != "session.ready":
            raise RuntimeError(
                f"Expected session.ready, got "
                f"{session_ready}"
            )

        session_id = session_ready["session_id"]

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

        message = await websocket.recv()

        if isinstance(message, bytes):
            raise RuntimeError(
                "Expected JSON turn.started."
            )

        turn_started = json.loads(message)

        print(
            f"[RECV] "
            f"{turn_started.get('type')}"
        )

        if turn_started.get("type") != "turn.started":
            raise RuntimeError(
                f"Expected turn.started, got "
                f"{turn_started}"
            )

        turn_id = turn_started["turn_id"]

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

        # ---------------------------------------------------------
        # Run sender and receiver concurrently
        # ---------------------------------------------------------
        receiver_task = asyncio.create_task(
            receive_messages(
                websocket,
                state,
            )
        )

        try:
            sent_chunks = await send_audio(
                websocket,
                turn_id,
            )
        except Exception:
            receiver_task.cancel()
            raise

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

        # ---------------------------------------------------------
        # Tell Gateway that user audio is complete.
        # ---------------------------------------------------------
        await websocket.send(
            json.dumps(
                {
                    "type": "audio.stop",
                    "turn_id": turn_id,
                }
            )
        )

        # ---------------------------------------------------------
        # Wait for the entire pipeline to finish.
        # ---------------------------------------------------------
        await asyncio.wait_for(
            receiver_task,
            timeout=120.0,
        )

        # ---------------------------------------------------------
        # Validate the complete STT -> Agent -> TTS pipeline.
        # ---------------------------------------------------------
        if not state["turn_ended"]:
            raise RuntimeError("Turn did not complete.")

        if not state["final_text"]:
            raise RuntimeError("STT final transcript was not received.")

        if not state["agent_response"]:
            raise RuntimeError("Agent response was not received.")

        if state["tts_chunks"] == 0:
            raise RuntimeError("No TTS audio was received.")

        if state["tts_bytes"] == 0:
            raise RuntimeError("TTS output is empty.")
            

    # -------------------------------------------------------------
    # Save TTS audio
    # -------------------------------------------------------------
    OUTPUT_FILE.write_bytes(
        bytes(state["tts_audio"])
    )

    print()
    print("=" * 40)
    print("RESULT")
    print("=" * 40)
    print(f"Session          : {session_id}")
    print(f"Turn             : {turn_id}")
    print(
        f"Audio sent      : {sent_chunks}"
    )
    print(
        f"Partial events  : {state['partial_count']}"
    )
    print(
        f"Final transcript: {state['final_text']}"
    )
    print(
        f"Agent response  : {state['agent_response']}"
    )
    print(
        f"TTS chunks      : {state['tts_chunks']}"
    )
    print(
        f"TTS bytes       : {state['tts_bytes']}"
    )
    print(
        f"TTS output      : {OUTPUT_FILE}"
    )
    print(
        f"Turn ended      : {state['turn_ended']}"
    )
    print()
    print(
        "Voice pipeline smoke test finished."
    )


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