import asyncio
import websockets
import os
from dotenv import load_dotenv

load_dotenv()
api_key = os.environ.get('RIME_API_KEY')
if not api_key:
    raise ValueError("RIME_API_KEY environment variable is not set")

FILE_PATH = "arcana_ws.wav"

class RimeClient:
    def __init__(self, speaker, api_key):
        self.url = f"wss://users.rime.ai/ws?speaker={speaker}&modelId=arcana&audioFormat=wav"
        self.auth_headers = {
            "Authorization": f"Bearer {api_key}"
        }
        self.audio_data = b''

    async def send_tokens(self, websocket, message):
        for token in message:
            await websocket.send(token)

    async def handle_audio(self, websocket):
        while True:
            try:
                audio = await websocket.recv()
            except websockets.exceptions.ConnectionClosedOK:
                break
            self.audio_data += audio

    async def run(self, message):
        async with websockets.connect(self.url, additional_headers=self.auth_headers) as websocket:
            await asyncio.gather(
                self.send_tokens(websocket, message),
                self.handle_audio(websocket),
            )

    def save_audio(self, file_path):
        with open(file_path, 'wb') as f:
            f.write(self.audio_data)

message = [
    "This ",
    "is ",
    "a ",
    "test ",
    "of ",
    "the ",
    "arcana ",
    "model ",
    "using ",
    "websockets ",
    "and ",
    "python.",
    "<EOS>",
]

client = RimeClient("astra", api_key=api_key)
asyncio.run(client.run(message))

print(f"Saving audio to {FILE_PATH}")
client.save_audio(FILE_PATH)