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

MODEL_NAME = "deepset/xlm-roberta-base-squad2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForQuestionAnswering.from_pretrained(MODEL_NAME)

EXAMPLES = [
    [
        "پاکستان جنوبی ایشیا میں واقع ایک ملک ہے۔ اس کا دارالحکومت اسلام آباد ہے۔ پاکستان کی آبادی تقریباً 22 کروڑ ہے۔ پاکستان 1947 میں برطانوی راج سے آزاد ہوا۔",
        "پاکستان کا دارالحکومت کیا ہے؟"
    ],
    [
        "کرکٹ پاکستان کا سب سے مقبول کھیل ہے۔ پاکستان کرکٹ ٹیم نے 1992 میں عمران خان کی قیادت میں ورلڈ کپ جیتا۔",
        "پاکستان نے ورلڈ کپ کب جیتا؟"
    ],
    [
        "فیصل مسجد اسلام آباد میں واقع ہے اور یہ پاکستان کی سب سے بڑی مسجد ہے۔ اس کی تعمیر 1986 میں مکمل ہوئی۔",
        "فیصل مسجد کہاں ہے؟"
    ],
]

def answer_question(context, question):
    if not context.strip() or not question.strip():
        return "", 0.0, ""

    inputs = tokenizer(
        question,
        context,
        return_tensors="pt",
        truncation=True,
        max_length=512
    )

    with torch.no_grad():
        outputs = model(**inputs)

    start = torch.argmax(outputs.start_logits)
    end = torch.argmax(outputs.end_logits) + 1
    tokens = inputs["input_ids"][0][start:end]
    answer = tokenizer.decode(tokens, skip_special_tokens=True)

    # Confidence from logits
    start_score = torch.softmax(outputs.start_logits, dim=1).max().item()
    end_score = torch.softmax(outputs.end_logits, dim=1).max().item()
    confidence = round((start_score + end_score) / 2 * 100, 2)

    highlighted = context.replace(answer, f"**{answer}**") if answer in context else context

    return answer, confidence, highlighted


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; }
.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;
}
.answer-box textarea {
    font-family: 'Noto Nastaliq Urdu', serif !important;
    font-size: 1.4rem !important; direction: rtl !important;
    text-align: right !important; background: #12121e !important;
    border: 1px solid #2a2a3e !important; color: #f59e0b !important;
    border-radius: 10px !important; font-weight: 600 !important;
}
button.primary { background: #f59e0b !important; border: none !important; color: #0a0a0f !important; font-weight: 600 !important; border-radius: 8px !important; }
footer { display: none !important; }
"""

with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    gr.HTML("""
        <div id="title">
            <h1>❓ Urdu Question Answering</h1>
            <p>اردو متن سے جوابات نکالیں — XLM-RoBERTa fine-tuned on SQuAD2</p>
        </div>
    """)

    with gr.Row():
        with gr.Column(scale=2):
            context_input = gr.Textbox(
                label="متن — Context", placeholder="یہاں اردو متن پیسٹ کریں...",
                lines=8, elem_classes="urdu-box"
            )
            question_input = gr.Textbox(
                label="سوال — Question", placeholder="یہاں سوال لکھیں...",
                lines=2, elem_classes="urdu-box"
            )
            ask_btn = gr.Button("جواب تلاش کریں — Find Answer", variant="primary")

        with gr.Column(scale=1):
            answer_out = gr.Textbox(label="جواب — Answer", interactive=False, lines=3, elem_classes="answer-box")
            confidence_out = gr.Number(label="Confidence %", interactive=False)
            highlighted_out = gr.Textbox(label="Answer in context", interactive=False, lines=6, elem_classes="urdu-box")

    gr.Examples(examples=EXAMPLES, inputs=[context_input, question_input], label="مثالیں — 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:#f59e0b">H-Layba</a> ·
            Model: XLM-RoBERTa · Fine-tuned on SQuAD2
        </div>
    """)

    ask_btn.click(fn=answer_question, inputs=[context_input, question_input], outputs=[answer_out, confidence_out, highlighted_out])
    question_input.submit(fn=answer_question, inputs=[context_input, question_input], outputs=[answer_out, confidence_out, highlighted_out])

demo.launch()