import os
import gradio as gr
import torch
from PIL import Image
from torchvision import transforms
from transformers import AutoConfig, AutoModelForImageClassification
from huggingface_hub import snapshot_download

# =========================
# إعداد البيئة لتقليل التحذيرات
# =========================
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# =========================
# ضبط استخدام CPU
# =========================
torch.set_num_threads(1)
torch.set_num_interop_threads(1)

# =========================
# إعداد المعالجة المسبقة
# =========================
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.Grayscale(num_output_channels=3),
    transforms.ToTensor()
])

# =========================
# متغيرات النموذج العالمية
# =========================
model = None
config = None
device = torch.device("cpu")

# =========================
# دالة تحميل النموذج عند الطلب
# =========================
def load_model():
    global model, config
    if model is None or config is None:
        print("🚀 تحميل النموذج لأول مرة من Hugging Face Hub...")
        model_path = snapshot_download(
            repo_id="shahad-alh/arabichar-finetuned-v2",
            local_dir="./model"
        )
        # تحميل الكود المخصص مع trust_remote_code=True
        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
        model = AutoModelForImageClassification.from_pretrained(
            model_path,
            config=config,
            trust_remote_code=True
        )
        model.to(device)
        model.eval()
        print("✅ تم تحميل النموذج بنجاح")

# =========================
# دالة التنبؤ
# =========================
def predict(image: Image.Image):
    try:
        if image is None:
            return "لا توجد صورة"

        # تحميل النموذج عند أول طلب
        load_model()

        tensor = transform(image).unsqueeze(0).to(device)

        with torch.no_grad():
            outputs = model(tensor)
            logits = outputs.logits
            predicted = torch.argmax(logits, dim=1).item()

        label = config.id2label[predicted]
        return label

    except Exception as e:
        print("❌ Prediction error:", e)
        return "خطأ في التنبؤ"

# =========================
# واجهة Gradio
# =========================
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="ارفع صورة حرف عربي"),
    outputs=gr.Textbox(label="الحرف المتوقع"),
    title="Arabic Character Recognition",
    description="نموذج للتعرف على الحروف العربية من الصور",
    flagging_mode="never"
).queue().launch()