import asyncio
import websockets
import os
from io import BytesIO
from pydub import AudioSegment
from utils import transcribe_audio, get_response_llm, play_text_to_speech, load_whisper

# Load Whisper Model
model = load_whisper()

# Store connected clients
clients = set()

async def handle_audio_stream(websocket):
    """Handles incoming live audio stream from users."""
    clients.add(websocket)
    try:
        while True:
            audio_chunk = await websocket.recv()
            if not audio_chunk:
                break

            # Convert received bytes to WAV file
            audio = AudioSegment.from_file(BytesIO(audio_chunk))
            audio.export("temp_audio_chunk.wav", format="wav")

            # Transcribe
            text = transcribe_audio(model, "temp_audio_chunk.wav")
            if text:
                print(f"User: {text}")

                # Generate response
                response_llm = get_response_llm(user_question=text)

                # Play response
                play_text_to_speech(text=response_llm)

                # Send response to WebSocket client
                await websocket.send(response_llm)
    except:
        pass
    finally:
        clients.remove(websocket)

async def main():
    async with websockets.serve(handle_audio_stream, "0.0.0.0", 8765):
        await asyncio.Future()  # Keeps the server running

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