import gradio as gr
import os
from coqui_tts import run_audio_generation_v1

def process_audio(input_text, speaker_audio, speaker_name, option_selected):
    try:
        # Ensure necessary directories exist
        os.makedirs("./tmp/audio/input_src/", exist_ok=True)
        os.makedirs("audio", exist_ok=True)
        
        # Save speaker audio to the required location
        speaker_audio_path = "./tmp/audio/input_src/0.wav"
        with open(speaker_audio_path, "wb") as f:
            f.write(speaker_audio)
        
        # Check selected option and execute corresponding function
        if option_selected == "Xtts_v2":
            # Generate TTS audio using run_audio_generation_v1
            run_audio_generation_v1(input_text)
        else:
            return f"The option '{option_selected}' is not implemented yet."
        
        # Save the output audio under the speaker's name
        speaker_output_path = f"audio/{speaker_name}.wav"
        os.rename("audio/output.wav", speaker_output_path)
        
        return speaker_output_path
    
    except Exception as e:
        return str(e)

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Audio Cloning and Text-to-Speech")

    with gr.Row():
        text_input = gr.Textbox(label="Input Text", placeholder="Enter your text here.")
    
    with gr.Row():
        speaker_audio = gr.Audio(label="Speaker Audio (to be cloned)",type='filepath', format='wav')
        speaker_name = gr.Textbox(label="Speaker Name", placeholder="Enter the speaker's name.")

    option_selected = gr.Dropdown(choices=["Xtts_v2", "metaVoice(not working at the moment)", "more"], label="Select an Option")

    submit_btn = gr.Button("Submit")

    output_audio = gr.Audio(label="Generated Audio Output")

    submit_btn.click(
        fn=process_audio,
        inputs=[text_input, speaker_audio, speaker_name, option_selected],
        outputs=output_audio,
    )

# Launch the Gradio app
demo.launch()