import gradio as gr
from transformers import pipeline
import torch
from huggingface_hub import login

# Placeholder for the model pipeline
pipe = None

def validate_and_load(token):
    """Handles the 'login' logic and loads the model."""
    global pipe
    try:
        # Authenticate
        login(token=token)
        
        # Load the model (Moving this here ensures it only loads after login)
        # Using bfloat16 and device_map for professional performance
        pipe = pipeline(
            "text-generation", 
            model="ibm-granite/granite-3.3-2b-base", 
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )
        # Return a success message and toggle UI visibility
        return gr.update(visible=False), gr.update(visible=True), "✅ Authentication Successful! Model Loaded."
    except Exception as e:
        return gr.update(visible=True), gr.update(visible=False), f"❌ Error: {str(e)}"

def translate_text(text, target_language):
    if pipe is None:
        return "Please login first."
    
    # Prompt engineering for the Base model
    prompt = f"Translate the following English text to {target_language}:\nEnglish: {text}\n{target_language}:"
    
    outputs = pipe(
        prompt, 
        max_new_tokens=150, 
        do_sample=False, 
        return_full_text=False
    )
    return outputs[0]['generated_text'].strip()

# --- UI DESIGN ---
with gr.Blocks(theme=gr.themes.Soft(), title="Granite Translator Pro") as demo:
    
    # Header Section
    gr.Markdown("""
    # 🌐 Granite Multi-Lingual Pro
    ### Enterprise-grade translation powered by IBM Granite 3.3 2B
    """)
    
    # 1. AUTHENTICATION SECTION (Visible by default)
    with gr.Column(visible=True) as auth_section:
        gr.Markdown("### 🔐 Authentication Required")
        hf_token = gr.Textbox(
            label="Hugging Face Access Token", 
            placeholder="hf_...", 
            type="password",
            info="Enter your read-access token to begin."
        )
        login_btn = gr.Button("Initialize Application", variant="primary")
        status_msg = gr.Markdown()

    # 2. MAIN APPLICATION SECTION (Hidden by default)
    with gr.Column(visible=False) as main_app:
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(
                    label="Input Text (English)", 
                    placeholder="Type something here...",
                    lines=5
                )
                target_lang = gr.Dropdown(
                    label="Target Language", 
                    choices=["Hindi (हिन्दी)", "Gujarati (ગુજરાતી)", "Spanish", "French", "German"], 
                    value="Hindi (हिन्दी)"
                )
                translate_btn = gr.Button("Translate Now", variant="primary")
            
            with gr.Column():
                output_text = gr.Textbox(label="Translated Result", lines=8, interactive=False)
        
        gr.ClearButton([input_text, output_text])

    # --- LOGIC FLOW ---
    # When login is clicked, validate token and switch views
    login_btn.click(
        fn=validate_and_load, 
        inputs=[hf_token], 
        outputs=[auth_section, main_app, status_msg]
    )

    # Translation trigger
    translate_btn.click(
        fn=translate_text, 
        inputs=[input_text, target_lang], 
        outputs=[output_text]
    )

if __name__ == "__main__":
    demo.launch()