import gradio as gr
from transformers import pipeline

# Load pipelines with small, CPU-friendly models
summarizer = pipeline("summarization", model="t5-small")
qa = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
sentiment = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
translator_en_id = pipeline("translation_en_to_id", model="Helsinki-NLP/opus-mt-en-id")
translator_id_en = pipeline("translation_id_to_en", model="Helsinki-NLP/opus-mt-id-en")
ner = pipeline("ner", grouped_entities=True)

# New: Text Generation pipeline
text_generator = pipeline("text-generation", model="gpt2")

# Functions for each feature
def summarize_text(text):
    if not text.strip():
        return "Please enter text to summarize."
    summary = summarizer(text, max_length=100, min_length=25, do_sample=False)
    return summary[0]['summary_text']

def answer_question(context, question):
    if not context.strip() or not question.strip():
        return "Please provide both context and a question."
    result = qa(question=question, context=context)
    return result['answer']

def analyze_sentiment(text):
    if not text.strip():
        return "Please enter text for sentiment analysis."
    result = sentiment(text)
    label = result[0]['label']
    score = result[0]['score']
    return f"Sentiment: {label} (confidence: {score:.2f})"

def translate_en_to_id(text):
    if not text.strip():
        return "Please enter English text to translate."
    translation = translator_en_id(text)
    return translation[0]['translation_text']

def translate_id_to_en(text):
    if not text.strip():
        return "Please enter Malay text to translate."
    translation = translator_id_en(text)
    return translation[0]['translation_text']

def extract_entities(text):
    if not text.strip():
        return "Please enter text for entity recognition."
    entities = ner(text)
    if not entities:
        return "No entities found."
    formatted = "\n".join([f"{e['entity_group']}: {e['word']}" for e in entities])
    return formatted

# New: Text Generation function
def generate_text(prompt):
    if not prompt.strip():
        return "Please enter a starting phrase."
    result = text_generator(prompt, max_length=100, num_return_sequences=1)
    return result[0]["generated_text"]

# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Multi-Function AI Assistant")
    gr.Markdown(
        "This app provides **Text Summarization**, **Question Answering**, **Sentiment Analysis**, "
        "**Translation**, **Named Entity Recognition**, and **Text Generation**. "
        "All models run efficiently on CPU, suitable for free tier deployment."
    )

    with gr.Tab("Summarization"):
        summ_input = gr.Textbox(label="Enter text to summarize", lines=5, placeholder="Paste your text here...")
        summ_output = gr.Textbox(label="Summary", lines=3)
        summ_button = gr.Button("Summarize")
        summ_button.click(summarize_text, inputs=summ_input, outputs=summ_output)

    with gr.Tab("Question Answering"):
        qa_context = gr.Textbox(label="Context", lines=5, placeholder="Enter context text here...")
        qa_question = gr.Textbox(label="Question", placeholder="Enter your question here...")
        qa_answer = gr.Textbox(label="Answer", lines=2)
        qa_button = gr.Button("Get Answer")
        qa_button.click(answer_question, inputs=[qa_context, qa_question], outputs=qa_answer)

    with gr.Tab("Sentiment Analysis"):
        sent_input = gr.Textbox(label="Enter text for sentiment analysis", lines=4, placeholder="Type text here...")
        sent_output = gr.Textbox(label="Sentiment Result")
        sent_button = gr.Button("Analyze Sentiment")
        sent_button.click(analyze_sentiment, inputs=sent_input, outputs=sent_output)

    with gr.Tab("Named Entity Recognition"):
        ner_input = gr.Textbox(label="Enter text for entity recognition", lines=5, placeholder="Type text here...")
        ner_output = gr.Textbox(label="Entities Found", lines=6)
        ner_button = gr.Button("Extract Entities")
        ner_button.click(extract_entities, inputs=ner_input, outputs=ner_output)

    with gr.Tab("Translation (EN → ID)"):
        trans_input = gr.Textbox(label="Enter English text", lines=4, placeholder="Type English text here...")
        trans_output = gr.Textbox(label="Indonesian Translation", lines=3)
        trans_button = gr.Button("Translate")
        trans_button.click(translate_en_to_id, inputs=trans_input, outputs=trans_output)

    with gr.Tab("Translation (ID → EN)"):
        trans_input = gr.Textbox(label="Enter Indonesian text", lines=4, placeholder="Type Indonesian text here...")
        trans_output = gr.Textbox(label="English Translation", lines=3)
        trans_button = gr.Button("Translate")
        trans_button.click(translate_id_to_en, inputs=trans_input, outputs=trans_output)

    # New: Text Generation Tab
    with gr.Tab("Text Generation"):
        gen_input = gr.Textbox(label="Enter a starting phrase", lines=3, placeholder="e.g., Once upon a time")
        gen_output = gr.Textbox(label="Generated Text", lines=6)
        gen_button = gr.Button("Generate")
        gen_button.click(generate_text, inputs=gen_input, outputs=gen_output)

    gr.Markdown(
        "### Notes:\n"
        "- Input text length is limited for performance.\n"
        "- All models are CPU-optimized for free tier deployment.\n"
        "- Refresh the page to reset the app.\n"
        "- Feel free to explore each tab for different AI functionalities."
    )

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