import gradio as gr
import json

# Translation dictionary for common phrases
TRANSLATIONS = {
    "hello": "مرحبا",
    "how are you": "كيف حالك",
    "thank you": "شكرا لك",
    "good morning": "صباح الخير",
    "good evening": "مساء الخير",
    "yes": "نعم",
    "no": "لا",
    "please": "من فضلك",
    "sorry": "آسف",
    "goodbye": "وداعا",
    "what is your name": "ما اسمك",
    "my name is": "اسمي",
    "where is the bathroom": "أين الحمام",
    "how much": "كم السعر",
    "i need help": "أحتاج مساعدة",
    "can you help me": "هل يمكنك مساعدتي",
    "i would like to book": "أود حجز",
    "set an alarm": "اضبط منبه",
    "call mom": "اتصل بأمي",
    "what time is it": "كم الساعة",
}

def translate_en_to_ar(text: str) -> str:
    """Translate English text to Arabic. Optimized for mobile-useful phrases.
    
    Use this tool when a user needs Arabic translation, especially for 
    common phrases, mobile interactions, or UAE/Arabic region communication.
    
    Args:
        text: English text to translate to Arabic
    
    Returns:
        JSON string with Arabic translation
    """
    lower = text.lower().strip()
    
    # Check exact match first
    if lower in TRANSLATIONS:
        return json.dumps({
            "english": text,
            "arabic": TRANSLATIONS[lower],
            "method": "dictionary_exact"
        }, ensure_ascii=False, indent=2)
    
    # Check partial matches
    for en, ar in TRANSLATIONS.items():
        if en in lower:
            return json.dumps({
                "english": text,
                "arabic": ar,
                "matched_phrase": en,
                "method": "dictionary_partial"
            }, ensure_ascii=False, indent=2)
    
    # Fallback: suggest using the Gemma-2B-Arabic-mobile model
    return json.dumps({
        "english": text,
        "arabic": None,
        "error": "No dictionary match. Use dispatchAI/Gemma-2B-Arabic-mobile model for full translation.",
        "model_url": "https://huggingface.co/dispatchAI/Gemma-2B-Arabic-mobile",
        "suggestion": "Load the model with: AutoModelForCausalLM.from_pretrained('dispatchAI/Gemma-2B-Arabic-mobile')"
    }, ensure_ascii=False, indent=2)

def translate_ar_to_en(text: str) -> str:
    """Translate Arabic text to English. Optimized for mobile-useful phrases.
    
    Args:
        text: Arabic text to translate to English
    
    Returns:
        JSON string with English translation
    """
    # Reverse lookup
    for en, ar in TRANSLATIONS.items():
        if ar in text:
            return json.dumps({
                "arabic": text,
                "english": en,
                "method": "dictionary"
            }, ensure_ascii=False, indent=2)
    
    return json.dumps({
        "arabic": text,
        "english": None,
        "error": "No dictionary match. Use dispatchAI/Gemma-2B-Arabic-mobile model for full translation.",
        "model_url": "https://huggingface.co/dispatchAI/Gemma-2B-Arabic-mobile"
    }, ensure_ascii=False, indent=2)

def list_phrases() -> str:
    """List all supported translation phrases.
    
    Returns:
        JSON string with all English-Arabic phrase pairs
    """
    return json.dumps(TRANSLATIONS, ensure_ascii=False, indent=2)

with gr.Blocks(title="dispatchAI Arabic Translate MCP") as demo:
    gr.Markdown("## 🌍 dispatchAI Arabic Translate (MCP Tool)")
    
    with gr.Tab("EN → AR"):
        en_input = gr.Textbox(label="English Text", placeholder="Hello, how are you?")
        ar_btn = gr.Button("Translate to Arabic", variant="primary")
        ar_output = gr.Textbox(label="Result (JSON)", lines=8)
        ar_btn.click(fn=translate_en_to_ar, inputs=en_input, outputs=ar_output)
    
    with gr.Tab("AR → EN"):
        ar_input = gr.Textbox(label="Arabic Text", placeholder="مرحبا، كيف حالك؟")
        en_btn = gr.Button("Translate to English", variant="primary")
        en_output = gr.Textbox(label="Result (JSON)", lines=8)
        en_btn.click(fn=translate_ar_to_en, inputs=ar_input, outputs=en_output)
    
    with gr.Tab("Phrases"):
        list_btn = gr.Button("List All Phrases")
        list_out = gr.Textbox(label="Phrase Dictionary (JSON)", lines=20)
        list_btn.click(fn=list_phrases, outputs=list_out)

demo.launch(mcp_server=True)