import gradio as gr
import requests
import json

# Language mapping from the provided data
LANGUAGE_KEY_TEXT = {
    "as": "Assamese - অসমীয়া",
    "bn": "Bangla - বাংলা",
    "brx": "Boro - बड़ो",
    "gu": "Gujarati - ગુજરાતી",
    "hi": "Hindi - हिंदी",
    "kn": "Kannada - ಕನ್ನಡ",
    "ml": "Malayalam - മലയാളം",
    "mni": "Manipuri - মিতৈলোন",
    "mr": "Marathi - मराठी",
    "or": "Odia - ଓଡ଼ିଆ",
    "raj": "Rajasthani - राजस्थानी",
    "ta": "Tamil - தமிழ்",
    "te": "Telugu - తెలుగు",
    "pa": "Panjabi - ਪੰਜਾਬੀ",
    "sa": "Sanskrit - संस्कृतम्",
    "ur": "Urdu - اُردُو",
    "en": "English - English",
    "ne": "Nepali - नेपाली",
    "si": "Sinhala - සිංහල",
}

def translate_text(text, source_lang, target_lang):
    if not text:
        return ""
        
    url = "https://demo-api.models.ai4bharat.org/inference/translation/v2"
    
    headers = {
        'Content-Type': 'application/json',
        'Accept': '*/*',
    }
    
    payload = {
        "controlConfig": {
            "dataTracking": True
        },
        "input": [{"source": text}],
        "config": {
            "serviceId": "",
            "language": {
                "sourceLanguage": source_lang,
                "targetLanguage": target_lang,
                "targetScriptCode": None,
                "sourceScriptCode": None
            }
        }
    }
    
    try:
        response = requests.post(url, headers=headers, json=payload)
        response.raise_for_status()
        result = response.json()
        return result['output'][0]['target']
    except Exception as e:
        return f"Error: {str(e)}"

def copy_to_clipboard(text):
    return text

def create_interface():
    with gr.Blocks(css="footer {display: none !important}") as interface:
        gr.Markdown(
        """
        # AI4Bharat Translation
        Translate text between multiple Indian languages and English
        """
        )
        
        with gr.Row():
            with gr.Column():
                source_lang = gr.Dropdown(
                    choices=list(LANGUAGE_KEY_TEXT.keys()),
                    value="en",
                    label="Source Language"
                )
                source_text = gr.TextArea(
                    label="Source Text",
                    placeholder="Enter text to translate...",
                    lines=5
                )
                
            with gr.Column():
                target_lang = gr.Dropdown(
                    choices=list(LANGUAGE_KEY_TEXT.keys()),
                    value="hi",
                    label="Target Language"
                )
                translated_text = gr.TextArea(
                    label="Translated Text",
                    lines=5,
                    interactive=False
                )
                
        with gr.Row():
            translate_btn = gr.Button("Translate")
            copy_btn = gr.Button("📋 Copy Translation")
            clear_btn = gr.Button("🗑️ Clear")
            
        # Swap language button
        swap_btn = gr.Button("🔄 Swap Languages")
        
        def swap_languages(src_lang, tgt_lang, src_text, tgt_text):
            return tgt_lang, src_lang, tgt_text, src_text
            
        def clear_fields():
            return "", ""
            
        # Event handlers
        translate_btn.click(
            fn=translate_text,
            inputs=[source_text, source_lang, target_lang],
            outputs=translated_text
        )
        
        swap_btn.click(
            fn=swap_languages,
            inputs=[source_lang, target_lang, source_text, translated_text],
            outputs=[source_lang, target_lang, source_text, translated_text]
        )
        
        clear_btn.click(
            fn=clear_fields,
            inputs=[],
            outputs=[source_text, translated_text]
        )
        
        # Add copy button functionality
        copy_btn.click(
            fn=copy_to_clipboard,
            inputs=[translated_text],
            outputs=[translated_text],
            js="async (text) => { await navigator.clipboard.writeText(text); return text; }"
        )
        
    return interface

# Launch the interface
if __name__ == "__main__":
    demo = create_interface()
    demo.launch()