import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Language options
LANGUAGES = {
    "English": "en",
    "French": "fr",
    "Spanish": "es",
    "German": "de",
    "Italian": "it",
    "Chinese": "zh",
    "Japanese": "ja",
    "Arabic": "ar",
    "Russian": "ru",
    "Hindi": "hi",
    "Urdu": "ur",
}

# Load model with caching
@st.cache_resource
def load_model(src_lang, tgt_lang):
    model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model

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

# App title and description
st.markdown(
    """
    <style>
    .title {
        font-size: 40px;
        text-align: center;
        color: #4CAF50;
    }
    .description {
        font-size: 16px;
        text-align: center;
        color: #555555;
    }
    </style>
    <div class="title">🌍 Language Translator</div>
    <p class="description">Easily translate text between multiple languages with this interactive app powered by Hugging Face Transformers.</p>
    """,
    unsafe_allow_html=True,
)

st.sidebar.title("Translator Settings")
source_lang = st.sidebar.selectbox("Source Language", list(LANGUAGES.keys()))
target_lang = st.sidebar.selectbox("Target Language", list(LANGUAGES.keys()))

st.markdown("---")

if source_lang != target_lang:
    text_to_translate = st.text_area(f"Enter text in {source_lang}:", height=150, placeholder="Type your text here...")
    
    if st.button("Translate", type="primary"):
        if text_to_translate.strip():
            with st.spinner("Translating... Please wait."):
                src_lang_code = LANGUAGES[source_lang]
                tgt_lang_code = LANGUAGES[target_lang]
                tokenizer, model = load_model(src_lang_code, tgt_lang_code)
                translated_text = translate_text(tokenizer, model, text_to_translate)
            st.success("Translation Complete!")
            st.markdown("### Translated Text:")
            st.text_area(f"In {target_lang}:", value=translated_text, height=150, disabled=True)
        else:
            st.error("Please enter text to translate.")
else:
    st.warning("Source and target languages must be different.")

st.markdown("---")
st.markdown(
    """
    <footer style="text-align: center; font-size: 14px; color: #888888;">
        Built with ❤️ using Streamlit and Hugging Face Transformers
    </footer>
    """,
    unsafe_allow_html=True,
)