import json
import torch
from TTS.api import TTS
import gradio as gr

# Load the language dictionary from the JSON file
with open("language_dictionary.json", "r") as f:
    language_dict = json.load(f)

# Get device
device = "cuda" if torch.cuda.is_available() else "cpu"

# Function to generate audio
def generate_audio(text, language):
    model_name = f"tts_models/{language}/fairseq/vits"
    tts = TTS(model_name=model_name, progress_bar=False).to(device)
    output_path = "output_file.wav"
    tts.tts_to_file(text=text, file_path=output_path)
    return output_path

# Create a list of language names for the dropdown
languages = list(language_dict.values())
iso_codes = list(language_dict.keys())

# Create Gradio interface
def run_interface():
    def update_iso_code(language):
        # Map selected language to its ISO code
        iso_code = iso_codes[languages.index(language)]
        return iso_code

    with gr.Blocks() as demo:
        with gr.Row():
            text_input = gr.Textbox(label="Enter Text", lines=4, placeholder="Type your text here...")
            language_dropdown = gr.Dropdown(label="Select Language", choices=languages)

        generate_button = gr.Button("Generate Audio")
        output_audio = gr.Audio(label="Generated Audio", type="filepath")

        generate_button.click(
            fn=lambda text, lang: generate_audio(text, update_iso_code(lang)),
            inputs=[text_input, language_dropdown],
            outputs=[output_audio]
        )

    demo.launch()

# Run the Gradio interface
if __name__ == "__main__":
    run_interface()