"""
Hugging Face Space Application for Arabic Claim & Newsworthiness Detection.
Backbone Model: MARBERTv2 (UBC-NLP/MARBERTv2)
Classes: claim, non_claim, gibberish
"""

import os
import json
import torch
import torch.nn.functional as F
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 1. Model Repository Configuration
# Replace with your Hugging Face model repo ID or set HF_MODEL_ID environment variable
MODEL_ID = os.getenv("HF_MODEL_ID", "ArabicNewsAnalyzer/NewsValidator-V2")


print(f"Loading tokenizer and model from: {MODEL_ID}...")

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
    print("Successfully loaded model from Hugging Face Hub.")
except Exception as e:
    print(f"Could not load from HF Hub ({e}). Attempting local fallback: {LOCAL_MODEL_DIR}...")


# Device setup
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

# Retrieve label mapping from model config
id2label = model.config.id2label if hasattr(model.config, "id2label") and model.config.id2label else {
    0: "claim",
    1: "gibberish",
    2: "non_claim"
}

# Ensure integer keys for id2label
id2label = {int(k): str(v) for k, v in id2label.items()}


def predict_claim(text: str):
    """
    Inference function for Arabic Claim Detection matching 02_train_claim_detection_model.ipynb.
    Takes raw text directly (without custom preprocessing) and computes model predictions.
    """
    if not text or not text.strip():
        return {"Error": "Please enter valid text."}, {}

    # Tokenization
    inputs = tokenizer(
        text,
        truncation=True,
        padding="max_length",
        max_length=128,
        return_tensors="pt"
    ).to(device)

    # Forward pass
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=-1).squeeze(0)
        pred_idx = int(torch.argmax(probs).item())

    # Extract class probabilities dict for Gradio Label component
    prob_dict = {id2label[i]: float(probs[i].item()) for i in range(len(id2label))}
    
    # Detailed summary dict
    top_label = id2label[pred_idx]
    top_confidence = float(probs[pred_idx].item())
    
    structured_result = {
        "Predicted Class": top_label,
        "Confidence Score": f"{top_confidence * 100:.2f}%",
        "Probabilities": {k: f"{v * 100:.2f}%" for k, v in prob_dict.items()}
    }

    return prob_dict, structured_result


# 2. Gradio Web Interface
custom_css = """
body {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}
.title-container {
    text-align: center;
    margin-bottom: 20px;
}
.title-container h1 {
    color: #1E3A8A;
    font-size: 2.2rem;
    font-weight: 700;
}
.title-container p {
    color: #4B5563;
    font-size: 1.1rem;
}
"""

with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <div class="title-container">
            <h1>🔍 تصنيف الادعاءات والأخبار العربية | Arabic Claim Detection</h1>
            <p>نموذج ذكاء اصطناعي قائم على <b>MARBERTv2</b> لتصنيف النصوص العربية إلى ادعاءات إخبارية، نصوص عادية، أو نصوص عشوائية (Gibberish)</p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                lines=5,
                placeholder="أدخل النص العربي للتحليل هنا... (Enter Arabic text here)",
                label="النص المدخل / Input Text",
                elem_id="input-text"
            )
            submit_btn = gr.Button("🚀 تحليل النص / Analyze Text", variant="primary")
            
            gr.Examples(
                examples=[
                    ["اندلع حريق هائل في أحد الأسواق التجارية الرئيسية وسط العاصمة صباح اليوم."],
                    ["وزارة الصحة أعلنت تسجيل عشر إصابات جديدة بفايروس كورونا."],
                    ["الفندق رائع جداً والخدمة ممتازة، أنصح بالزيارة بالتأكيد."],
                    ["أنا أحب هذا الكتاب كثيراً، أسلوب الكاتب ساحر ومشوق للغاية."],
                    ["نيتبمبنب ههههههههههههههه !!!!! 123123 😂😂😂"],
                    ["https://example.com/test?id=123 SELECT * FROM users;"]
                ],
                inputs=[input_text],
                label="أمثلة للاختبار / Sample Examples"
            )

        with gr.Column(scale=1):
            label_output = gr.Label(
                label="الفئة المتوقعة / Predicted Class & Probabilities",
                num_top_classes=3
            )
            json_output = gr.JSON(
                label="التفاصيل الهيكلية / Structured Breakdown"
            )

    submit_btn.click(
        fn=predict_claim,
        inputs=[input_text],
        outputs=[label_output, json_output]
    )

if __name__ == "__main__":
    demo.launch()