import gradio as gr
import requests
import html
import os

# LibreTranslate instance
API_URL = "https://imsidag-community-libretranslate-kabyle.hf.space/translate"
API_KEY = os.environ["API_KEY"]

# Language codes
langs = {
    "English": "en",
    "French": "fr",
    "Italian": "it",
    "Occitan (26000)": "oc_comp2",
    "Taqbaylit (LTCommunity)": "kab",
    "Taqbaylit (Tasenselkimt)": "kab_comp",
    "Taqbaylit (51000)": "kab_comp2",
    "Taqbaylit (OS)": "kab_os",
    "Taqbaylit (Google)": "google",
    "Taqbaylit (All Variants)": "kab_all"
}

kabyle_variants = {
    "Taqbaylit (LTCommunity)": "kab",
    "Taqbaylit (Tasenselkimt)": "kab_comp",
    "Taqbaylit (51000)": "kab_comp2",
    "Taqbaylit (OS)": "kab_os",
    "Taqbaylit (Google)": "google"
}

def highlight_differences(translations_dict):
    token_lists = {k: v.split() for k, v in translations_dict.items()}
    max_len = max(len(tokens) for tokens in token_lists.values())
    aligned = {k: [] for k in token_lists.keys()}

    for i in range(max_len):
        words_at_i = [token_lists[k][i] if i < len(token_lists[k]) else "" for k in token_lists]
        if len(set(words_at_i)) == 1:
            for k in token_lists:
                aligned[k].append(words_at_i[0])
        else:
            for k in token_lists:
                word = token_lists[k][i] if i < len(token_lists[k]) else ""
                # Use gold color for highlights that works in both themes
                highlighted = f"<span style='color:#FFD700; font-weight:bold'>{word}</span>"
                aligned[k].append(highlighted)

    return {k: " ".join(tokens) for k, tokens in aligned.items()}


def translate(text, source_lang, target_lang):
    source_code = langs.get(source_lang, "en")
    target_code = langs.get(target_lang)

    if target_code == "kab_all":
        # Escape and format the original text with dark-theme friendly styling
        escaped_text = html.escape(text).replace('\n', '<br>')
        header = f"""
        <div style='
            margin-bottom: 25px; 
            padding: 15px;
            border-radius: 8px;
            background-color: rgba(120, 120, 120, 0.15);
            border: 1px solid rgba(255, 255, 255, 0.1);
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        '>
            <div style='font-size: 0.9em; opacity: 0.8; margin-bottom: 5px;'>ORIGINAL TEXT</div>
            <div style='font-size: 1.1em;'>{escaped_text}</div>
            <div style='margin-top: 8px; font-size: 0.85em; opacity: 0.7;'>
                From: {source_lang} → To: {target_lang}
            </div>
        </div>
        """
        
        translations_dict = {}
        alternatives_dict = {}

        for name, code in kabyle_variants.items():
            if code == "google":
                # Fallback source for Google
                google_source = source_code if source_code in ["en", "fr", "it", "es"] else "fr"
                params = {
                    "engine": "google",
                    "from": google_source,
                    "to": "ber-Latn",
                    "text": text
                }
                try:
                    response = requests.get("https://mozhi.pussthecat.org/api/translate", params=params)
                    result = response.json()
                    translations_dict[name] = result.get("translated-text", "<i>No result</i>")
                    alternatives_dict[name] = ""
                except Exception as e:
                    translations_dict[name] = f"<i>Google Error: {str(e)}</i>"
                    alternatives_dict[name] = ""
            else:
                data = {
                    "q": text,
                    "source": source_code,
                    "target": code,
                    "format": "text",
                    "alternatives": 3,
                    "api_key": API_KEY
                }
                try:
                    response = requests.post(API_URL, data=data, headers={"Content-Type": "application/x-www-form-urlencoded"})
                    result = response.json()

                    if "translatedText" in result:
                        translations_dict[name] = result["translatedText"]
                        alternatives_dict[name] = ", ".join(result.get("alternatives", []))
                    elif "error" in result:
                        translations_dict[name] = f"<i>API Error: {result['error']}</i>"
                        alternatives_dict[name] = ""
                    else:
                        translations_dict[name] = "<i>Unexpected response</i>"
                        alternatives_dict[name] = ""
                except Exception as e:
                    translations_dict[name] = f"<i>Exception: {str(e)}</i>"
                    alternatives_dict[name] = ""

        highlighted = highlight_differences(translations_dict)

        outputs = []
        for name in kabyle_variants:
            output = f"<div style='margin-bottom: 15px;'><b>{name}</b>:<br>Translation: {highlighted[name]}"
            if alternatives_dict.get(name):
                output += f"<br>Alternatives: <span style='opacity:0.8;'>{alternatives_dict[name]}</span>"
            output += "</div>"
            outputs.append(output)

        return header + "".join(outputs)

    elif target_code == "google":
        google_source = source_code if source_code in ["en", "fr", "it", "es"] else "fr"
        params = {
            "engine": "google",
            "from": google_source,
            "to": "ber-Latn",
            "text": text
        }
        try:
            response = requests.get("https://mozhi.pussthecat.org/api/translate", params=params)
            result = response.json()
            return f"<b>Translation:</b> {result.get('translated-text', '<i>No result</i>')}"
        except Exception as e:
            return f"<i>Google Error: {str(e)}</i>"

    else:
        data = {
            "q": text,
            "source": source_code,
            "target": target_code,
            "format": "text",
            "alternatives": 3,
            "api_key": API_KEY
        }

        try:
            response = requests.post(API_URL, data=data, headers={"Content-Type": "application/x-www-form-urlencoded"})
            result = response.json()

            if "translatedText" in result:
                output = f"<b>Translation:</b> {result['translatedText']}"
                if result.get("alternatives"):
                    alt = ", ".join(result["alternatives"])
                    output += f"<br><b>Alternatives:</b> {alt}"
                return output
            elif "error" in result:
                return f"<i>API Error: {result['error']}</i>"
            else:
                return f"<i>Unexpected response: {result}</i>"
        except Exception as e:
            return f"<i>Exception: {str(e)}</i>"


# Gradio UI with dark theme friendly styling
with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.Markdown("# Kabyle Translator")
    gr.Markdown("Compare Kabyle translations from LibreTranslate (variants) and Google Translate via Mozhi. Differences are highlighted.")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Enter text to translate", lines=3)
            source_lang = gr.Dropdown(choices=list(langs.keys()), label="From language", value="English")
            target_lang = gr.Dropdown(choices=list(langs.keys()), label="To language", value="Taqbaylit (All Variants)")
            submit_btn = gr.Button("Translate")
        
        with gr.Column():
            output = gr.HTML(label="Translation Results")
    
    submit_btn.click(
        fn=translate,
        inputs=[input_text, source_lang, target_lang],
        outputs=output
    )

demo.launch()