import gradio as gr
from transformers import pipeline

# === Pipelines Setup ===
grammar_corrector = pipeline("text2text-generation", model="vennify/t5-base-grammar-correction")
paraphraser = pipeline("text2text-generation", model="eugenesiow/bart-paraphrase")
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
translator_en2ur = pipeline("translation_en_to_ur", model="Helsinki-NLP/opus-mt-en-ur")
translator_ur2en = pipeline("translation_ur_to_en", model="Helsinki-NLP/opus-mt-ur-en")

# === Functions ===
def correct_grammar(text):
    result = grammar_corrector(text, max_length=256, num_beams=4, clean_up_tokenization_spaces=True)
    return result[0]['generated_text']

def paraphrase_text(text):
    result = paraphraser("paraphrase: " + text, max_length=256, num_beams=5, clean_up_tokenization_spaces=True)
    return result[0]['generated_text']

def summarize_text(text):
    result = summarizer(text, max_length=130, min_length=30, do_sample=False)
    return result[0]['summary_text']

def translate_to_urdu(text):
    result = translator_en2ur(text)
    return result[0]['translation_text']

def translate_to_english(text):
    result = translator_ur2en(text)
    return result[0]['translation_text']

# === Gradio Interface ===
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown(
        """
        <h1 style='text-align:center; color:#2E86C1;'>🪄 AI Writing Assistant - Grammar, Paraphrase, Translate & Summarize</h1>
        <p style='text-align:center;'>Your all-in-one free QuillBot-style tool – powered by Hugging Face 🤖</p>
        """
    )

    with gr.Tab("🧹 Grammar Correction"):
        text_in = gr.Textbox(label="Enter Text", placeholder="Write or paste text here...")
        text_out = gr.Textbox(label="Corrected Grammar")
        btn = gr.Button("Correct Grammar ✅")
        btn.click(correct_grammar, text_in, text_out)

    with gr.Tab("🔁 Paraphrasing"):
        text_in2 = gr.Textbox(label="Enter Text", placeholder="Write text to rephrase...")
        text_out2 = gr.Textbox(label="Paraphrased Text")
        btn2 = gr.Button("Paraphrase 🪄")
        btn2.click(paraphrase_text, text_in2, text_out2)

    with gr.Tab("📝 Summarization"):
        text_in3 = gr.Textbox(label="Enter Text", placeholder="Paste long text to summarize...")
        text_out3 = gr.Textbox(label="Summary")
        btn3 = gr.Button("Summarize 📄")
        btn3.click(summarize_text, text_in3, text_out3)

    with gr.Tab("🌐 Translation"):
        with gr.Row():
            text_in4 = gr.Textbox(label="Enter English or Urdu Text")
        with gr.Row():
            text_out4_en2ur = gr.Textbox(label="English → Urdu")
            text_out4_ur2en = gr.Textbox(label="Urdu → English")

        btn4a = gr.Button("Translate English → Urdu 🇺🇸➡🇵🇰")
        btn4b = gr.Button("Translate Urdu → English 🇵🇰➡🇺🇸")

        btn4a.click(translate_to_urdu, text_in4, text_out4_en2ur)
        btn4b.click(translate_to_english, text_in4, text_out4_ur2en)

app.launch()