import streamlit as st
from transformers import pipeline
import re

st.set_page_config(page_title="Grammar Agent - Grammarly Alternative", layout="wide")
st.title("📝 Grammar Agent - English Text Corrector")

# Load grammar correction model
@st.cache_resource
def load_model():
    return pipeline("text2text-generation", model="vennify/t5-base-grammar-correction")

grammar_corrector = load_model()

# Utility functions
def capitalize_sentences(text):
    sentences = re.split(r'(?<=[.!?]) +', text)
    return ' '.join(s[0].upper() + s[1:] if s else '' for s in sentences)

def correct_text(text):
    corrected = grammar_corrector(text, max_length=512)[0]['generated_text']
    corrected = capitalize_sentences(corrected)
    return corrected

def word_by_word_diff(original, corrected):
    original_words = original.split()
    corrected_words = corrected.split()
    diffs = []
    for i in range(min(len(original_words), len(corrected_words))):
        if original_words[i] != corrected_words[i]:
            diffs.append((original_words[i], corrected_words[i]))
    if len(original_words) != len(corrected_words):
        extra = corrected_words[len(original_words):]
        for word in extra:
            diffs.append(("", word))
    return diffs

def analyze_differences(original, corrected):
    diffs = word_by_word_diff(original, corrected)
    analysis = []
    for orig, corr in diffs:
        reason = "Grammar, Spelling or Punctuation correction"
        if orig == "":
            reason = "Missing word"
        elif orig.lower() != corr.lower():
            if orig.lower() == corr:
                reason = "Capitalization"
            elif re.sub(r'[.,!?;]', '', orig.lower()) == re.sub(r'[.,!?;]', '', corr.lower()):
                reason = "Punctuation"
            elif len(orig) <= 3 or len(corr) <= 3:
                reason = "Spelling"
        analysis.append({
            "Original": orig,
            "Correction": corr,
            "Reason": reason
        })
    return analysis

# Input Area
input_text = st.text_area("Enter your sentence, paragraph, or essay:", height=250)

if st.button("Correct and Analyze") and input_text:
    with st.spinner("Analyzing and correcting text..."):
        corrected = correct_text(input_text)
        analysis = analyze_differences(input_text, corrected)

    st.subheader("✅ Corrected Text:")
    st.success(corrected)

    st.subheader("📊 Error Analysis:")
    if analysis:
        for item in analysis:
            st.markdown(f"**Original:** `{item['Original']}` → **Correction:** `{item['Correction']}`")
            st.markdown(f"_Reason_: {item['Reason']}")
            st.markdown("---")
    else:
        st.write("No major corrections found. Great job!")

    st.subheader("📃 Side-by-Side Comparison")
    st.columns(2)[0].markdown(f"**Original:**\n\n{input_text}")
    st.columns(2)[1].markdown(f"**Corrected:**\n\n{corrected}")
else:
    st.info("Enter text above and click 'Correct and Analyze' to begin.")