import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Set page configuration
st.set_page_config(page_title="AI Translator", page_icon="🌍", layout="centered")

# Custom styling
st.markdown(
    """
    <style>
    body {
        background-color: #f4f4f4;
    }
    .stTextArea textarea {
        font-size: 18px;
        height: 150px !important;
    }
    .stSelectbox, .stButton button {
        font-size: 18px;
    }
    .footer {
        text-align: center;
        margin-top: 20px;
        font-size: 14px;
        color: gray;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# Language pairs (including Urdu & Arabic)
language_pairs = {
    "English → Arabic": ("en", "ar"),
    "Arabic → English": ("ar", "en"),
    "English → Urdu": ("en", "ur"),
    "Urdu → English": ("ur", "en"),
    "Urdu → Arabic": ("ur", "ar"),
    "Arabic → Urdu": ("ar", "ur"),
    "English → French": ("en", "fr"),
    "French → English": ("fr", "en"),
    "English → Spanish": ("en", "es"),
    "Spanish → English": ("es", "en"),
    "English → German": ("en", "de"),
    "German → English": ("de", "en"),
    "English → Italian": ("en", "it"),
    "Italian → English": ("it", "en"),
    "English → Dutch": ("en", "nl"),
    "Dutch → English": ("nl", "en"),
    "English → Russian": ("en", "ru"),
    "Russian → English": ("ru", "en"),
    "English → Chinese": ("en", "zh"),
    "Chinese → English": ("zh", "en"),
    "English → Hindi": ("en", "hi"),
    "Hindi → English": ("hi", "en"),
}

# Load model (cached for performance)
@st.cache_resource
def load_model(src_lang, tgt_lang):
    model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
    try:
        tokenizer = MarianTokenizer.from_pretrained(model_name)
        model = MarianMTModel.from_pretrained(model_name)
        return tokenizer, model
    except Exception as e:
        st.error(f"❌ Model {model_name} not found. Try another language pair.")
        return None, None

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

# UI Design
st.title("🌍 AI-Powered Translator")
st.subheader("Translate text instantly with open-source AI models.")

# Language selection
selected_pair = st.selectbox("🔄 Select Language Pair:", list(language_pairs.keys()))
src_lang, tgt_lang = language_pairs[selected_pair]

# Load the translation model
tokenizer, model = load_model(src_lang, tgt_lang)

# Input text
text_input = st.text_area("✍️ Enter text to translate:", "")

# Translate button
if st.button("🚀 Translate Now"):
    if text_input and tokenizer and model:
        translated_text = translate_text(text_input, tokenizer, model)
        st.subheader("✅ Translated Text:")
        st.success(translated_text)
    else:
        st.warning("⚠️ Please enter text to translate.")

# Footer
st.markdown('<p class="footer">🌟 Made with ❤️ using Helsinki-NLP models</p>', unsafe_allow_html=True)