import asyncio
import websockets
import numpy as np
import pyaudio
import signal
import sys

RATE = 16000
CHUNK = 4096
CHANNELS = 1

WS_URL = "wss://fredyhoundayi-relay.hf.space/ws"

running = True

# -------- Stop propre CTRL+C --------
def stop(sig, frame):
    global running
    running = False
    print("\nStopping...")
signal.signal(signal.SIGINT, stop)


# -------- Audio Capture --------
async def send_audio(ws):
    p = pyaudio.PyAudio()

    stream = p.open(
        format=pyaudio.paInt16,
        channels=CHANNELS,
        rate=RATE,
        input=True,
        frames_per_buffer=CHUNK
    )

    print("🎙️ Micro streaming started")

    try:
        while running:
            data = stream.read(CHUNK, exception_on_overflow=False)

            # int16 → float32 normalisé [-1,1]
            audio = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0

            await ws.send(audio.tobytes())
            await asyncio.sleep(0)  # yield event loop

    finally:
        stream.stop_stream()
        stream.close()
        p.terminate()


# -------- Réception texte --------
async def receive_text(ws):
    try:
        async for message in ws:
            print("📝", message)
    except websockets.ConnectionClosed:
        print("Connection closed by server")


# -------- Main --------
async def main():
    print("Connecting to Relay server...")

    async with websockets.connect(
        WS_URL,
        ping_interval=20,
        ping_timeout=20,
        max_size=10_000_000
    ) as ws:

        send_task = asyncio.create_task(send_audio(ws))
        recv_task = asyncio.create_task(receive_text(ws))

        done, pending = await asyncio.wait(
            [send_task, recv_task],
            return_when=asyncio.FIRST_COMPLETED
        )

        for task in pending:
            task.cancel()


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