import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Define available target languages and their Hugging Face model names
language_options = {
    "French": "Helsinki-NLP/opus-mt-en-fr",
    "German": "Helsinki-NLP/opus-mt-en-de",
    "Spanish": "Helsinki-NLP/opus-mt-en-es",
    "Italian": "Helsinki-NLP/opus-mt-en-it",
    "Dutch": "Helsinki-NLP/opus-mt-en-nl",
    "Urdu": "Helsinki-NLP/opus-mt-en-ur",
    "Punjabi": "Helsinki-NLP/opus-mt-en-pa",  # Punjabi language model
    "Pashto": "Helsinki-NLP/opus-mt-en-ps",  # Pashto language model
}

# Function to load the model based on the selected language
@st.cache_resource
def load_model(model_name):
    model = MarianMTModel.from_pretrained(model_name)
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    return model, tokenizer

# Function to translate text
def translate_text(model, tokenizer, text):
    inputs = tokenizer.encode(text, return_tensors='pt', truncation=True, padding=True)
    translated = model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Streamlit app layout with dark gradient colors
st.markdown("""
    <style>
        body {
            background: linear-gradient(135deg, #1f1c2c, #928dab);
            color: #e0e0e0;
        }
        .stApp {
            background: linear-gradient(135deg, #1f1c2c, #928dab);
        }
        .css-1d391kg {
            background-color: rgba(31, 28, 44, 0.8);
            color: #e0e0e0;
            border-radius: 10px;
            padding: 20px;
            border: 1px solid #333;
        }
        .stButton > button {
            background-color: #333;
            background-image: linear-gradient(to right, #4facfe, #00f2fe);
            color: #e0e0e0;
            border: none;
            border-radius: 8px;
            padding: 10px 20px;
        }
        .stButton > button:hover {
            background-image: linear-gradient(to right, #00c6ff, #0072ff);
            border: none;
        }
        .stTextArea textarea {
            background-color: #333;
            background-image: linear-gradient(to right, #4facfe, #00f2fe);
            color: #e0e0e0;
            border: none;
            border-radius: 8px;
        }
        .stSidebar {
            background-color: rgba(31, 28, 44, 0.9);
        }
    </style>
""", unsafe_allow_html=True)

# App title and subtitle with dark gradient
st.markdown("<h1 style='text-align: center; color: #e0e0e0;'>🌍 Polyglot Translator 🌍</h1>", unsafe_allow_html=True)
st.markdown("<p style='text-align: center; color: #b0b0b0; font-size: 20px;'>Empowering You to Speak Any Language!</p>", unsafe_allow_html=True)

# Sidebar for language selection and instructions
st.sidebar.title("Language Options")
language = st.sidebar.selectbox("Choose target language", list(language_options.keys()))

st.sidebar.markdown("### How to use")
st.sidebar.write("1. Enter English text in the box below.")
st.sidebar.write("2. Select the target language from the options.")
st.sidebar.write("3. Click **Translate** to get the result.")

# Input text
st.markdown("<h3 style='color: #e0e0e0;'>Enter text in English to translate:</h3>", unsafe_allow_html=True)
text = st.text_area("", height=150)

# Character count
st.write(f"<p style='color: #e0e0e0;'>Character count: {len(text)}</p>", unsafe_allow_html=True)

# Button to translate
if st.button("Translate"):
    if text:
        # Show a spinner during the translation process
        with st.spinner('Translating...'):
            # Load model and tokenizer based on selected language
            model_name = language_options[language]
            model, tokenizer = load_model(model_name)
            
            # Perform translation
            translated_text = translate_text(model, tokenizer, text)
        
        # Display the translation
        st.markdown(f"<h3 style='color: #e0e0e0;'>Translated Text ({language}):</h3>", unsafe_allow_html=True)
        st.success(translated_text)
    else:
        st.error("Please enter text to translate.")