import os
import logging
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import pipeline
import torch

app = Flask(__name__)
CORS(app)

# تهيئة السجل (Logging)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# الفئات المحددة مسبقًا
PREDEFINED_LABELS = [
    "مطار", "مطعم", "سوق", "حديقة", "ملعب", "مسجد",
    "طريق", "مخبز", "صيدلية", "مستشفى", "مصنع",
    "محطة وقود", "جامعة", "مطبخ", "غرفة نوم",
    "حمام", "غرفة معيشة", "شرفة", "مكتب", "صف دراسي"
]

# متغيرات عامة للنموذج والتهيئة
MODEL = None
DEVICE = None
LABELS_ENCODED = None  # سيتم تخزين تمثيل الـ Labels هنا

def initialize():
    global MODEL, DEVICE, LABELS_ENCODED
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    logger.info(f"جار التحميل على الجهاز: {DEVICE}")
    
    MODEL = pipeline(
        task="zero-shot-classification",
        model="joeddav/xlm-roberta-large-xnli",
        device=DEVICE,
        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
    )
    
    # تهيئة النموذج مع الـ Labels كاملة (Warm-up)
    logger.info("جار تهيئة النموذج مع الـ Labels...")
    dummy_result = MODEL(
        "تهيئة النموذج",
        PREDEFINED_LABELS,
        multi_label=False
    )
    logger.info("تم تحميل النموذج والـ Labels بنجاح")

# استدعاء التهيئة عند بدء التشغيل
initialize()

@app.route('/classify', methods=['POST'])
def classify():
    try:
        text = request.json.get('text', '').strip()
        if not text:
            return jsonify({"error": "يجب تقديم نص للتصنيف"}), 400
        
        # الاستدعاء باستخدام الـ Labels المخزنة مسبقًا
        result = MODEL(text, PREDEFINED_LABELS, multi_label=False)
        
        return jsonify({
            "prediction": {
                "label": result["labels"][0],
                "score": float(result["scores"][0])
            }
        })
    
    except Exception as e:
        logger.error(f"خطأ في التصنيف: {str(e)}")
        return jsonify({"error": "حدث خطأ أثناء المعالجة"}), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)