import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

MODEL_NAME = "Helsinki-NLP/opus-mt-ur-en"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

EXAMPLES = [
    ["پاکستان ایک خوبصورت ملک ہے۔ اس کے پہاڑ اور دریا بہت مشہور ہیں۔"],
    ["اسلام آباد پاکستان کا دارالحکومت ہے۔ یہ ایک جدید اور صاف شہر ہے۔"],
    ["کرکٹ پاکستان کا سب سے مقبول کھیل ہے۔ پاکستانی عوام کرکٹ سے بہت محبت کرتے ہیں۔"],
    ["تعلیم ہر انسان کا بنیادی حق ہے۔ علم کی روشنی سے دنیا کو بہتر بنایا جا سکتا ہے۔"],
]

def translate(text, num_beams):
    if not text.strip():
        return "", 0

    inputs = tokenizer(
        text.strip(),
        return_tensors="pt",
        truncation=True,
        max_length=512
    )

    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_length=256,
            num_beams=int(num_beams),
            early_stopping=True
        )

    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    word_count = len(translation.split())
    return translation, word_count


css = """
@import url('https://fonts.googleapis.com/css2?family=Noto+Nastaliq+Urdu&family=Inter:wght@400;500;600&display=swap');

body, .gradio-container { background: #0a0a0f !important; font-family: 'Inter', sans-serif; }

#title { text-align: center; padding: 2rem 1rem 1rem; }
#title h1 { font-size: 2rem; font-weight: 600; color: #e4e4f0; margin-bottom: 0.4rem; }
#title p { color: #6b6b8a; font-size: 0.9rem; }

#stats {
    display: flex; justify-content: center; gap: 2rem;
    padding: 0.5rem; margin-bottom: 1rem;
}
.stat-card {
    background: #12121e; border: 1px solid #2a2a3e;
    border-radius: 10px; padding: 0.6rem 1.2rem; text-align: center;
}
.stat-num { font-size: 1.2rem; font-weight: 600; color: #3b82f6; }
.stat-label { font-size: 0.7rem; color: #6b6b8a; margin-top: 2px; }

.urdu-box textarea {
    font-family: 'Noto Nastaliq Urdu', serif !important;
    font-size: 1.2rem !important; direction: rtl !important;
    text-align: right !important; line-height: 2.2 !important;
    background: #12121e !important; border: 1px solid #2a2a3e !important;
    color: #e4e4f0 !important; border-radius: 10px !important;
}
.urdu-box textarea:focus {
    border-color: #3b82f6 !important;
    box-shadow: 0 0 0 2px rgba(59, 130, 246, 0.15) !important;
}
.english-box textarea {
    font-family: 'Inter', sans-serif !important;
    font-size: 1.1rem !important;
    background: #12121e !important; border: 1px solid #2a2a3e !important;
    color: #60a5fa !important; border-radius: 10px !important;
    line-height: 1.8 !important;
}

button.primary {
    background: #3b82f6 !important; border: none !important;
    color: white !important; font-weight: 500 !important;
    border-radius: 8px !important;
}
button.primary:hover { background: #2563eb !important; }

footer { display: none !important; }
"""

with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:

    gr.HTML("""
        <div id="title">
            <h1>🌐 Urdu → English Translator</h1>
            <p>اردو سے انگریزی ترجمہ — Helsinki-NLP opus-mt-ur-en trained on 2.7M sentence pairs</p>
        </div>
        <div id="stats">
            <div class="stat-card">
                <div class="stat-num">2.7M</div>
                <div class="stat-label">Training sentence pairs</div>
            </div>
            <div class="stat-card">
                <div class="stat-num">opus-mt</div>
                <div class="stat-label">Model architecture</div>
            </div>
            <div class="stat-card">
                <div class="stat-num">Marian</div>
                <div class="stat-label">Framework</div>
            </div>
        </div>
    """)

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="اردو متن — Urdu Text",
                placeholder="یہاں اردو میں لکھیں یا پیسٹ کریں...",
                lines=7,
                elem_classes="urdu-box"
            )
            num_beams = gr.Slider(
                label="Beam search width (higher = better quality, slower)",
                minimum=1, maximum=8, value=4, step=1
            )
            translate_btn = gr.Button("ترجمہ کریں — Translate", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(
                label="English Translation",
                lines=7,
                interactive=False,
                elem_classes="english-box"
            )
            word_count = gr.Number(label="Translation word count", interactive=False)

    gr.Examples(
        examples=EXAMPLES,
        inputs=input_text,
        label="مثالیں — Try these examples"
    )

    gr.HTML("""
        <div style="text-align:center; padding:1.5rem; color:#3a3a5a; font-size:0.8rem;">
            Built by <a href="https://huggingface.co/H-Layba" style="color:#3b82f6">H-Layba</a> ·
            Model: Helsinki-NLP/opus-mt-ur-en · OPUS parallel corpus
        </div>
    """)

    translate_btn.click(
        fn=translate,
        inputs=[input_text, num_beams],
        outputs=[output_text, word_count]
    )
    input_text.submit(
        fn=translate,
        inputs=[input_text, num_beams],
        outputs=[output_text, word_count]
    )

demo.launch()