import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Map of language codes to full language names and flag images
language_dict = {
    "en": "English",
    "de": "German",
    "fr": "French",
    "es": "Spanish",
    "it": "Italian",
    "nl": "Dutch",
    "pt": "Portuguese",
    "ro": "Romanian",
    "ru": "Russian",
    "zh": "Chinese"
}

# Reverse map for language selection
language_codes = {v: k for k, v in language_dict.items()}

# Add custom CSS for styling
st.markdown("""
    <style>
        .stTextInput, .stTextArea {
            background-color: #f4f6f8;
            padding: 10px;
            border-radius: 8px;
        }
        .stButton {
            background-color: #4CAF50;
            color: black;
            padding: 10px 20px;
            border-radius: 5px;
        }
        .stButton:hover {
            background-color: #45a049;
        }
        .stText {
            font-size: 1.2em;
            font-family: 'Helvetica', sans-serif;
        }
        .footer {
            padding: 10px;
            background-color: #f1f1f1;
            text-align: center;
        }
        h1 {
            color: #4CAF50;
        }
    </style>
""", unsafe_allow_html=True)

# Custom Exception for Same Source and Target Language
class SameLanguageException(Exception):
    pass

# Function to load model and tokenizer
def load_translation_model(src_lang, tgt_lang):
    model_name = f'Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}'
    model = MarianMTModel.from_pretrained(model_name)
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    return model, tokenizer

# Function to translate text
def translate_text(text, src_lang, tgt_lang):
    if src_lang == tgt_lang:
        raise SameLanguageException("Source and target languages must be different.")
    
    model, tokenizer = load_translation_model(src_lang, tgt_lang)

    # Tokenize the input text
    inputs = tokenizer(text, return_tensors="pt", padding=True)

    # Translate the text
    translated = model.generate(**inputs)

    # Decode the translated text
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)

    return translated_text

# Language selector with flags and names
def language_selector():
    # Use two columns to display the language selectors
    col1, col2 = st.columns(2)

    with col1:
        source_language = st.selectbox(
            "Select source language",
            list(language_dict.values()),
            key="source_language"
        )

    with col2:
        target_language = st.selectbox(
            "Select target language",
            list(language_dict.values()),
            key="target_language"
        )

    return language_codes[source_language], language_codes[target_language]

# Main application function
def run_app():
    st.title("🌍 Interactive Language Translation App")

    st.subheader("Translate your text in real-time with ease! 🌐")

    # Select source and target languages
    src_lang_code, tgt_lang_code = language_selector()

    # Text input field with placeholder
    text_input = st.text_area("Enter text to translate", placeholder="Type here...", height=150)

    # Option to upload a text file for input
    uploaded_file = st.file_uploader("Or upload a text file", type=["txt"])

    if uploaded_file is not None:
        text_input = uploaded_file.read().decode("utf-8")
        st.text_area("File content", value=text_input, height=150)

    # Button for translation
    if st.button("🔄 Translate", key="translate"):
        if text_input:
            try:
                with st.spinner("🔄 Translating... Please wait..."):
                    # Translate the text
                    translated_text = translate_text(text_input, src_lang_code, tgt_lang_code)
                    # Display translated text directly in the success container
                    st.success(f"✅ Translation: {translated_text}")
                st.balloons()  # Show animation when translation is complete
            except SameLanguageException as e:
                st.error(f"❌ Error: {str(e)}")
        else:
            st.warning("⚠️ Please enter or upload some text to translate.")

# Run the Streamlit app
if __name__ == "__main__":
    run_app()