import gradio as gr
from transformers import pipeline

# --- 1. Load the two successful models when the app starts ---
print("Loading Translation model...")
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-swc", device=-1)
print("✅ Translation model ready.")

print("Loading Summarization model...")
summarizer = pipeline("summarization", model="csebuetnlp/mT5_multilingual_XLSum", device=-1)
print("✅ Summarization model ready.")


# --- 2. Define the core functions we tested in our notebook ---
def translate_english_to_swahili(text_to_translate):
    """Takes English text and returns the Swahili translation."""
    if not text_to_translate:
        return ""
    
    # Use the loaded translator model
    result = translator(text_to_translate)
    return result[0]['translation_text']

def summarize_kiswahili_text(text_to_summarize):
    """Takes Kiswahili text and returns a concise summary."""
    if not text_to_summarize:
        return ""
        
    # Use the loaded summarizer model
    result = summarizer(text_to_summarize, max_length=60, min_length=15, do_sample=False)
    return result[0]['summary_text']


# --- 3. Build the Gradio Interface ---
# We use gr.Blocks() for a custom layout with tabs.
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Karani v0: AI Language Assistant for Kiswahili 🇰🇪🇹🇿🇺🇬")
    gr.Markdown("An AI assistant for translation and summarization tasks. Select a tab below to begin.")

    # First Tab for Translation
    with gr.Tab("Tafsiri (Translate)"):
        gr.Markdown("### English to Kiswahili Translation")
        with gr.Row():
            english_input = gr.Textbox(lines=5, label="Enter English Text Here", placeholder="e.g., Artificial intelligence is transforming our world.")
            swahili_output = gr.Textbox(lines=5, label="Swahili Result", interactive=False)
        translate_button = gr.Button("Tafsiri (Translate)")

    # Second Tab for Summarization
    with gr.Tab("Fupisha (Summarize)"):
        gr.Markdown("### Kiswahili Text Summarization")
        with gr.Row():
            kiswahili_input = gr.Textbox(lines=7, label="Enter Long Kiswahili Text Here")
            summary_output = gr.Textbox(lines=3, label="Summary", interactive=False)
        summarize_button = gr.Button("Fupisha (Summarize)")

    # --- 4. Define what happens when buttons are clicked ---
    translate_button.click(
        fn=translate_english_to_swahili,
        inputs=english_input,
        outputs=swahili_output
    )

    summarize_button.click(
        fn=summarize_kiswahili_text,
        inputs=kiswahili_input,
        outputs=summary_output
    )

# --- 5. Launch the App ---
demo.launch()