import gradio as gr 
from transformers import pipeline

# Initialize translation pipelines
en_to_fr = pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr")
fr_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-en")
en_to_ja = pipeline("translation", model="Helsinki-NLP/opus-mt-en-jap")
ja_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-jap-en")
en_to_es = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
es_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-es-en")
en_to_zh = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh")
zh_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
fr_to_es = pipeline("translation", model="Helsinki-NLP/opus-mt-fr-es")
es_to_fr = pipeline("translation", model="Helsinki-NLP/opus-mt-es-fr")

# Dictionary to map language pairs to translation pipelines
direct_translators = {
    "English-French": en_to_fr,
    "French-English": fr_to_en,
    "English-Japanese": en_to_ja,
    "Japanese-English": ja_to_en,
    "English-Spanish": en_to_es,
    "Spanish-English": es_to_en,
    "English-Mandarin": en_to_zh,
    "Mandarin-English": zh_to_en,
    "French-Spanish": fr_to_es,
    "Spanish-French": es_to_fr
}

# Two-step translation pairs (using English as pivot)
two_step_pairs = [
    "French-Mandarin", "Mandarin-French",
    "French-Japanese", "Japanese-French",
    "Spanish-Japanese", "Japanese-Spanish",
    "Spanish-Mandarin", "Mandarin-Spanish",
    "Japanese-Mandarin", "Mandarin-Japanese"
]

# Language emoji mapping for visual appeal
language_emojis = {
    "English": "🇬🇧",
    "French": "🇫🇷",
    "Japanese": "🇯🇵",
    "Spanish": "🇪🇸",
    "Mandarin": "🇨🇳"
}

def translate_two_step(text, source_lang, target_lang):
    """Perform translation in two steps using English as an intermediate language."""
    # First step: translate to English
    source_to_en_key = f"{source_lang}-English"
    source_to_en = direct_translators.get(source_to_en_key)
    if not source_to_en:
        return f"Error: Cannot translate from {source_lang} to English."
    
    # Translate source to English
    english_text = source_to_en(text)[0]["translation_text"]
    
    # Second step: translate from English to target
    en_to_target_key = f"English-{target_lang}"
    en_to_target = direct_translators.get(en_to_target_key)
    if not en_to_target:
        return f"Error: Cannot translate from English to {target_lang}."
    
    # Translate English to target
    return en_to_target(english_text)[0]["translation_text"]

def get_translator(source_lang, target_lang):
    """Get the appropriate translation method based on source and target languages."""
    key = f"{source_lang}-{target_lang}"
    
    # Direct translation available
    if key in direct_translators:
        return lambda text: direct_translators[key](text)[0]["translation_text"]
    
    # Two-step translation needed
    elif key in two_step_pairs:
        return lambda text: translate_two_step(text, source_lang, target_lang)
    
    # No translation path available
    else:
        return None

def translate(text, source_lang, target_lang):
    """Translate text from source language to target language."""
    if not text:
        return ""
    
    if source_lang == target_lang:
        return text
    
    translator = get_translator(source_lang, target_lang)
    if translator:
        return translator(text)
    else:
        return f"Translation from {source_lang} to {target_lang} is not supported."

def update_translation_info(source, target):
    """Update information about the translation path with friendly, colorful message."""
    if source == target:
        return f"<div style='padding: 10px; background-color: #ffe6e6; border-radius: 8px; text-align: center;'>✨ Both languages are the same! No translation needed. ✨</div>"
    
    key = f"{source}-{target}"
    emoji_source = language_emojis.get(source, "")
    emoji_target = language_emojis.get(target, "")
    
    if key in direct_translators:
        return f"<div style='padding: 10px; background-color: #e6ffe6; border-radius: 8px; text-align: center;'>✅ Direct translation: {emoji_source} {source} → {emoji_target} {target}</div>"
    elif key in two_step_pairs:
        return f"<div style='padding: 10px; background-color: #e6f2ff; border-radius: 8px; text-align: center;'>🔄 Two-step translation: {emoji_source} {source} → 🇬🇧 English → {emoji_target} {target}</div>"
    else:
        return f"<div style='padding: 10px; background-color: #fff2e6; border-radius: 8px; text-align: center;'>⚠️ Translation from {emoji_source} {source} to {emoji_target} {target} is not supported</div>"

def swap_languages(source, target):
    """Swap source and target languages."""
    return target, source

# Custom CSS
custom_css = """
.gradio-container {
    background: linear-gradient(135deg, #f5f7fa 0%, #e4ecfb 100%);
}
.main-div {
    border-radius: 20px;
    box-shadow: 0 10px 25px rgba(0,0,0,0.1);
    background-color: white;
    padding: 20px;
    margin: 10px;
}
.title {
    text-align: center;
    color: #4a69bd;
    font-size: 2.5em;
    margin-bottom: 10px;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
}
.subtitle {
    text-align: center;
    color: #6a89cc;
    margin-bottom: 20px;
}
.translate-button {
    background-color: #4a69bd !important;
    color: white !important;
    border-radius: 30px !important;
    padding: 10px 20px !important;
    font-size: 1.2em !important;
    transition: transform 0.3s ease !important;
}
.translate-button:hover {
    transform: scale(1.05) !important;
    background-color: #6a89cc !important;
}
.swap-button {
    background-color: #fbc531 !important;
    border-radius: 50% !important;
    width: 50px !important;
    height: 50px !important;
    padding: 0 !important;
    display: flex !important;
    justify-content: center !important;
    align-items: center !important;
    transition: transform 0.3s ease !important;
}
.swap-button:hover {
    transform: rotate(180deg) !important;
}
.language-select {
    border-radius: 10px !important;
    border: 2px solid #b2bec3 !important;
}
.text-area {
    border-radius: 10px !important;
    border: 2px solid #b2bec3 !important;
    font-size: 1.1em !important;
}
.result-area {
    border-radius: 10px !important;
    border: 2px solid #6a89cc !important;
    background-color: #f8f9fa !important;
    font-size: 1.1em !important;
}
.footer {
    text-align: center;
    margin-top: 20px;
    color: #7f8c8d;
    font-size: 0.9em;
}
.examples-container {
    margin-top: 20px;
    padding: 15px;
    background-color: #f8f9fa;
    border-radius: 10px;
    border: 2px dashed #b2bec3;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_classes="main-div"):
        gr.Markdown("# 🌍 Happy Translator 🌎", elem_classes="title")
        gr.Markdown("### Connect with the world through language! ✨", elem_classes="subtitle")
        
        with gr.Row():
            with gr.Column(scale=2):
                source_lang = gr.Dropdown(
                    choices=["English", "French", "Japanese", "Spanish", "Mandarin"],
                    label=f"From Language",
                    value="English",
                    elem_classes="language-select"
                )
            
            with gr.Column(scale=1):
                swap_btn = gr.Button("🔄", elem_classes="swap-button")
            
            with gr.Column(scale=2):
                target_lang = gr.Dropdown(
                    choices=["English", "French", "Japanese", "Spanish", "Mandarin"],
                    label=f"To Language",
                    value="French",
                    elem_classes="language-select"
                )
        
        translation_info = gr.HTML()
        
        # Update translation path info when languages change
        source_lang.change(
            update_translation_info,
            inputs=[source_lang, target_lang],
            outputs=translation_info
        )
        
        target_lang.change(
            update_translation_info,
            inputs=[source_lang, target_lang],
            outputs=translation_info
        )
        
        # Handle language swap
        swap_btn.click(
            swap_languages,
            inputs=[source_lang, target_lang],
            outputs=[source_lang, target_lang]
        ).then(
            update_translation_info,
            inputs=[source_lang, target_lang],
            outputs=translation_info
        )
        
        with gr.Row():
            with gr.Column():
                source_text = gr.Textbox(
                    label="Type your text here",
                    placeholder="Enter text to translate...",
                    lines=5,
                    elem_classes="text-area"
                )
                
            with gr.Column():
                target_text = gr.Textbox(
                    label="Translation Result",
                    lines=5,
                    elem_classes="result-area"
                )
        
        translate_btn = gr.Button("✨ Translate ✨", elem_classes="translate-button")
        
        # Fun greeting examples with language-appropriate phrases
        with gr.Column(elem_classes="examples-container"):
            gr.Markdown("### Try these fun examples! 🎉", elem_classes="subtitle")
            
            examples = [
                ["Hello, how are you today? I'm having a wonderful day!"],
                ["I would like to visit Japan someday and see the cherry blossoms."],
                ["Bonjour, comment allez-vous aujourd'hui? Je passe une journée merveilleuse!"],
                ["J'aimerais visiter le Japon un jour et voir les cerisiers en fleurs."],
                ["こんにちは、今日はお元気ですか？素晴らしい一日を過ごしています！"],
                ["いつか日本を訪れて桜を見たいです。"],
                ["Hola, ¿cómo estás hoy? ¡Estoy teniendo un día maravilloso!"],
                ["Me gustaría visitar Japón algún día y ver los cerezos en flor."],
                ["你好，今天好吗？我今天过得很愉快！"],
                ["我希望有一天能去日本看樱花。"]
            ]
            
            gr.Examples(
                examples=examples,
                inputs=source_text
            )
        
        # Set up translation function
        translate_btn.click(
            translate,
            inputs=[source_text, source_lang, target_lang],
            outputs=target_text
        )
        
        # Also translate when Enter is pressed in the source text box
        source_text.submit(
            translate,
            inputs=[source_text, source_lang, target_lang],
            outputs=target_text
        )
        
        gr.Markdown("### Made with ❤️ for language lovers everywhere", elem_classes="footer")
        
    # Display initial translation info on load
    demo.load(
        update_translation_info,
        inputs=[source_lang, target_lang],
        outputs=translation_info
    )

if __name__ == "__main__":
    demo.launch(share=True)