import streamlit as st
from transformers import pipeline
from langdetect import detect
from datetime import datetime
import base64
import os
import PyPDF2
from docx import Document

# Initialize models
summarizer = pipeline("summarization")
translator_en_ur = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ur")
translator_ur_en = pipeline("translation", model="Helsinki-NLP/opus-mt-ur-en")
grammar_corrector = pipeline("text2text-generation", model="vennify/t5-base-grammar-correction")

st.set_page_config(page_title="CSS AI Assistant by Komal", layout="wide")

st.title("📚 CSS AI Assistant 🇵🇰")
st.write("Helping students with news summaries, precis, essays & more — for free!")

menu = st.sidebar.radio("Select Feature", [
    "Daily News Summary",
    "Precis Evaluation",
    "Essay Feedback",
    "Saved Notes",
    "File Upload"
])

# Local file to save notes
NOTES_FILE = "saved_notes.txt"

def save_note(title, content):
    with open(NOTES_FILE, "a", encoding='utf-8') as f:
        f.write(f"\n## {title}\n{content}\n")

def load_notes():
    if os.path.exists(NOTES_FILE):
        with open(NOTES_FILE, "r", encoding='utf-8') as f:
            return f.read()
    return "No saved notes yet."

def generate_download_link(text, filename):
    b64 = base64.b64encode(text.encode()).decode()
    return f'<a href="data:file/txt;base64,{b64}" download="{filename}">📥 Download Notes</a>'

def extract_text_from_pdf(uploaded_file):
    pdf_reader = PyPDF2.PdfReader(uploaded_file)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

def extract_text_from_word(uploaded_file):
    doc = Document(uploaded_file)
    text = ""
    for para in doc.paragraphs:
        text += para.text
    return text

# --- Daily News Summary ---
if menu == "Daily News Summary":
    st.header("📰 AI-Powered News Summarizer")
    user_news = st.text_area("Paste or write today's news:", height=300)
    lang = st.radio("Language", ["English", "Urdu"])

    if st.button("Summarize") and user_news:
        summary = summarizer(user_news, max_length=500, min_length=100, do_sample=False)[0]['summary_text']
        if lang == "Urdu":
            translated = translator_en_ur(summary)[0]['translation_text']
            st.success("Summary in Urdu:")
            st.write(translated)
            if st.button("Save Urdu Summary"):
                save_note(f"Urdu Summary - {datetime.now().date()}", translated)
        else:
            st.success("Summary in English:")
            st.write(summary)
            if st.button("Save English Summary"):
                save_note(f"English Summary - {datetime.now().date()}", summary)

# --- Precis Evaluation ---
elif menu == "Precis Evaluation":
    st.header("✍️ Precis Evaluation Tool")
    precis_input = st.text_area("Enter your precis for evaluation:", height=300)
    if st.button("Evaluate Precis") and precis_input:
        corrected = grammar_corrector("grammar: " + precis_input)[0]['generated_text']
        st.success("✅ Corrected Precis:")
        st.write(corrected)
        if st.button("Save Precis"):
            save_note(f"Precis - {datetime.now().date()}", corrected)

# --- Essay Feedback ---
elif menu == "Essay Feedback":
    st.header("🖋️ Essay Evaluation Tool")
    essay_input = st.text_area("Paste your essay here:", height=400)
    if st.button("Evaluate Essay") and essay_input:
        corrected = grammar_corrector("grammar: " + essay_input)[0]['generated_text']
        st.success("✅ Suggestions:")
        st.write(corrected)
        if st.button("Save Essay"):
            save_note(f"Essay - {datetime.now().date()}", corrected)

# --- Saved Notes ---
elif menu == "Saved Notes":
    st.header("💾 Saved Notes")
    notes = load_notes()
    st.markdown(notes)
    st.markdown(generate_download_link(notes, "css_notes.txt"), unsafe_allow_html=True)

# --- File Upload ---
elif menu == "File Upload":
    st.header("📄 File Upload - PDF or Word")
    uploaded_file = st.file_uploader("Upload a PDF or Word file", type=["pdf", "docx"])
    
    if uploaded_file is not None:
        file_type = uploaded_file.type
        if file_type == "application/pdf":
            file_text = extract_text_from_pdf(uploaded_file)
        elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
            file_text = extract_text_from_word(uploaded_file)

        st.text_area("Extracted Text", file_text, height=300)

        # Now you can apply the summarization, translation, or grammar correction on the extracted text
        lang = st.radio("Language for Summary/Translation", ["English", "Urdu"])
        if st.button("Summarize and Translate") and file_text:
            summary = summarizer(file_text, max_length=500, min_length=100, do_sample=False)[0]['summary_text']
            if lang == "Urdu":
                translated = translator_en_ur(summary)[0]['translation_text']
                st.success("Summary in Urdu:")
                st.write(translated)
            else:
                st.success("Summary in English:")
                st.write(summary)