import gradio as gr
import torch
import re
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from arabert.preprocess import ArabertPreprocessor
import torch.nn.functional as F


# =========================
# Model configuration
# =========================
MODEL_PATH = "arabert_sentiment_model"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
model.to(device)
model.eval()

arabert_prep = ArabertPreprocessor(
    model_name="aubmindlab/bert-base-arabertv02"
)

label_map = {0: "negative", 1: "positive"}


# =========================
# Arabic validation (REGEX)
# =========================
def is_mostly_arabic(text, threshold=0.6):
    if not isinstance(text, str) or len(text.strip()) == 0:
        return False

    arabic_pattern = re.compile(
        r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDFF\uFE70-\uFEFF]'
    )

    arabic_chars = arabic_pattern.findall(text)
    letters = re.findall(r'\w', text)

    if len(letters) == 0:
        return False

    return len(arabic_chars) / len(letters) >= threshold


# =========================
# Prediction function
# =========================
def predict(text):
    text = text.strip()

    # 🔒 Validation arabe
    if not is_mostly_arabic(text):
        return "Arabic text only ❌", "—"

    if len(text.split()) < 3:
        return "Sentence too short ❌", "—"

    # 🧹 AraBERT preprocessing
    clean_text = arabert_prep.preprocess(text)

    # 🔢 Tokenization
    inputs = tokenizer(
        clean_text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=128
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}

    # 🤖 Prediction
    with torch.no_grad():
        probs = torch.softmax(model(**inputs).logits, dim=1)
        conf, pred = torch.max(probs, dim=1)

    return (
        label_map[pred.item()].capitalize(),
        f"{round(conf.item() * 100, 2)} %"
    )


# =========================
# Gradio Interface
# =========================
gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        lines=3,
        placeholder="أدخل جملة عربية هنا..."
    ),
    outputs=[
        gr.Textbox(label="Sentiment"),
        gr.Textbox(label="Confidence")
    ],
    title="Arabic Sentiment Analysis (AraBERT)",
    description="تحليل المشاعر للنصوص العربية باستخدام AraBERT"
).launch()