from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline


class EndpointHandler:
    def __init__(self, path=""):
        # تحميل النموذج مع تخطي اختلاف الأوزان
        self.model = AutoModelForSequenceClassification.from_pretrained(
            path,
            ignore_mismatched_sizes=True
        )
        self.tokenizer = AutoTokenizer.from_pretrained(path)

        # إعداد البايبلاين للتصنيف النصي
        self.pipeline = pipeline(
            "text-classification",
            model=self.model,
            tokenizer=self.tokenizer
        )

        # أسماء التصنيفات لو حبيت ترجعها بشكل واضح
        self.labels = {
            "AI": "نص مكتوب بالذكاء الاصطناعي 🤖",
            "HUMAN": "نص بشري حقيقي 🧠"
        }

    def __call__(self, data):
        # الحصول على النص
        inputs = data.get("inputs", data if isinstance(data, str) else "")

        # تشغيل النموذج
        outputs = self.pipeline(inputs)

        # استخراج النتيجة الأولى فقط (لأن فيه احتمال يرجع لستة)
        result = outputs[0]
        label = result["label"]
        score = round(result["score"] * 100, 2)  # تحويل للنسبة المئوية

        # رسالة مخصصة باللغة العربية
        message = f"{self.labels.get(label, label)} (نسبة التأكد: {score}%)"

        # إعادة الإخراج بصيغة JSON Array كما تتوقع واجهة HF
        return [{
            "label": label,
            "score": score,
            "message": message
        }]