import streamlit as st
from transformers import pipeline
from gtts import gTTS
import uuid
import os

# Language options and model map (now with 10 languages)
MODEL_MAP = {
    ("en", "fr"): "Helsinki-NLP/opus-mt-en-fr",
    ("fr", "en"): "Helsinki-NLP/opus-mt-fr-en",
    ("en", "de"): "Helsinki-NLP/opus-mt-en-de",
    ("de", "en"): "Helsinki-NLP/opus-mt-de-en",
    ("en", "es"): "Helsinki-NLP/opus-mt-en-es",
    ("es", "en"): "Helsinki-NLP/opus-mt-es-en",
    ("en", "it"): "Helsinki-NLP/opus-mt-en-it",
    ("it", "en"): "Helsinki-NLP/opus-mt-it-en",
    ("en", "zh"): "Helsinki-NLP/opus-mt-en-zh",  # Chinese
    ("zh", "en"): "Helsinki-NLP/opus-mt-zh-en",  # Chinese
    ("en", "ur"): "Helsinki-NLP/opus-mt-en-ur",  # Urdu
    ("ur", "en"): "Helsinki-NLP/opus-mt-ur-en",  # Urdu
    ("en", "ar"): "Helsinki-NLP/opus-mt-en-ar",  # Arabic
    ("ar", "en"): "Helsinki-NLP/opus-mt-ar-en",  # Arabic
    ("en", "fa"): "Helsinki-NLP/opus-mt-en-fa",  # Persian
    ("fa", "en"): "Helsinki-NLP/opus-mt-fa-en",  # Persian
    ("en", "tr"): "Helsinki-NLP/opus-mt-en-tr",  # Turkish
    ("tr", "en"): "Helsinki-NLP/opus-mt-tr-en",  # Turkish
}

TRANSLATORS = {}

# Language map for the dropdowns
lang_map = {
    "English": "en",
    "French": "fr",
    "German": "de",
    "Spanish": "es",
    "Italian": "it",
    "Chinese": "zh",  # Simplified Chinese
    "Urdu": "ur",
    "Arabic": "ar",
    "Persian": "fa",
    "Turkish": "tr"
}

# UI
st.title("🌍 Multilingual Translator with Text-to-Speech")

lang_options = list(lang_map.keys())
src_lang_name = st.selectbox("Select source language", lang_options, index=0)
tgt_lang_name = st.selectbox("Select target language", lang_options, index=1)

# Map language names to codes
src_lang = lang_map[src_lang_name]
tgt_lang = lang_map[tgt_lang_name]

text_input = st.text_area("Enter text to translate:")

if st.button("Translate & Speak"):
    if not text_input.strip():
        st.warning("Please enter some text.")
    elif src_lang == tgt_lang:
        st.warning("Source and target languages must be different.")
    else:
        # Load model lazily
        if (src_lang, tgt_lang) not in TRANSLATORS:
            with st.spinner(f"Loading model: {MODEL_MAP[(src_lang, tgt_lang)]}"):
                TRANSLATORS[(src_lang, tgt_lang)] = pipeline("translation", model=MODEL_MAP[(src_lang, tgt_lang)])

        translator = TRANSLATORS[(src_lang, tgt_lang)]
        with st.spinner("Translating..."):
            translated_text = translator(text_input)[0]["translation_text"]
            st.success("Translation complete!")

        st.text_area("Translated Text:", translated_text)

        try:
            audio_filename = f"{uuid.uuid4().hex}.mp3"
            tts = gTTS(text=translated_text, lang=tgt_lang)
            tts.save(audio_filename)
            audio_file = open(audio_filename, "rb")
            st.audio(audio_file.read(), format="audio/mp3")
            audio_file.close()
            os.remove(audio_filename)
        except Exception as e:
            st.warning(f"Translation OK, but TTS failed: {e}")