import gradio as gr
from transformers import pipeline
import re

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

# Load English NLP model for intent classification
intent_classifier_en = pipeline("text-classification", model="facebook/bart-large-mnli")

# Load language detection model
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."
}

# Map intents to responses
INTENT_KEYWORDS = {
    "balance": ["balance", "رصيد", "حساب"],
    "lost_card": ["lost", "card", "stolen", "فقدت", "بطاقة", "مسروقة"],
    "loan": ["loan", "borrow", "قرض", "استدانة"],
    "transfer": ["transfer", "send money", "تحويل", "ارسال"],
    "new_account": ["account", "open", "حساب", "فتح"],
    "interest_rates": ["interest", "rate", "فائدة", "نسبة"],
    "branches": ["branch", "location", "فرع", "موقع"],
    "working_hours": ["hours", "time", "ساعات", "وقت"],
    "contact": ["contact", "phone", "email", "اتصال", "هاتف", "بريد"]
}

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, language: str):
    # Use appropriate classifier based on language
    if language == "ar":
        # For Arabic
        result = intent_classifier_ar(message)
        intent = result[0]['label']
    else:
        # For English
        result = intent_classifier_en(message, candidate_labels=list(INTENT_KEYWORDS.keys()))
        intent = result["labels"][0]
    
    # Fallback to keyword matching
    if intent not in INTENT_KEYWORDS:
        for intent_key, keywords in INTENT_KEYWORDS.items():
            for keyword in keywords:
                if re.search(r'\b' + re.escape(keyword.lower()) + r'\b', message.lower()):
                    return intent_key
        return "unknown"
    
    return intent

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
    intent = classify_intent(message, language)
    
    # Prepare responses in both languages
    responses = {
        "ar": "",
        "en": ""
    }
    
    # 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
        ar_options = ", ".join(list(ONB_GUIDELINES_AR.keys()))
        en_options = ", ".join(list(ONB_GUIDELINES_EN.keys()))
        
        responses["ar"] = f"عذرًا، لم أفهم سؤالك. الرجاء إعادة الصياغة أو اختيار أحد الخيارات التالية: {ar_options}"
        responses["en"] = f"Sorry, I didn't understand your question. Please rephrase or choose one of the following options: {en_options}"
    
    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;
}
"""

# 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(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
    
    # 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]
    )
    
    # 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
    )