import os
os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"

import streamlit as st
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import pytesseract
from PIL import Image
import difflib
from newspaper import Article

model_name = "google/gemma-2b-it"

@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
    return pipe

pipe = load_model()

st.set_page_config(page_title="Blog/Image Reviewer", layout="centered")
st.title("🧠 BlogChecker AI")
st.markdown("Paste blog text, upload an image, or give a URL — the AI will review and improve it.")

mode = st.radio("Choose Input Mode:", ["Text", "Image", "URL"])
input_text = ""

if mode == "Text":
    input_text = st.text_area("Paste your blog here:", height=300)

elif mode == "Image":
    uploaded_image = st.file_uploader("Upload an image with text", type=["png", "jpg", "jpeg"])
    if uploaded_image is not None:
        image = Image.open(uploaded_image)
        input_text = pytesseract.image_to_string(image)
        st.image(image, caption="Uploaded Image", width=400)
        st.markdown("**Extracted Text:**")
        st.text_area("OCR Text:", value=input_text, height=200)

elif mode == "URL":
    blog_url = st.text_input("Paste blog/article URL:")
    if blog_url:
        try:
            article = Article(blog_url)
            article.download()
            article.parse()
            input_text = article.text
            st.success("✅ Blog content extracted successfully.")
            st.text_area("Extracted Blog Content:", value=input_text, height=300)
        except:
            st.error("⚠️ Failed to extract content from URL.")

def generate_review(original):
    prompt = f"""
You are an AI blog reviewer. Improve the content below:
- Fix unclear, biased, or emotional language
- Correct grammar and spelling
- Identify and address sensitive or policy-violating content
- Suggest better alternatives
- Keep useful content unchanged

Blog:
{original}
"""
    response = pipe(prompt, max_new_tokens=512, do_sample=True, temperature=0.7)
    improved = response[0]['generated_text'].split("Blog:")[-1].strip()
    return improved

def highlight_changes(original, improved):
    d = difflib.Differ()
    diff = list(d.compare(original.split(), improved.split()))
    highlighted = ""
    for word in diff:
        if word.startswith("- "):
            highlighted += f"~~{word[2:]}~~ "
        elif word.startswith("+ "):
            highlighted += f"**{word[2:]}** "
        elif word.startswith("  "):
            highlighted += word[2:] + " "
    return highlighted

if st.button("🔍 Review Content"):
    if input_text.strip():
        with st.spinner("Reviewing..."):
            improved_text = generate_review(input_text)
            highlighted = highlight_changes(input_text, improved_text)

        st.subheader("🧠 AI Suggestions:")
        st.markdown(highlighted)

        if st.button("✅ Accept Changes"):
            st.subheader("✍️ Final Content:")
            st.text_area("Updated Blog:", improved_text, height=300)

        if st.button("🪄 Auto-Correct"):
            st.success("✅ AI-corrected content:")
            st.text_area("Auto-corrected Text:", improved_text, height=300)
    else:
        st.warning("⚠️ Please provide input text, image, or URL.")