import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Supported language codes for Helsinki-NLP models
LANGUAGES = {
    "English": "en",
    "French": "fr",
    "German": "de",
    "Spanish": "es",
    "Italian": "it",
    "Portuguese": "pt",
    "Russian": "ru",
    "Chinese": "zh",
    "Arabic": "ar",
    "Hindi": "hi",
    "Urdu": "ur"
}

def get_model_name(src_lang, tgt_lang):
    return f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"

@st.cache_resource(show_spinner=False)
def load_model_and_tokenizer(src_lang, tgt_lang):
    model_name = get_model_name(src_lang, tgt_lang)
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model

def translate(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    translated = model.generate(**inputs)
    return tokenizer.decode(translated[0], skip_special_tokens=True)

# Streamlit UI
st.set_page_config(page_title="🌐 Multi-Language Translator", layout="centered")

st.title("🌐 Multi-Language Translator")
st.markdown("Translate text instantly between multiple languages using open-source AI (Hugging Face Helsinki-NLP models).")

with st.form("translator_form"):
    col1, col2 = st.columns(2)
    with col1:
        src_lang_name = st.selectbox("Select Input Language", list(LANGUAGES.keys()), index=0)
    with col2:
        tgt_lang_name = st.selectbox("Select Output Language", list(LANGUAGES.keys()), index=1)

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

    submit_btn = st.form_submit_button("Translate")

if submit_btn:
    if src_lang_name == tgt_lang_name:
        st.warning("Source and target languages are the same. Please choose different languages.")
    elif not text_input.strip():
        st.warning("Please enter text to translate.")
    else:
        src_lang_code = LANGUAGES[src_lang_name]
        tgt_lang_code = LANGUAGES[tgt_lang_name]
        try:
            tokenizer, model = load_model_and_tokenizer(src_lang_code, tgt_lang_code)
            output = translate(text_input, tokenizer, model)
            st.success("✅ Translation Completed:")
            st.text_area("Translated Text:", value=output, height=150)
        except Exception as e:
            st.error(f"Error: Could not load model for {src_lang_code} to {tgt_lang_code}.")
            st.exception(e)