import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from audiotokenizer import AudioTokenizer

#list of available speakers
AVAILABLE_SPEAKERS = [
    "zainab", "idera", "regina", "chinenye", "joke", "remi",
    "jude", "tayo", "umar", "osagie", "onye", "emma"
]

def load_model_and_tokenizer():
    hf_path = "saheedniyi/YarnGPT"
    wav_tokenizer_config_path = "./models/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
    wav_tokenizer_model_path = "./models/wavtokenizer_large_speech_320_24k.ckpt"
    
    audio_tokenizer = AudioTokenizer(hf_path, wav_tokenizer_model_path, wav_tokenizer_config_path)
    model = AutoModelForCausalLM.from_pretrained(hf_path, torch_dtype="auto")
    model = model.to(audio_tokenizer.device)
    return model, audio_tokenizer

def generate_speech(text, speaker="idera", temperature=0.1, repetition_penalty=1.1, max_length=4000):
    if speaker not in AVAILABLE_SPEAKERS:
        raise ValueError(f"Speaker must be one of: {', '.join(AVAILABLE_SPEAKERS)}")
    
    model, audio_tokenizer = load_model_and_tokenizer()
    prompt = audio_tokenizer.create_prompt(text, speaker)
    input_ids = audio_tokenizer.tokenize_prompt(prompt)
    output = model.generate(
        input_ids=input_ids,
        temperature=temperature,
        repetition_penalty=repetition_penalty,
        max_length=max_length
    )
    codes = audio_tokenizer.get_codes(output)
    audio = audio_tokenizer.get_audio(codes)
    return audio
