import gradio as gr
import base64
import numpy as np
from scipy.io import wavfile
from voice_processing import parallel_tts_wrapper, get_model_names, voice_mapping
from io import BytesIO
import asyncio
import tempfile

# Define get_models function
def get_models():
    return get_model_names()

# Define get_voices function
def get_voices():
    return list(voice_mapping.keys())

async def convert_tts(model_name, tts_text, selected_voice, slang_rate, use_uploaded_voice, voice_upload):
    try:
        edge_tts_voice = voice_mapping.get(selected_voice)
        if not edge_tts_voice:
            return {"error": f"Invalid voice '{selected_voice}'."}, None

        voice_upload_file = None
        if use_uploaded_voice and voice_upload is not None:
            with open(voice_upload.name, 'rb') as f:
                voice_upload_file = f.read()

        task = (
            model_name, tts_text, edge_tts_voice, slang_rate, use_uploaded_voice, voice_upload_file
        )
        
        results = await parallel_tts_wrapper([task])
        info, _, (tgt_sr, audio_output) = results[0]

        # Save audio to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio_file:
            if isinstance(audio_output, np.ndarray):
                wavfile.write(temp_audio_file.name, tgt_sr, audio_output)
            else:
                temp_audio_file.write(audio_output)

        return {"info": info}, temp_audio_file.name

    except Exception as e:
        print(f"Error in convert_tts: {e}")
        import traceback
        traceback.print_exc()
        return {"error": str(e)}, None

# Update the Gradio interface
iface = gr.Interface(
    fn=convert_tts,
    inputs=[
        gr.Dropdown(choices=get_models(), label="Model", interactive=True),
        gr.Textbox(label="Text", placeholder="Enter text here"),
        gr.Dropdown(choices=get_voices(), label="Voice", interactive=True),
        gr.Slider(minimum=0, maximum=1, step=0.01, label="Slang Rate"),
        gr.Checkbox(label="Use Uploaded Voice"),
        gr.File(label="Voice File")
    ],
    outputs=[
        gr.JSON(label="Info"),
        gr.Audio(label="Generated Audio")
    ],
    title="Text-to-Speech Conversion"
)

# Launch the interface
if __name__ == "__main__":
    iface.launch(debug=True)

def cleanup_temp_files():
    temp_dir = tempfile.gettempdir()
    for filename in os.listdir(temp_dir):
        if filename.endswith(".wav"):
            file_path = os.path.join(temp_dir, filename)
            try:
                os.remove(file_path)
            except Exception as e:
                print(f"Error removing temporary file {file_path}: {e}")