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

# ── Model & Tokenizer Configuration ──
MODEL_ID = os.getenv("HF_MODEL_ID", "ArabicNewsAnalyzer/MARBERTv2-Single-Arabic-Dialect-Scrapped-MSA-Mask-epoch3")
LOCAL_MODEL_PATH = os.path.join(os.path.dirname(__file__), "..", "results", "best_model")
MAX_LENGTH = 128

COUNTRY_DETAILS = {
    "Algeria": {"flag": "🇩🇿", "name_ar": "الجزائر", "region": "Maghrebi"},
    "Bahrain": {"flag": "🇧🇭", "name_ar": "البحرين", "region": "Gulf"},
    "Egypt": {"flag": "🇪🇬", "name_ar": "مصر", "region": "Egyptian"},
    "Iraq": {"flag": "🇮🇶", "name_ar": "العراق", "region": "Iraqi"},
    "Jordan": {"flag": "🇯🇴", "name_ar": "الأردن", "region": "Levantine"},
    "Kuwait": {"flag": "🇰🇼", "name_ar": "الكويت", "region": "Gulf"},
    "Lebanon": {"flag": "🇱🇧", "name_ar": "لبنان", "region": "Levantine"},
    "Libya": {"flag": "🇱🇾", "name_ar": "ليبيا", "region": "Maghrebi"},
    "Morocco": {"flag": "🇲🇦", "name_ar": "المغرب", "region": "Maghrebi"},
    "Oman": {"flag": "🇴🇲", "name_ar": "عُمان", "region": "Gulf"},
    "Palestine": {"flag": "🇵🇸", "name_ar": "فلسطين", "region": "Levantine"},
    "Qatar": {"flag": "🇶🇦", "name_ar": "قطر", "region": "Gulf"},
    "Saudi_Arabia": {"flag": "🇸🇦", "name_ar": "السعودية", "region": "Gulf"},
    "Sudan": {"flag": "🇸🇩", "name_ar": "السودان", "region": "Nilo-Saharan"},
    "Syria": {"flag": "🇸🇾", "name_ar": "سوريا", "region": "Levantine"},
    "Tunisia": {"flag": "🇹🇳", "name_ar": "تونس", "region": "Maghrebi"},
    "UAE": {"flag": "🇦🇪", "name_ar": "الإمارات", "region": "Gulf"},
    "Yemen": {"flag": "🇾🇪", "name_ar": "اليمن", "region": "Peninsular"},
    "MSA": {"flag": "📝", "name_ar": "الفصحى", "region": "MSA"}
}

# ── Cache Loading ──
@functools.lru_cache(maxsize=1)
def load_model_and_tokenizer():
    path_to_load = MODEL_ID
    if os.path.exists(LOCAL_MODEL_PATH) and os.path.exists(os.path.join(LOCAL_MODEL_PATH, "config.json")):
        path_to_load = LOCAL_MODEL_PATH

    tokenizer = AutoTokenizer.from_pretrained(path_to_load)
    model = AutoModelForSequenceClassification.from_pretrained(path_to_load)
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    id2label = model.config.id2label
    if not id2label or isinstance(list(id2label.keys())[0], str):
        id2label = {int(k): str(v) for k, v in id2label.items()}

    return tokenizer, model, id2label, device

# ── Text Preprocessing ──
def preprocess_text(text: str) -> str:
    if not text:
        return ""
    text = re.sub(r"[\u200b\u200c\u200d\u200e\u200f\ufeff\u00ad]", "", text)
    return re.sub(r"\s+", " ", text).strip()

# ── Inference Function ──
@torch.inference_mode()
def predict_dialect(text: str) -> dict:
    cleaned_text = preprocess_text(text)
    if not cleaned_text:
        return {"error": "Invalid or empty text input."}

    tokenizer, model, id2label, device = load_model_and_tokenizer()

    inputs = tokenizer(
        cleaned_text,
        return_tensors="pt",
        truncation=True,
        max_length=MAX_LENGTH
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}

    outputs = model(**inputs)
    probs = torch.softmax(outputs.logits, dim=-1).squeeze(0).cpu().numpy()

    # Sort probabilities in descending order
    sorted_indices = probs.argsort()[::-1]

    all_predictions = []
    for idx in sorted_indices:
        label = id2label[int(idx)]
        details = COUNTRY_DETAILS.get(label, {"flag": "🏳️", "name_ar": label, "region": "Unknown"})
        all_predictions.append({
            "country": label,
            "name_ar": details["name_ar"],
            "flag": details["flag"],
            "region": details["region"],
            "confidence": round(float(probs[idx]), 4)
        })

    top_prediction = all_predictions[0]

    return {
        "text": cleaned_text,
        "prediction": top_prediction["country"],
        "confidence": top_prediction["confidence"],
        "details": top_prediction,
        "probabilities": {p["country"]: p["confidence"] for p in all_predictions}
    }

# ── Examples List ──
examples = [
    ["واش كاين الزين فهاد البلاد غاب علينا هاد الأيام"],
    ["يا زلمة شو هالحكي الفاضي شو عم تحكي انت"],
    ["شنو هذا الحجي ما يصير هيج ابد والله"],
    ["ايوا كده الكلام ده حلو اوي تسلم ايدك يا فنان"],
    ["والله يا خوي هالمطعم أكله طيب بزاف تبارك الله"],
    ["ايه والله كلامك صح يالغالي ما قصرت ربي يحفظك"],
    ["كيفك حبيبي شلونك اليوم عساك طيب"],
    ["هاذي السالفة ما تنفع خلاص لازم نلقى حل سريع"],
    ["وكالة تسنيم عن مصادر إيرانية: دوي انفجارين في جزيرة قشم سببه التصدي لأهداف معادية عند مدخل مضيق هرمز"],
]

# ── Minimal Gradio Interface ──
demo = gr.Interface(
    fn=predict_dialect,
    inputs=gr.Textbox(lines=3, placeholder="أدخل النص العربي هنا...", label="Input Text"),
    outputs=gr.JSON(label="API Response"),
    examples=examples,
    cache_examples=False,
    title="Arabic Dialect Classifier",
    api_name="predict"
)

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