import gradio as gr
from huggingface_hub import InferenceClient
import logging  # For analytics

# Use Arabic-optimized model
client = InferenceClient("aubmindlab/aragpt2-base")

# 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."
}

# Sudanese dialect examples
SUDANESE_EXAMPLES = {
    "أزهريتك؟": "أهلاً وسهلاً! كيف يمكنني مساعدتك اليوم؟",
    "شبيك؟": "أنا بخير، شكراً! كيف يمكنني مساعدتك؟",
    "فين الفرع؟": "أقرب فرع موجود في أم درمان، الخرطوم، وبورتسودان."
}

# Set up logging for analytics
logging.basicConfig(filename='chatbot_queries.log', level=logging.INFO)

def respond(
    message: str,
    history: list[list[str]],
    system_message="أنت مساعد مصرفي لبنك أم درمان الوطني. قدم معلومات عامة فقط.",
    max_tokens=256,
    temperature=0.5,
    top_p=0.9,
):
    # Log the query
    logging.info(f"Query: {message}")
    
    # Check for Sudanese dialect phrases
    if message in SUDANESE_EXAMPLES:
        yield SUDANESE_EXAMPLES[message]
        return
    
    # Force Arabic responses
    prompt = f"باللغة العربية: {message}"
    
    # Format history in messages format
    messages = [{"role": "system", "content": system_message}]
    
    for user_msg, bot_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})

    messages.append({"role": "user", "content": prompt})

    # Generate Arabic response
    response = ""
    for chunk in client.text_generation(
        prompt=prompt,
        max_new_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        repetition_penalty=1.2,
    ):
        response += chunk
        yield response

# Omdurman National Bank-specific interface
with gr.Blocks(css=".gradio-container {direction: rtl;}") as demo:
    gr.Markdown("# <center>بنك أم درمان الوطني - المساعد المصرفي</center>")
    
    with gr.Tab("المحادثة"):
        chatbot = gr.ChatInterface(
            respond,
            examples=[
                "كيف يمكنني التحقق من رصيدي؟",  # Check balance
                "أريد الإبلاغ عن فقدان بطاقتي",  # Report lost card
                "ما هي شروط الحصول على قرض؟",  # Loan eligibility
                "ما هي ساعات العمل؟",  # Working hours
                "أين يوجد أقرب فرع؟"  # Branch locations
            ]
        )
    
    with gr.Tab("الإرشادات المصرفية"):
        gr.Markdown("## إرشادات بنك أم درمان الوطني")
        for key, value in ONB_GUIDELINES.items():
            gr.Markdown(f"**{key.capitalize()}**: {value}")

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