import gradio as gr
from transformers import pipeline
import torch

# ----------------------------------------
# Supported Languages
# ----------------------------------------
languages = [
    "English",
    "Hindi",
    "Tamil",
    "Telugu",
    "Kannada",
    "Malayalam",
    "Bengali",
    "Marathi",
    "Gujarati",
    "Punjabi",
    "Urdu"
]

# ----------------------------------------
# Translation Function
# ----------------------------------------
def translate_text(hf_token, source_lang, target_lang, input_text):

    if not hf_token.strip():
        return "Please enter Hugging Face Access Token."

    if not input_text.strip():
        return "Please enter text."

    try:

        # Load IBM Granite model
        pipe = pipeline(
            "text-generation",
            model="ibm-granite/granite-3.3-2b-base",
            token=hf_token,
            device_map="auto",
            torch_dtype=torch.float16
        )

        # Translation Prompt
        prompt = f"""
You are an expert multilingual translator.

Translate the below text from {source_lang} to {target_lang}.

Only provide translated text.

Text:
{input_text}
"""

        # Generate Translation
        result = pipe(
            prompt,
            max_new_tokens=200,
            temperature=0.2
        )

        translated_text = result[0]["generated_text"]

        # Remove original prompt if generated
        translated_text = translated_text.replace(prompt, "").strip()

        return translated_text

    except Exception as e:
        return f"Error: {str(e)}"


# ----------------------------------------
# Professional UI Styling
# ----------------------------------------
custom_css = """
body {
    background: linear-gradient(to right, #0f2027, #203a43, #2c5364);
}

.gradio-container {
    font-family: 'Segoe UI', sans-serif;
}

.main-title {
    text-align: center;
    font-size: 42px;
    font-weight: bold;
    color: white;
}

.sub-title {
    text-align: center;
    color: #dfe6e9;
    font-size: 18px;
    margin-bottom: 20px;
}

textarea {
    border-radius: 12px !important;
}

footer {
    visibility: hidden;
}
"""

# ----------------------------------------
# Gradio Interface
# ----------------------------------------
with gr.Blocks(
    theme=gr.themes.Soft(),
    css=custom_css
) as demo:

    gr.Markdown(
        """
        <div class="main-title">
            🌍 AI Multilingual Translator
        </div>

        <div class="sub-title">
            IBM Granite + Hugging Face + Gradio
        </div>
        """
    )

    with gr.Row():

        with gr.Column(scale=1):

            hf_token = gr.Textbox(
                label="🔑 Hugging Face Access Token",
                placeholder="Paste your HF token here",
                type="password"
            )

            source_lang = gr.Dropdown(
                choices=languages,
                value="English",
                label="🌐 Source Language"
            )

            target_lang = gr.Dropdown(
                choices=languages,
                value="Hindi",
                label="🎯 Target Language"
            )

        with gr.Column(scale=2):

            input_text = gr.Textbox(
                label="📝 Enter Text",
                lines=8,
                placeholder="Type text to translate..."
            )

            output_text = gr.Textbox(
                label="✅ Translated Output",
                lines=8
            )

    translate_btn = gr.Button(
        "🚀 Translate",
        variant="primary"
    )

    translate_btn.click(
        fn=translate_text,
        inputs=[
            hf_token,
            source_lang,
            target_lang,
            input_text
        ],
        outputs=output_text
    )

# Launch App
demo.launch(share=True)