import os
import gradio as gr
from transformers import pipeline

# === Load Hugging Face token (only needed if Space is private) ===
auth_token = os.getenv("HF_API_TOKEN")

# === Choose your model (pick one that works best for you) ===
MODEL_NAME = "Helsinki-NLP/opus-mt-en-bn"   # English → Bengali
# MODEL_NAME = "shhossain/opus-mt-en-to-bn"  # Alternate
# MODEL_NAME = "facebook/nllb-200-distilled-600M"  # Multi-language

# === Load translation pipeline ===
try:
    translator = pipeline(
        "translation",
        model=MODEL_NAME,
        token=auth_token  # Works for private Spaces
    )
except Exception as e:
    translator = None
    print("❌ Error loading model:", e)


# === Define translation function ===
def translate_text(text, direction):
    if not translator:
        return "⚠️ Model not loaded. Check HF_API_TOKEN for private Spaces."

    try:
        if direction == "English → Bengali":
            result = translator(text, src="en", tgt="bn")
        elif direction == "Bengali → English":
            # For reverse, swap model if needed
            rev_translator = pipeline(
                "translation",
                model="Helsinki-NLP/opus-mt-bn-en",
                token=auth_token
            )
            result = rev_translator(text)
        else:
            return "⚠️ Unknown direction."
        return result[0]['translation_text']
    except Exception as e:
        return f"⚠️ Translation failed: {str(e)}"


# === Gradio UI ===
with gr.Blocks(title="Private English ↔ Bengali Translator") as demo:
    gr.Markdown("## 🌐 English ↔ Bengali Translator (Private Space Ready)")

    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Enter text",
                placeholder="Type here...",
                lines=5
            )
            direction = gr.Radio(
                ["English → Bengali", "Bengali → English"],
                value="English → Bengali",
                label="Select Translation Direction"
            )
            translate_btn = gr.Button("Translate", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(
                label="Translation Result",
                lines=5
            )

    translate_btn.click(
        fn=translate_text,
        inputs=[input_text, direction],
        outputs=[output_text]
    )

# === Launch app ===
if __name__ == "__main__":
    demo.launch()