from fastapi import FastAPI, Query
from transformers import AutoModelForImageClassification, AutoImageProcessor
from PIL import Image, UnidentifiedImageError
import requests
from io import BytesIO
import torch

# إنشاء التطبيق
app = FastAPI(title="Image to Text API", description="تصنيف الصور باستخدام Hugging Face")

# تحميل النموذج
model = AutoModelForImageClassification.from_pretrained(
    "asyafalni/arabichar-v3",
    trust_remote_code=True
)

# استخدام معالج صور جاهز
processor = AutoImageProcessor.from_pretrained(
    "google/vit-base-patch16-224-in21k"
)

# نقطة نهاية الجذر
@app.get("/")
def root():
    return {
        "message": "مرحبًا بك في Image-to-Text API!",
        "usage": "أرسل رابط صورة إلى /classify عبر POST"
    }

@app.post("/classify")
def classify_image(url: str = Query(..., description="رابط الصورة")):
    try:
        # تحميل الصورة مع timeout ومع التحقق
        response = requests.get(url, timeout=10, stream=True)
        response.raise_for_status()  # رفع استثناء لو كان الرابط غير صالح

        # قراءة الصورة
        try:
            image = Image.open(BytesIO(response.content)).convert("RGB")
        except UnidentifiedImageError:
            return {"error": "الرابط لا يحتوي على صورة صالحة"}

        # تجهيز المدخلات
        inputs = processor(images=image, return_tensors="pt")

        # تمريرها للنموذج
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_class = logits.argmax(-1).item()

        return {
            "url": url,
            "predicted_class": model.config.id2label.get(predicted_class, "Unknown"),
            "scores": logits.softmax(-1).tolist()
        }

    except requests.exceptions.RequestException as e:
        return {"error": f"تعذر الوصول إلى الرابط: {e}"}
    except Exception as e:
        return {"error": str(e)}