import gradio as gr
from transformers import pipeline

# Preload models
summarizer = pipeline("summarization")
sentiment_analyzer = pipeline("sentiment-analysis", model="cardiffnlp/twitter-roberta-base-sentiment")
translator_hi = pipeline("translation", model="Helsinki-NLP/opus-mt-en-hi")
translator_fr = pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr")
translator_de = pipeline("translation", model="Helsinki-NLP/opus-mt-en-de")
translator_es = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
translator_ta = pipeline("translation", model="facebook/nllb-200-distilled-600M", src_lang="eng_Latn", tgt_lang="tam_Taml")
speech_to_text = pipeline("automatic-speech-recognition", model="openai/whisper-small")
question_generator = pipeline("text2text-generation", model="valhalla/t5-base-e2e-qg")

# Functional logic
def summarize(text):
    return summarizer(text, max_length=60, min_length=20, do_sample=False)[0]['summary_text']

def analyze_sentiment(text):
    result = sentiment_analyzer(text)[0]
    label = result["label"]
    if label == "LABEL_1":
        return "Neutral"
    elif label == "LABEL_2":
        return "Positive"
    else:
        return "Negative"

def translate(text, lang):
    if lang == "Tamil":
        return translator_ta(text)[0]["translation_text"]
    elif lang == "Hindi":
        return translator_hi(text)[0]["translation_text"]
    elif lang == "French":
        return translator_fr(text)[0]["translation_text"]
    elif lang == "German":
        return translator_de(text)[0]["translation_text"]
    elif lang == "Spanish":
        return translator_es(text)[0]["translation_text"]
    else:
        return "Unsupported Language"

def transcribe(audio):
    return speech_to_text(audio)["text"]

def generate_questions(text):
    prompt = "generate questions: " + text
    result = question_generator(prompt, max_length=256, do_sample=False)
    return result[0]["generated_text"]

# UI App
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🌐 AI Buddy — Multi-task App l")
    gr.Markdown("Choose a task below:")

    # Navigation buttons
    with gr.Row():
        task = gr.State(value="")
        btn_summ = gr.Button("Text Summarization")
        btn_sent = gr.Button("Sentiment Analysis")
        btn_trans = gr.Button("Translation")
        btn_speech = gr.Button("Speech-to-Text")
        btn_qgen = gr.Button("Question Generation")

    # Tab sections
    with gr.Column(visible=False) as summarize_tab:
        text = gr.Textbox(label="Enter paragraph", lines=10)
        out = gr.Textbox(label="Summary", lines=4)
        gr.Button("Summarize").click(summarize, inputs=text, outputs=out)

    with gr.Column(visible=False) as sentiment_tab:
        sent_in = gr.Textbox(label="Enter sentence", lines=3)
        sent_out = gr.Textbox(label="Sentiment")
        gr.Button("Analyze Sentiment").click(analyze_sentiment, inputs=sent_in, outputs=sent_out)

    with gr.Column(visible=False) as translate_tab:
        tran_in = gr.Textbox(label="Enter English text", lines=3)
        lang = gr.Dropdown(["Tamil", "Hindi", "French", "German", "Spanish"], value="Tamil", label="Language")
        tran_out = gr.Textbox(label="Translated text", lines=3)
        gr.Button("Translate").click(translate, inputs=[tran_in, lang], outputs=tran_out)

    with gr.Column(visible=False) as speech_tab:
        audio = gr.Audio(type="filepath", label="Record or upload")
        speech_out = gr.Textbox(label="Recognized Text")
        gr.Button("Convert Speech to Text").click(transcribe, inputs=audio, outputs=speech_out)

    with gr.Column(visible=False) as question_tab:
        ques_in = gr.Textbox(label="Enter a paragraph", lines=10)
        ques_out = gr.Textbox(label="Generated Questions", lines=10)
        gr.Button("Generate Questions").click(generate_questions, inputs=ques_in, outputs=ques_out)

    # Logic to show/hide tabs
    def show_tab(tab_name):
        return [
            gr.update(visible=(tab_name == "summarize")),
            gr.update(visible=(tab_name == "sentiment")),
            gr.update(visible=(tab_name == "translate")),
            gr.update(visible=(tab_name == "speech")),
            gr.update(visible=(tab_name == "question")),
        ]

    # Shared hidden input to control tab switching
    hidden_tab_name = gr.Textbox(value="", visible=False)

    btn_summ.click(lambda: "summarize", outputs=hidden_tab_name).then(
        show_tab, inputs=hidden_tab_name, outputs=[summarize_tab, sentiment_tab, translate_tab, speech_tab, question_tab]
    )
    btn_sent.click(lambda: "sentiment", outputs=hidden_tab_name).then(
        show_tab, inputs=hidden_tab_name, outputs=[summarize_tab, sentiment_tab, translate_tab, speech_tab, question_tab]
    )
    btn_trans.click(lambda: "translate", outputs=hidden_tab_name).then(
        show_tab, inputs=hidden_tab_name, outputs=[summarize_tab, sentiment_tab, translate_tab, speech_tab, question_tab]
    )
    btn_speech.click(lambda: "speech", outputs=hidden_tab_name).then(
        show_tab, inputs=hidden_tab_name, outputs=[summarize_tab, sentiment_tab, translate_tab, speech_tab, question_tab]
    )
    btn_qgen.click(lambda: "question", outputs=hidden_tab_name).then(
        show_tab, inputs=hidden_tab_name, outputs=[summarize_tab, sentiment_tab, translate_tab, speech_tab, question_tab]
    )

demo.launch()