import gradio as gr
import re

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

# Omdurman National Bank-specific guidelines in Arabic
ONB_GUIDELINES_AR = {
    "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."
}

# Omdurman National Bank-specific guidelines in English
ONB_GUIDELINES_EN = {
    "balance": "You can check your balance online or via the ONB mobile app.",
    "lost_card": "In case of a lost card, call 249-123-456-789 immediately.",
    "loan": "Loan requirements include minimum income (5000 SDG) and good credit history.",
    "transfer": "To transfer funds, use the mobile app or online banking service.",
    "new_account": "To open a new account, visit your nearest branch with your passport or national ID.",
    "interest_rates": "Interest rates on deposits range from 5% to 10% annually.",
    "branches": "Our branches are located in Omdurman, Khartoum, and Port Sudan. Visit our website for details.",
    "working_hours": "Working hours are from 8 AM to 3 PM, Sunday to Thursday.",
    "contact": "Contact us at 249-123-456-789 or via email at info@onb.sd."
}

# Menu options in both languages
MENU_AR = """
قائمة الخدمات المصرفية:
1. رصيد - استعلام عن رصيد حسابك
2. بطاقة - الإبلاغ عن بطاقة مفقودة
3. قرض - معلومات عن القروض
4. تحويل - تحويل الأموال
5. حساب - فتح حساب جديد
6. فائدة - أسعار الفائدة
7. فرع - مواقع الفروع
8. ساعات - ساعات العمل
9. اتصال - معلومات الاتصال
"""

MENU_EN = """
Banking Services Menu:
1. balance - Check your account balance
2. card - Report a lost card
3. loan - Information about loans
4. transfer - Transfer funds
5. account - Open a new account
6. interest - Interest rates
7. branch - Branch locations
8. hours - Working hours
9. contact - Contact information
"""

# Map intents to keywords (enhanced)
INTENT_KEYWORDS = {
    "balance": ["balance", "check balance", "account balance", "how much", "رصيد", "حساب", "كم المبلغ", "1"],
    "lost_card": ["lost", "card", "stolen", "missing", "فقدت", "بطاقة", "مسروقة", "ضائعة", "2"],
    "loan": ["loan", "borrow", "borrowing", "credit", "قرض", "استدانة", "إئتمان", "3"],
    "transfer": ["transfer", "send money", "payment", "تحويل", "ارسال", "دفع", "4"],
    "new_account": ["account", "open", "create", "new", "حساب", "فتح", "جديد", "إنشاء", "5"],
    "interest_rates": ["interest", "rate", "rates", "return", "فائدة", "نسبة", "عائد", "6"],
    "branches": ["branch", "location", "where", "office", "فرع", "موقع", "أين", "مكتب", "7"],
    "working_hours": ["hours", "time", "open", "close", "ساعات", "وقت", "مفتوح", "مغلق", "8"],
    "contact": ["contact", "phone", "email", "call", "اتصال", "هاتف", "بريد", "اتصل", "9"]
}

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):
    # Check for menu request
    menu_keywords = ["menu", "options", "help", "قائمة", "خيارات", "مساعدة"]
    message_lower = message.lower()
    
    for keyword in menu_keywords:
        if keyword in message_lower:
            return "menu"
    
    # Use keyword matching for intent classification
    for intent_key, keywords in INTENT_KEYWORDS.items():
        for keyword in keywords:
            if keyword.lower() in message_lower:
                return intent_key
                
    return "unknown"

def respond(message: str):
    if not message.strip():
        return {
            "ar": "الرجاء كتابة سؤالك.",
            "en": "Please type your question."
        }
    
    # Detect language
    language = detect_language(message)
    
    # If the language is neither Arabic nor English, default to English
    if language != "ar" and language != "en":
        language = "en"
    
    # Classify the user's intent using keyword matching
    intent = classify_intent(message)
    
    # Prepare responses in both languages
    responses = {
        "ar": "",
        "en": ""
    }
    
    # Special handling for menu request
    if intent == "menu":
        responses["ar"] = MENU_AR
        responses["en"] = MENU_EN
        return responses
    
    # If intent is recognized, return the corresponding response
    if intent != "unknown":
        responses["ar"] = ONB_GUIDELINES_AR.get(intent, "عذرًا، لم يتم التعرف على الخيار المحدد.")
        responses["en"] = ONB_GUIDELINES_EN.get(intent, "Sorry, the selected option was not recognized.")
    else:
        # Default response if no intent is matched - show menu
        responses["ar"] = "عذرًا، لم أفهم سؤالك. إليك قائمة بالخدمات المتاحة:" + MENU_AR
        responses["en"] = "Sorry, I didn't understand your question. Here's a menu of available services:" + MENU_EN
    
    return responses

# Custom CSS for better UI
custom_css = """
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}

.chat-message {
    padding: 1rem;
    border-radius: 10px;
    margin-bottom: 1rem;
    max-width: 80%;
}

.user-message {
    background-color: #e6f7ff;
    margin-left: auto;
    text-align: right;
}

.bot-message {
    background-color: #f0f0f0;
    margin-right: auto;
    text-align: left;
}

.bot-message-ar {
    background-color: #f0f0f0;
    margin-left: auto;
    text-align: right;
}

.header-section {
    background-color: #1a5276;
    color: white;
    padding: 1rem;
    border-radius: 10px;
    margin-bottom: 1rem;
    text-align: center;
}

.footer-section {
    font-size: 0.8rem;
    text-align: center;
    margin-top: 2rem;
    color: #666;
}

.lang-selector {
    text-align: right;
    margin-bottom: 1rem;
}

.menu-button {
    margin-top: 0.5rem;
}
"""

# Chat interface with enhanced UI
with gr.Blocks(css=custom_css) as demo:
    # Store conversation history
    state = gr.State(value=[])
    # Store selected language
    selected_lang = gr.State(value="ar")
    
    with gr.Row(elem_classes="header-section"):
        with gr.Column():
            gr.Markdown("# Omdurman National Bank | بنك أم درمان الوطني")
            gr.Markdown("### Virtual Banking Assistant | المساعد المصرفي الافتراضي")
    
    with gr.Row():
        with gr.Column(elem_classes="lang-selector"):
            language_btn = gr.Radio(
                ["العربية", "English"], 
                value="العربية", 
                label="Language | اللغة"
            )
    
    with gr.Row():
        chat_box = gr.Chatbot(elem_id="chatbox", height=400)
    
    with gr.Row():
        with gr.Column(scale=8):
            text_input = gr.Textbox(
                placeholder="Type your question here | اكتب سؤالك هنا", 
                label="", 
                elem_id="chat-input"
            )
        with gr.Column(scale=1):
            submit_btn = gr.Button("Send | إرسال", variant="primary")
    
    with gr.Row():
        with gr.Column(elem_classes="menu-button"):
            menu_btn = gr.Button("Show Menu | إظهار القائمة")
    
    with gr.Row(elem_classes="footer-section"):
        gr.Markdown("© 2025 Omdurman National Bank. All Rights Reserved. | جميع الحقوق محفوظة لبنك أم درمان الوطني ٢٠٢٥ ©")
    
    # Update language state when language is changed
    def update_language(lang):
        if lang == "العربية":
            return "ar"
        else:
            return "en"
    
    language_btn.change(
        fn=update_language,
        inputs=language_btn,
        outputs=selected_lang
    )
    
    # Handle message submission
    def on_submit(message, chat_history, lang):
        if not message.strip():
            return "", chat_history
        
        # Add user message to chat history
        chat_history.append([message, None])
        
        # Get response
        responses = respond(message)
        
        # Select response based on language
        response = responses[lang]
        
        # Update bot response in chat history
        chat_history[-1][1] = response
        
        return "", chat_history
    
    # Handle menu button click
    def show_menu(chat_history, lang):
        menu_responses = {
            "ar": MENU_AR,
            "en": MENU_EN
        }
        
        # Add system message showing the menu
        menu_text = menu_responses[lang]
        chat_history.append([None, menu_text])
        
        return chat_history
    
    # Link inputs and button to response function
    submit_btn.click(
        fn=on_submit,
        inputs=[text_input, chat_box, selected_lang],
        outputs=[text_input, chat_box]
    )
    
    # Link menu button to show menu function
    menu_btn.click(
        fn=show_menu,
        inputs=[chat_box, selected_lang],
        outputs=[chat_box]
    )
    
    # Also trigger on Enter key
    text_input.submit(
        fn=on_submit,
        inputs=[text_input, chat_box, selected_lang],
        outputs=[text_input, chat_box]
    )

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