import os
import torch
import requests
from transformers import AutoModelForCausalLM
from audiotokenizer import AudioTokenizer
from tqdm import tqdm

#define model storage directory
MODEL_DIR = os.path.expanduser("~/.yarngpt/models")
os.makedirs(MODEL_DIR, exist_ok=True)

#define file paths
CONFIG_PATH = os.path.join(MODEL_DIR, "wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml")
MODEL_PATH = os.path.join(MODEL_DIR, "wavtokenizer_large_speech_320_24k.ckpt")

#urls from Hugging Face
CONFIG_URL = "https://huggingface.co/novateur/WavTokenizer-medium-speech-75token/resolve/main/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
MODEL_URL = "https://huggingface.co/novateur/WavTokenizer-large-speech-75token/resolve/main/wavtokenizer_large_speech_320_24k.ckpt"

def download_file(url, dest_path):
    """Downloads a file with a progress bar if it doesn't already exist."""
    if os.path.exists(dest_path):
        print(f"{dest_path} already exists. Skipping download.")
        return

    print(f"Downloading {url} to {dest_path}...")

    response = requests.get(url, stream=True)
    total_size = int(response.headers.get('content-length', 0))

    with open(dest_path, "wb") as f, tqdm(
        total=total_size, unit="B", unit_scale=True, desc=os.path.basename(dest_path)
    ) as progress_bar:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
            progress_bar.update(len(chunk))

    print("Download complete.")

#ensure model files are available
download_file(CONFIG_URL, CONFIG_PATH)
download_file(MODEL_URL, MODEL_PATH)

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

def load_model_and_tokenizer():
    """Loads the YarnGPT model and tokenizer."""
    hf_path = "saheedniyi/YarnGPT"

    #initialize tokenizer
    audio_tokenizer = AudioTokenizer(hf_path, MODEL_PATH, CONFIG_PATH)

    #load model using Hugging Face's caching system
    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):
    """Generates speech audio from input text using the selected speaker."""
    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
