import gradio as gr
from transformers import pipeline

# Load Arabic NLP model for intent classification
intent_classifier = pipeline("text-classification", model="aubmindlab/bert-base-arabertv02")

# Load language detection model
language_detector = pipeline("text-classification", model="papluca/xlm-roberta-base-language-detection")

# Omdurman National Bank-specific guidelines
ONB_GUIDELINES = {
    "balance": "يمكنك التحقق من رصيدك عبر الإنترنت أو عبر تطبيق الهاتف الخاص ببنك أم درمان الوطني.",
    "lost_card": "في حالة فقدان البطاقة، اتصل بالرقم 249-123-456-789 فورًا.",
    "loan": "شروط القرض تشمل الحد الأدنى للدخل (5000 جنيه سوداني) وتاريخ ائتماني جيد.",
    "transfer": "لتحويل الأموال، استخدم تطبيق الهاتف أو الخدمة المصرفية عبر الإنترنت.",
    "new_account": "لفتح حساب جديد، قم بزيارة أقرب فرع مع جواز سفرك أو هويتك الوطنية.",
    "interest_rates": "أسعار الفائدة على الودائع تتراوح بين 5% إلى 10% سنويًا.",
    "branches": "فروعنا موجودة في أم درمان، الخرطوم، وبورتسودان. زيارة موقعنا للتفاصيل.",
    "working_hours": "ساعات العمل من 8 صباحًا إلى 3 مساءً من الأحد إلى الخميس.",
    "contact": "الاتصال بنا على الرقم 249-123-456-789 أو عبر البريد الإلكتروني info@onb.sd."
}

# Map intents to responses
INTENT_TO_RESPONSE = {
    "balance": "balance",
    "lost_card": "lost_card",
    "loan": "loan",
    "transfer": "transfer",
    "new_account": "new_account",
    "interest_rates": "interest_rates",
    "branches": "branches",
    "working_hours": "working_hours",
    "contact": "contact"
}

def detect_language(text):
    # Use Hugging Face language detection model
    result = language_detector(text)
    language = result[0]['label']
    return language

def classify_intent(message: str):
    # Use NLP model to classify the user's intent
    result = intent_classifier(message)
    intent = result[0]['label']
    return INTENT_TO_RESPONSE.get(intent, "unknown")

def respond(message: str):
    # Detect language
    language = detect_language(message)
    
    # If the language is not Arabic, return a default response
    if language != "ar":
        return "عذرًا، هذا المساعد يدعم اللغة العربية فقط. الرجاء إعادة الصياغة بالعربية."
    
    # Classify the user's intent using NLP
    intent = classify_intent(message)
    
    # If intent is recognized, return the corresponding response
    if intent != "unknown":
        return ONB_GUIDELINES.get(intent, "عذرًا، لم يتم التعرف على الخيار المحدد.")
    
    # Fallback to keyword matching if NLP doesn't recognize the intent
    for keyword, key in INTENT_TO_RESPONSE.items():
        if keyword in message:
            return ONB_GUIDELINES.get(key, "عذرًا، لم يتم التعرف على الخيار المحدد.")
    
    # Default response if no intent or keyword is matched
    return "عذرًا، لم أفهم سؤالك. الرجاء إعادة الصياغة أو اختيار أحد الخيارات التالية: " + ", ".join(INTENT_TO_RESPONSE.keys())

# Chat interface
with gr.Blocks(css=".gradio-container {direction: rtl;}") as demo:
    gr.Markdown("# <center>بنك أم درمان الوطني - المساعد المصرفي</center>")
    
    with gr.Tab("المحادثة"):
        gr.Markdown("## اكتب سؤالك هنا:")
        
        # Text input
        text_input = gr.Textbox(label="السؤال")
        
        # Submit button
        submit_btn = gr.Button("إرسال")
        
        # Output textbox for responses
        output = gr.Textbox(label="الرد", interactive=False)
        
        # Link inputs and button to response function
        submit_btn.click(
            fn=respond,
            inputs=text_input,
            outputs=output
        )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True  # Enable public link
    )