import gradio as gr
from transformers import pipeline

# Load models
grammar_model = pipeline("text2text-generation", model="vennify/t5-base-grammar-correction")
meaning_model = pipeline("text2text-generation", model="facebook/bart-large-cnn")

def dictionary_grammar_bot(user_input):
    try:
        # --- Grammar Correction ---
        grammar_output = grammar_model(
            user_input,
            max_new_tokens=128,
            clean_up_tokenization_spaces=True
        )[0]['generated_text']

        # --- Meaning / Definition ---
        meaning_prompt = f"Explain the meaning and definition of the word or phrase: {user_input}"
        meaning_output = meaning_model(
            meaning_prompt,
            max_new_tokens=200,
            clean_up_tokenization_spaces=True
        )[0]['generated_text']

        # --- Phrase Suggestions ---
        phrase_prompt = f"Suggest grammatically correct alternative phrases or rephrasings for: {user_input}"
        phrase_output = meaning_model(
            phrase_prompt,
            max_new_tokens=120,
            clean_up_tokenization_spaces=True
        )[0]['generated_text']

        # --- Example Sentence ---
        sentence_prompt = f"Write an example sentence using the phrase: {user_input}"
        sentence_output = meaning_model(
            sentence_prompt,
            max_new_tokens=120,
            clean_up_tokenization_spaces=True
        )[0]['generated_text']

        # Format structured output
        response = f"""
### 🧩 Grammar Correction
{grammar_output}

---

### 📘 Meaning / Definition
{meaning_output}

---

### 💬 Phrase Suggestions
{phrase_output}

---

### ✏️ Example Sentence
{sentence_output}
"""

        return response

    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# --- Gradio UI ---
iface = gr.Interface(
    fn=dictionary_grammar_bot,
    inputs=gr.Textbox(
        label="Enter a word, phrase, or sentence",
        placeholder="Example: resistance movement failed"
    ),
    outputs=gr.Markdown(label="AI Response"),
    title="Dictionary & Grammar Chatbot",
    description="AI-powered chatbot that provides meanings, grammar corrections, phrase suggestions, and example sentences.",
    theme="gradio/soft"
)

iface.launch()