import streamlit as st
from transformers import MarianMTModel, MarianTokenizer, pipeline

# Set page configuration
st.set_page_config(page_title="Language Translation App", layout="centered")

# Title of the app
st.title("🌍 Language Translation App")
st.markdown("""
**Purpose**: Translate text between multiple languages  
**Use Case**: Help users who want to learn or speak different languages or communicate with others in their preferred language.
""")

# Language mapping to language codes
language_mapping = {
    "English": "en",
    "French": "fr",
    "German": "de",
    "Hindi": "hi",
    "Spanish": "es",
    "Italian": "it",
}

# Language pair selection
source_lang = st.selectbox("Select Source Language", list(language_mapping.keys()))
target_lang = st.selectbox("Select Target Language", list(language_mapping.keys()))

# Correctly construct the model name for Hugging Face
model_name = f"Helsinki-NLP/opus-mt-{language_mapping[source_lang]}-{language_mapping[target_lang]}"

# Function to load the translation model
@st.cache_resource
def load_pipeline(model_name):
    try:
        model = MarianMTModel.from_pretrained(model_name)
        tokenizer = MarianTokenizer.from_pretrained(model_name)
        return pipeline("translation", model=model, tokenizer=tokenizer)
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

# Initialize the translator
translator = load_pipeline(model_name)

# Text input field for translation
text_input = st.text_area("Enter text to translate", height=150)

# Translate and display result
if st.button("Translate"):
    if text_input.strip():
        with st.spinner("Translating..."):
            result = translator(text_input)
            translated_text = result[0]['translation_text']
            st.success("Translation complete!")
            st.text_area("Translated Text", translated_text, height=150)
    else:
        st.warning("Please enter some text to translate.")