from transformers import pipeline
import arabic_reshaper
from bidi.algorithm import get_display

class ArabicChatModel:
    def __init__(self):
        self.model = None
        self.load_model()
    
    def load_model(self):
        """تحميل النموذج مع التعامل مع الأخطاء"""
        try:
            self.model = pipeline(
                "text2text-generation", 
                model="MBZUAI/LaMini-Flan-T5-248M"  # نموذج خفيف
            )
        except Exception as e:
            print(f"خطأ في تحميل النموذج: {e}")
            raise
    
    def format_arabic(self, text):
        """تنسيق النص العربي"""
        reshaped = arabic_reshaper.reshape(text)
        return get_display(reshaped)
    
    def generate_response(self, message, history):
        """توليد رد مع التعامل مع الأخطاء"""
        try:
            # بناء السياق من تاريخ المحادثة
            context = "\n".join(
                [f"User: {h[0]}\nAI: {h[1]}" for h in history] + 
                [f"User: {message}"]
            )
            
            # توليد الإجابة
            response = self.model(context, max_length=100)[0]['generated_text']
            
            # تنظيف الإجابة وإرجاعها
            return self.format_arabic(response.split("AI:")[-1].strip())
        
        except Exception as e:
            print(f"خطأ في توليد الرد: {e}")
            return "عذراً، حدث خطأ. يرجى المحاولة لاحقاً."