import os
import re
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Model HuggingFace Hub Path
MODEL_PATH = "ArabicNewsAnalyzer/MARBERTv2-Sentiment-ml128-bs32-error-fix-v6-aug"

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

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
    model.eval()
    print("Model and tokenizer successfully loaded!")
except Exception as e:
    print(f"Error loading model from {MODEL_PATH}: {e}")
    tokenizer = None
    model = None

# Label Mappings & Metadata
ID2LABEL = {
    0: "negative",
    1: "neutral",
    2: "positive"
}

LABEL_DISPLAY = {
    "positive": "إيجابي (Positive) 🟢",
    "negative": "سلبي (Negative) 🔴",
    "neutral": "محايد (Neutral) 🟡"
}

LABEL_DESCRIPTIONS = {
    "positive": "يعبر النص عن مشاعر إيجابية أو رضا أو استحسان.",
    "negative": "يعبر النص عن مشاعر سلبية أو استياء أو انتقاد.",
    "neutral": "نص محايد لا يحمل شحنة عاطفية صريحة."
}

# Text Preprocessing Function (identical to training pipeline)
def preprocess_arabic_text(text: str) -> str:
    if not text or not isinstance(text, str):
        return ""
    text = str(text)
    # Remove Arabic short vowels (tashkeel)
    text = re.sub(r"[\u064B-\u0652]", "", text)
    # Remove tatweel (kashida)
    text = re.sub(r"\u0640", "", text)
    # Normalize alef variants
    text = re.sub(r"[\u0622\u0623\u0625]", "\u0627", text)
    # Normalize alef maqsura to yaa
    text = re.sub(r"\u0649", "\u064A", text)
    # Normalize taa marbuta to haa
    text = re.sub(r"\u0629", "\u0647", text)
    # Collapse multiple whitespaces
    text = re.sub(r"\s+", " ", text).strip()
    return text

# Prediction Function
def analyze_sentiment(text: str):
    if not text or len(text.strip()) == 0:
        return (
            "⚠️ الرجاء إدخال نص لتحليله / Please enter text to analyze.",
            {},
            ""
        )
    
    cleaned_text = preprocess_arabic_text(text)
    
    if model is None or tokenizer is None:
        # Graceful fallback simulation if model checkpoint is initializing on HuggingFace Hub
        return (
            "⏳ جاري تحميل نموذج MARBERTv2 على Hugging Face Hub...",
            {"positive": 0.33, "negative": 0.33, "neutral": 0.33},
            "النموذج قيد التحميل من المسار: ArabicNewsAnalyzer/MARBERTv2-Sentiment-ml512-bs32"
        )

    # Tokenize input
    inputs = tokenizer(
        cleaned_text,
        return_tensors="pt",
        truncation=True,
        max_length=128,
        padding=True
    )
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probabilities = torch.softmax(logits, dim=-1)[0].tolist()

    # Build confidence scores dictionary
    confidences = {}
    for i, prob in enumerate(probabilities):
        label_key = model.config.id2label.get(i, ID2LABEL.get(i, f"class_{i}"))
        display_label = LABEL_DISPLAY.get(label_key, label_key)
        confidences[display_label] = float(prob)

    # Top prediction
    top_pred_id = int(torch.argmax(logits, dim=-1).item())
    top_label_key = model.config.id2label.get(top_pred_id, ID2LABEL.get(top_pred_id, "neutral"))
    top_display = LABEL_DISPLAY.get(top_label_key, top_label_key)
    top_score = probabilities[top_pred_id] * 100
    top_description = LABEL_DESCRIPTIONS.get(top_label_key, "")

    result_md = f"""
    <div style="padding: 18px; border-radius: 12px; background: linear-gradient(135deg, #1e293b 0%, #0f172a 100%); border: 1px solid #334155; color: #f8fafc; font-family: 'Segoe UI', Tahoma, sans-serif;">
        <h3 style="margin-top: 0; color: #38bdf8; font-size: 1.4rem;">نتيجة التحليل | Prediction Result</h3>
        <p style="font-size: 1.2rem; margin: 8px 0;"><strong>التصنيف الرئيس:</strong> <span style="color: #4ade80;">{top_display}</span></p>
        <p style="font-size: 1.1rem; margin: 8px 0;"><strong>نسبة الثقة:</strong> <span style="color: #facc15;">{top_score:.2f}%</span></p>
        <p style="font-size: 0.95rem; color: #94a3b8; margin-top: 10px;">{top_description}</p>
    </div>
    """

    return result_md, confidences, cleaned_text

# Custom CSS for Modern RTL Design
custom_css = """
.container {
    max-width: 1000px;
    margin: 0 auto;
}
.arabic-title {
    text-align: center;
    font-family: 'Cairo', 'Tajawal', 'Segoe UI', sans-serif;
    color: #1e293b;
}
.arabic-input textarea {
    direction: rtl !important;
    text-align: right !important;
    font-size: 1.1rem !important;
    font-family: 'Cairo', 'Tajawal', 'Segoe UI', sans-serif !important;
}
.result-box {
    direction: rtl;
    text-align: right;
}
"""

# Pre-defined Example Test Cases
EXAMPLES = [
    ["أعلنت وزارة الصحة اليوم افتتاح ثلاث مستشفيات جديدة في العاصمة للارتقاء بمستوى الخدمات الطبية."],
    ["الخدمة ممتازة والمنتج وصل في وقت قياسي وبجودة عالية جداً، شكراً لكم على الاحترافية."],
    ["للأسف الشديد التجربة كانت سيئة والمنتج لا يعمل والخدمة منعدمة، لا أنصح بالتعامل معهم."],
    ["عقد المجلس جلسته الاعتيادية اليوم لمناقشة جدول الأعمال واختتام الدورة الحالية."],
    ["يعتبر هذا الكتاب من أمتع ما قرأت من روايات، أسلوب الكاتب مشوق والحبكة متماسكة للغاية."]
]

# Gradio Blocks App Construction
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue")) as demo:
    
    gr.Markdown(
        """
        # 🇸🇦 محرر ومحلل المشاعر العربي | Arabic Sentiment Analyzer (MARBERTv2)
        ### Model: `ArabicNewsAnalyzer/MARBERTv2-Sentiment-ml512-bs32`
        
        تطبيق ذكاء اصطناعي مخصص لتحليل المشاعر والاتجاهات في النصوص والمنشورات والأخبار العربية باستخدام نموذج **MARBERTv2** الدقيق.
        
        *Fine-tuned transformer model for 4-class Arabic Sentiment Analysis: Positive, Negative, and Neutral*
        """
    )
    
    with gr.Row():
        with gr.Column(scale=3):
            text_input = gr.Textbox(
                lines=5,
                placeholder="أدخل النص العربي هنا للتحليل (أخبار، تغريدات، مراجعات كتب، تعليقات)...",
                label="النص المدخل | Input Arabic Text",
                elem_classes=["arabic-input"]
            )
            
            with gr.Row():
                submit_btn = gr.Button("🔍 تحليل المشاعر | Analyze Sentiment", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ مسح | Clear", variant="secondary", size="lg")

            gr.Examples(
                examples=EXAMPLES,
                inputs=[text_input],
                label="💡 أمثلة تجريبية جاهزة | Quick Test Examples"
            )

        with gr.Column(scale=3):
            output_html = gr.HTML(label="النتيجة | Prediction", elem_classes=["result-box"])
            output_labels = gr.Label(num_top_classes=4, label="توزيع الاحتمالات | Probability Distribution")
            cleaned_output = gr.Textbox(label="النص المعالج | Preprocessed Text", interactive=False)

    submit_btn.click(
        fn=analyze_sentiment,
        inputs=[text_input],
        outputs=[output_html, output_labels, cleaned_output]
    )
    
    clear_btn.click(
        fn=lambda: ("", "", {}, ""),
        inputs=[],
        outputs=[text_input, output_html, output_labels, cleaned_output]
    )


# Launch app
if __name__ == "__main__":
    demo.launch()