import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# ---------------- PAGE CONFIG ----------------
st.set_page_config(
    page_title="AI Language Translator",
    page_icon="🌍",
    layout="wide"
)

# ---------------- CUSTOM CSS ----------------
st.markdown("""
<style>
.stApp {
    background: linear-gradient(135deg, #0f172a, #1e293b);
}

.title {
    text-align:center;
    font-size:48px;
    font-weight:bold;
    color:white;
}

.subtitle {
    text-align:center;
    color:#cbd5e1;
    margin-bottom:20px;
}

.result-box {
    background:#1e293b;
    padding:20px;
    border-radius:15px;
    color:white;
    font-size:18px;
    border:1px solid #334155;
}
</style>
""", unsafe_allow_html=True)

# ---------------- TITLE ----------------
st.markdown(
    '<p class="title">🌍 AI Language Translator</p>',
    unsafe_allow_html=True
)

st.markdown(
    '<p class="subtitle">Powered by Hugging Face NLLB-200</p>',
    unsafe_allow_html=True
)

# ---------------- LOAD MODEL ----------------
@st.cache_resource
def load_model():
    model_name = "facebook/nllb-200-distilled-600M"

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

    return tokenizer, model

tokenizer, model = load_model()

# ---------------- TRANSLATION FUNCTION ----------------
def translate_text(text, src_lang, tgt_lang):

    tokenizer.src_lang = src_lang

    encoded = tokenizer(
        text,
        return_tensors="pt",
        truncation=True
    )

    generated_tokens = model.generate(
        **encoded,
        forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
        max_length=256
    )

    translated_text = tokenizer.batch_decode(
        generated_tokens,
        skip_special_tokens=True
    )[0]

    return translated_text

# ---------------- LANGUAGES ----------------
languages = {
    "English": "eng_Latn",
    "Urdu": "urd_Arab",
    "Arabic": "arb_Arab",
    "French": "fra_Latn",
    "German": "deu_Latn",
    "Spanish": "spa_Latn",
    "Hindi": "hin_Deva",
    "Chinese": "zho_Hans",
    "Japanese": "jpn_Jpan",
    "Turkish": "tur_Latn",
    "Russian": "rus_Cyrl"
}

# ---------------- UI ----------------
col1, col2 = st.columns(2)

with col1:
    source_lang = st.selectbox(
        "Source Language",
        list(languages.keys())
    )

with col2:
    target_lang = st.selectbox(
        "Target Language",
        list(languages.keys()),
        index=1
    )

text = st.text_area(
    "Enter Text",
    height=200,
    placeholder="Type your text here..."
)

# ---------------- TRANSLATE BUTTON ----------------
if st.button("🚀 Translate", use_container_width=True):

    if text.strip():

        with st.spinner("Translating..."):

            translated_text = translate_text(
                text,
                languages[source_lang],
                languages[target_lang]
            )

            st.markdown(
                f"""
                <div class="result-box">
                    <h3>Translation</h3>
                    <p>{translated_text}</p>
                </div>
                """,
                unsafe_allow_html=True
            )

            st.success("Translation Completed!")

    else:
        st.warning("Please enter text to translate.")