import os
import spaces
import gradio as gr
import tempfile
from importlib.resources import files
from silma_tts.api import SilmaTTS


CURR_BASE_DIR = os.getcwd()

print("Loading SILMA TTS model...", flush=True)
silma_tts = SilmaTTS()


@spaces.GPU(duration=120)
def do_inference(ref_audio, ref_text, gen_text):
    if not ref_audio:
        return None

    _, output_wav_path = tempfile.mkstemp(suffix=".wav")

    wav, sr, _ = silma_tts.infer(
        ref_file=ref_audio,
        ref_text=ref_text,
        gen_text=gen_text,
        file_wave=output_wav_path,
    )


    return output_wav_path



custom_css = """
    .gradio-container{
        background-color: unset;
    }
    input, textarea {
        font-family: 'Noto Naskh Arabic', 'Arial', sans-serif !important;
    }

    button.secondary:hover{
        background-color:steelblue !important;
        color:white;
    }
   
    #tool-header{
        padding:0px 0px 0px 0px !important;
        font-family: sans-serif;
    }
    #tool-header h1{
        display: flex;
        align-items: center;
    }
    #tool-header img{
        width: 80px;
        display: inline-block;
        margin-right: 10px;
        border-radius: 5px;
    }

    .gradio-style a{
        padding: 0px !important;
    }
"""

with gr.Blocks(css=custom_css) as demo:

    gr.HTML(
        f"""<h1><img src='/gradio_api/file={CURR_BASE_DIR}/images/silma-logo.png'/>SILMA TTS Demo</h1>
        <br>
        <p style="font-size:16px">
            SILMA TTS v1 is a high-performance, 150M-parameter bilingual (Arabic/English) TTS model developed by <a href="https://silma.ai">SILMA.AI</a>.    
            Check out the <a href="https://huggingface.co/silma-ai/silma-tts" target="_blank">model page</a> or explore the repository on <a href="https://github.com/SILMA-AI/silma-tts" target="_blank">Github</a> for more details.
        </p>
        """,
        elem_id="tool-header"
    )
    
    with gr.Row():
        with gr.Column():
            # Define inputs
            ref_audio_input = gr.Audio(
                label="Reference Audio", 
                type="filepath", 
                value="ar.ref.24k.wav"
            )
            ref_text_input = gr.Textbox(
                label="Reference Text", 
                value="ويدقق النظر في القرآن الكريم وسائر الكتب السماوية ويتبع مسالك الرسل العظام عليهم الصلاة والسلام."
            )
            gen_text_input = gr.Textbox(
                label="Generation Text", 
                lines=5,
                value="""
أنا نموذج جديد من سلمى لتحويل النص إلى كلام، يمكنني التحدث باللغة العربية مع أو بدون علامات التشكيل.
I am the new SILMA model for converting text to speech, I can speak Arabic with or without diacritics.
""".strip()
            )
            submit_btn = gr.Button("Generate Speech")
            
        with gr.Column():
            audio_output = gr.Audio(label="Generated Speech")


    gr.Markdown("Note: enter text without diacritics, and our model will add them automatically. If you include full Tashkeel, auto-diacritization is disabled to preserve your original input")

    # When ref_audio_input changes, we update ref_text_input with an empty string
    ref_audio_input.input(
        fn=lambda: "", 
        inputs=None, 
        outputs=ref_text_input
    )
    
    # Set up the click event for the button
    submit_btn.click(
        fn=do_inference,
        inputs=[ref_audio_input, ref_text_input, gen_text_input],
        outputs=audio_output
    )

def main():
    print("Starting app...")

    demo.queue().launch(ssr_mode=False, allowed_paths=[CURR_BASE_DIR+"/images/"])

if __name__ == "__main__":
    main()