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

# ✅ Page config MUST be first Streamlit command
st.set_page_config(page_title="AI Study Plan Assistant", layout="wide")

# Load models with caching
@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()

# Extract text from supported formats
def extract_text(file):
    ext = os.path.splitext(file.name)[1].lower()
    if ext == ".txt":
        return file.read().decode("utf-8")
    elif ext == ".docx":
        doc = Document(file)
        return "\n".join([para.text for para in doc.paragraphs])
    elif ext == ".pdf":
        pdf_reader = PdfReader(file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text() + "\n"
        return text
    else:
        raise ValueError("Only .txt, .docx, and .pdf files are supported.")

# Generate study plan
def generate_plan(file, hours_per_day, exam_date, language):
    try:
        content = extract_text(file)
        today = datetime.now().date()
        exam = datetime.strptime(exam_date.strip(), "%Y-%m-%d").date()
        days = (exam - today).days

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

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

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

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

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

# Streamlit UI
st.sidebar.title("📚 Study Assistant Options")
uploaded_file = st.sidebar.file_uploader("Upload syllabus (.txt, .docx, .pdf)", type=["txt", "docx", "pdf"])
study_hours = st.sidebar.number_input("Study hours per day", min_value=1, max_value=12, value=3)
exam_date = st.sidebar.text_input("Exam Date (YYYY-MM-DD)", value="2025-06-30")
language = st.sidebar.selectbox("Select Language", ["English", "Urdu"])

st.title("🧠 AI Study Plan & QA Assistant")

tab1, tab2 = st.tabs(["📅 Generate Study Plan", "❓ Ask a Question"])

with tab1:
    st.subheader("Generate a Personalized Study Plan")
    if uploaded_file and st.button("Generate Study Plan"):
        result = generate_plan(uploaded_file, study_hours, exam_date, language)
        st.text_area("Study Plan", result, height=400)

with tab2:
    st.subheader("Ask Questions from Uploaded Material")
    question = st.text_input("Enter your question:")
    if uploaded_file and question and st.button("Get Answer"):
        answer = ask_question(uploaded_file, question)
        st.text_area("Answer", answer, height=200)