import base64
import io
import os
from typing import Dict, Any, List
import torch
from PIL import Image
from transformers import pipeline, AutoConfig

class EndpointHandler:
    def __init__(self, model_dir: str) -> None:
        print(f"بدء تهيئة النموذج من المسار: {model_dir}")
        print(f"قائمة الملفات في المسار: {os.listdir(model_dir)}")

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"استخدام الجهاز: {self.device}")

        try:
            print("تحميل النموذج باستخدام pipeline")

            self.classifier = pipeline(
                task="image-classification",
                model="yaya36095/ai-source-detector",
                device=0 if torch.cuda.is_available() else -1,
                torch_dtype=torch.float16
            )

            print("تم تحميل النموذج بنجاح")
            self.fallback_mode = False

        except Exception as e:
            print(f"خطأ أثناء تهيئة النموذج: {e}")
            try:
                print("محاولة تحميل النموذج بطريقة بديلة...")
                config = AutoConfig.from_pretrained("yaya36095/ai-source-detector")
                self.fallback_mode = True
                self.config = config
                print("تم التحويل إلى وضع المحاكاة البسيطة")
            except Exception as e2:
                print(f"فشلت المحاولة البديلة أيضًا: {e2}")
                raise

    def __call__(self, data: Any) -> List[Dict[str, Any]]:
        print(f"استدعاء __call__ مع نوع البيانات: {type(data)}")

        img = None
        try:
            if isinstance(data, Image.Image):
                print("البيانات هي صورة PIL")
                img = data

            elif isinstance(data, dict):
                payload = data.get("inputs") or data.get("image")

                if isinstance(payload, Image.Image):
                    img = payload

                elif isinstance(payload, bytes):
                    try:
                        img = Image.open(io.BytesIO(payload)).convert("RGB")
                        print("تم فتح الصورة من Bytes بنجاح")
                    except Exception as e:
                        print(f"فشل في فتح الصورة من Bytes: {e}")

                elif isinstance(payload, str):
                    try:
                        # محاولة فك base64
                        try:
                            img = Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
                            print("تم فتح الصورة من base64 بنجاح")
                        except Exception:
                            # لو فشل، نحاول اعتباره مسار ملف محلي
                            if os.path.exists(payload):
                                img = Image.open(payload).convert("RGB")
                                print("تم فتح الصورة من مسار محلي بنجاح")
                    except Exception as e:
                        print(f"فشل في فتح الصورة من النص: {e}")

            if img is None:
                print("لم يتم العثور على صورة صالحة")
                return [{"label": "error", "score": 1.0}]

            if self.fallback_mode:
                print("استخدام وضع المحاكاة البسيطة")
                results = [
                    {"label": "real", "score": 0.5},
                    {"label": "stable_diffusion", "score": 0.2},
                    {"label": "midjourney", "score": 0.15},
                    {"label": "dalle", "score": 0.1},
                    {"label": "other_ai", "score": 0.05}
                ]
                results.sort(key=lambda x: x["score"], reverse=True)
                return [results[0]]

            print("تصنيف الصورة باستخدام النموذج الحقيقي")
            results = self.classifier(img)

            if isinstance(results, list) and len(results) > 0:
                return [results[0]]
            else:
                print("لم يتم الحصول على نتائج صالحة من النموذج")
                return [{"label": "error", "score": 1.0}]

        except Exception as e:
            print(f"حدث استثناء: {e}")
            return [{"label": "real", "score": 0.5}]