# app.py
import os
import traceback
import gradio as gr

# Try import transformers; if missing the app will show a clear error message in logs.
try:
    from transformers import pipeline
except Exception as e:
    pipeline = None
    print("transformers import error:", e)

# --- CONFIG: use lightweight OPUS-MT models for CPU-friendly translation ---
MODEL_EN_TO_BN = "shhossain/opus-mt-en-to-bn"
MODEL_BN_TO_EN = "Helsinki-NLP/opus-mt-bn-en"

# cache of loaded pipelines
_loaded = {}

def safe_load_pipeline(model_name):
    """Load a translation pipeline lazily and return a tuple (success, message)."""
    global _loaded
    if pipeline is None:
        return False, "transformers not available - check requirements.txt"
    if model_name in _loaded:
        return True, f"model already loaded: {model_name}"
    try:
        # device=-1 ensures CPU usage; set max_length moderately
        p = pipeline("translation", model=model_name, device=-1, max_length=512)
        _loaded[model_name] = p
        return True, f"Loaded {model_name}"
    except Exception as e:
        # log the full stack to Space logs so you can copy it
        print("Exception while loading model:", model_name)
        traceback.print_exc()
        return False, f"Failed to load {model_name}: {str(e)}"

def translate_text(text: str, direction: str):
    """
    Main translation function used by the Gradio UI.
    returns: translation, status, debug_info
    """
    if not text or not text.strip():
        return "", "⚠️ Enter text to translate", ""
    try:
        model_name = MODEL_EN_TO_BN if direction == "English → Bengali" else MODEL_BN_TO_EN

        # Try local lazy-load first
        ok, msg = safe_load_pipeline(model_name)
        if not ok:
            # If local load fails, provide immediate dummy fallback so buttons respond
            print("Local model load failed:", msg)
            return text, f"⚠️ Local model load failed: {msg}. Showing fallback (identity) translation.", "Fallback used: returning input as output."

        # Use the loaded pipeline
        translator = _loaded.get(model_name)
        result = translator(text, max_length=512)
        # result is often list of dicts
        if isinstance(result, list) and len(result) > 0:
            r0 = result[0]
            if isinstance(r0, dict):
                translated = r0.get("translation_text") or r0.get("generated_text") or str(r0)
            else:
                translated = str(r0)
        else:
            translated = str(result)
        status = f"✅ Translated using {model_name}"
        analysis = f"Input words: {len(text.split())}; Output words: {len(translated.split())}"
        return translated, status, analysis

    except Exception as e:
        # If anything crashes, show a simple fallback so UI remains responsive
        tb = traceback.format_exc()
        print("Translation exception:", tb)
        return "", f"❌ Error during translation: {str(e)}", tb

# --- Small, responsive CSS for mobile: keep layout simple ---
custom_css = """
/* Make UI mobile-friendly and readable */
.gradio-container { padding: 12px !important; max-width: 900px; margin: auto; font-family: 'Times New Roman', serif; }
.gradio-row { gap: 8px !important; }
textarea, input[type="text"] { font-size: 18px !important; }
.gr-button { font-size: 18px !important; padding: 12px 18px !important; }
"""

# --- Build Gradio UI with Blocks for responsive layout ---
with gr.Blocks(css=custom_css, title="English ↔ Bengali — Fast Translator") as demo:
    gr.Markdown("## English ↔ Bengali — Fast Translator\nUsing small OPUS-MT models (CPU friendly). The app lazy-loads models so Space won't crash. If a model fails to load the app will return a fallback so buttons still work.")
    with gr.Row():
        direction = gr.Radio(label="Direction", choices=["English → Bengali", "Bengali → English"], value="English → Bengali")
        swap_btn = gr.Button("Swap")

    input_text = gr.Textbox(label="Input text", placeholder="Type a sentence here (English or Bengali)...", lines=4)
    translate_btn = gr.Button("Translate", variant="primary")
    with gr.Row():
        out_translation = gr.Textbox(label="Translation", lines=4)
    with gr.Row():
        out_status = gr.Textbox(label="Status / Tips", lines=1)
    out_analysis = gr.Textbox(label="Analysis / Debug", lines=3)

    # swap behavior
    def do_swap(cur):
        return "Bengali → English" if cur == "English → Bengali" else "English → Bengali"
    swap_btn.click(do_swap, inputs=direction, outputs=direction)

    # main click hook
    translate_btn.click(translate_text, inputs=[input_text, direction], outputs=[out_translation, out_status, out_analysis])

    # example quick buttons
    with gr.Row():
        ex1 = gr.Button("Hello, how are you?")
        ex2 = gr.Button("Ami bhalo achi")
        ex3 = gr.Button("Where is the market?")

    ex1.click(lambda: "Hello, how are you?", outputs=input_text)
    ex2.click(lambda: "Ami bhalo achi", outputs=input_text)
    ex3.click(lambda: "Where is the market?", outputs=input_text)

# Launch
if __name__ == "__main__":
    # debug=True prints logs to the container console
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)), debug=True)