# tts_utils.py
import os
import time
import logging
import requests
import streamlit as st

logger = logging.getLogger(__name__)

def generate_tts_audio(text, voice="af_heart", speed=1.0):
    RUNPOD_API_TOKEN = os.getenv("RUNPOD_GPU")
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {RUNPOD_API_TOKEN}'
    }
    data_payload = {"input": {"text": text, "voice": voice, "speed": speed}}
    response = requests.post('https://api.runpod.ai/v2/ozz8w092oprwqx/run', headers=headers, json=data_payload)
    if response.status_code != 200:
        raise Exception(f"RunPod API call failed with status {response.status_code}: {response.text}")
    run_id = response.json().get("id")
    status_url = f"https://api.runpod.ai/v2/ozz8w092oprwqx/status/{run_id}"
    st.toast("TTS generation started, please wait...")
    while True:
        time.sleep(5)
        status_response = requests.post(status_url, headers=headers, json=data_payload)
        status_json = status_response.json()
        logger.debug("TTS status: %s", status_json.get("status"))
        if status_json.get("status") == "COMPLETED":
            download_url = status_json.get("output", {}).get("download_url")
            if download_url:
                mp3_response = requests.get(download_url)
                if mp3_response.status_code == 200:
                    return mp3_response.content
                else:
                    raise Exception(f"Failed to download audio: {mp3_response.status_code}")
        elif status_json.get("status") in ["FAILED", "ERROR"]:
            logger.error("TTS generation failed.")
            st.error("TTS generation failed. Please try again later.")
            raise Exception("TTS generation failed.")