import os
from typing import Optional
import gradio as gr
from synthesis import synthesize_text

# Configure correct paths
MODEL_DIR = os.path.join(os.path.dirname(__file__), "trained_model")
MODEL_PATH = os.path.join(MODEL_DIR, "checkpoint_latest.pth.tar")
HPARAMS_PATH = os.path.join(MODEL_DIR, "hparams.yml")

# Check if model files exist
if not os.path.exists(MODEL_PATH):
    print(f"ERROR: Model checkpoint not found at {MODEL_PATH}")
    
if not os.path.exists(HPARAMS_PATH):
    print(f"ERROR: Hyperparameters file not found at {HPARAMS_PATH}")

# Synthesis function wrapper
def generate_speech(text: str) -> Optional[str]:
    try:
        if not text.strip():
            return None
            
        # Pass correct paths to the synthesis function
        audio_output = synthesize_text(
            text, 
            checkpoint_path=MODEL_PATH,
            hparams_path=HPARAMS_PATH
        )
        return audio_output
    except Exception as e:
        print(f"Error generating speech: {str(e)}")
        return None

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Myanmar Text-to-Speech Demo")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Enter Myanmar text",
                placeholder="မြန်မာစာ ရိုက်ထည့်ပါ",
                lines=3
            )
            submit_btn = gr.Button("Generate Speech")
            
        with gr.Column():
            audio_output = gr.Audio(label="Generated Speech")
    
    submit_btn.click(
        fn=generate_speech,
        inputs=text_input,
        outputs=audio_output
    )
    
    gr.Markdown("""
    ## Example Phrases
    - မင်္ဂလာပါ (Hello)
    - မြန်မာစကားပြောစနစ်ကို ကြိုဆိုပါတယ် (Welcome to the Myanmar speech system)
    - ဒီစနစ်ဟာ မြန်မာစာကို အသံအဖြစ် ပြောင်းပေးနိုင်ပါတယ် (This system can convert Myanmar text to speech)
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()