import streamlit as st
from transformers import pipeline
from docx import Document
from PyPDF2 import PdfReader
from datetime import datetime
from fpdf import FPDF
import tempfile
import os

# Load models (you can replace with smaller or local models if needed)
@st.cache_resource
def load_models():
    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
    qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
    return summarizer, qa_pipeline

summarizer, qa_pipeline = load_models()

# Text extraction
def extract_text(uploaded_file):
    ext = os.path.splitext(uploaded_file.name)[-1].lower()
    if ext == ".txt":
        return uploaded_file.read().decode("utf-8")
    elif ext == ".docx":
        doc = Document(uploaded_file)
        return "\n".join([para.text for para in doc.paragraphs])
    elif ext == ".pdf":
        pdf = PdfReader(uploaded_file)
        return "\n".join([page.extract_text() for page in pdf.pages if page.extract_text()])
    else:
        return None

# Generate Study Plan
def generate_study_plan(text, hours, exam_date, language):
    today = datetime.now().date()
    try:
        exam = datetime.strptime(exam_date.strip(), "%Y-%m-%d").date()
        days = (exam - today).days
    except:
        return "❌ Invalid date format. Use YYYY-MM-DD."

    if days <= 0:
        return "❌ Exam date must be in the future."

    prompt = f"Create a {days}-day study plan for this syllabus:\n{text}\nDaily study time: {hours} hours."
    if language == "Urdu":
        prompt += "\nTranslate the plan into Urdu."

    summary = summarizer(prompt, max_length=500, min_length=100, do_sample=False)
    return summary[0]['summary_text']

# Generate PDF
def create_pdf(content):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=12)
    for line in content.split("\n"):
        pdf.multi_cell(0, 10, line)
    temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
    pdf.output(temp_path.name)
    return temp_path.name

# Ask Question
def answer_question(context, question):
    try:
        result = qa_pipeline({"context": context, "question": question})
        return result["answer"]
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# UI
st.set_page_config(page_title="AI Study Plan Assistant", layout="wide")

st.title("📚 AI Study Plan Assistant")

with st.sidebar:
    st.header("Upload Syllabus & Set Options")
    uploaded_file = st.file_uploader("Upload a .pdf, .docx or .txt file", type=["pdf", "docx", "txt"])
    study_hours = st.number_input("Study Hours per Day", min_value=1, max_value=24, value=2)
    exam_date = st.text_input("Exam Date (YYYY-MM-DD)")
    language = st.selectbox("Language of Plan", ["English", "Urdu"])

    generate_btn = st.button("Generate Study Plan")

# Display Plan
if generate_btn and uploaded_file:
    with st.spinner("🔍 Reading content..."):
        raw_text = extract_text(uploaded_file)
    if not raw_text:
        st.error("⚠️ Could not extract text from the file.")
    else:
        with st.spinner("✍️ Generating study plan..."):
            plan = generate_study_plan(raw_text, study_hours, exam_date, language)
        st.subheader("📄 Generated Study Plan")
        st.text_area("Study Plan", value=plan, height=400)

        pdf_path = create_pdf(plan)
        with open(pdf_path, "rb") as f:
            st.download_button(label="📥 Download Study Plan as PDF", data=f, file_name="study_plan.pdf")

# Ask a question
st.markdown("---")
st.subheader("❓ Ask a Question From Your Syllabus")
if uploaded_file:
    question = st.text_input("Type your question:")
    if st.button("Get Answer"):
        context = extract_text(uploaded_file)
        with st.spinner("🤖 Searching for the answer..."):
            answer = answer_question(context, question)
        st.success(f"Answer: {answer}")