import os
from datetime import datetime
from PyPDF2 import PdfReader
from docx import Document
import streamlit as st
from groq import Groq
import textwrap

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

# ✅ Load Groq Client
@st.cache_resource
def load_groq_client():
    return Groq(api_key=os.getenv("GROQ_API_KEY"))

groq_client = load_groq_client()

# ✅ File Text Extraction
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:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
        return text
    else:
        raise ValueError("Only .txt, .docx, and .pdf files are supported.")

# ✅ Chunking helper
def chunk_text(text, chunk_size=1500):
    return textwrap.wrap(text, width=chunk_size, break_long_words=False, replace_whitespace=False)

# ✅ Query Groq LLM
def query_groq(prompt, temperature=0.4):
    try:
        chat = groq_client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama3-8b-8192",
            temperature=temperature
        )
        return chat.choices[0].message.content.strip()
    except Exception as e:
        return f"⚠️ Groq API Error: {str(e)}"

# ✅ Generate Study Plan with chunking
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."

        chunks = chunk_text(content, chunk_size=1500)
        plan_parts = []

        for idx, chunk in enumerate(chunks):
            prompt = f"""This is part {idx + 1} of {len(chunks)} of a syllabus.
Create a study plan segment for this syllabus. Total study duration is {days} days with {hours_per_day} hours/day.
Write the study plan in {language}.

Syllabus:
\"\"\"
{chunk}
\"\"\"
"""
            response = query_groq(prompt)
            plan_parts.append(response)

        return "\n\n".join(plan_parts)

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

# ✅ Ask Question with context chunking
def ask_question(file, question):
    try:
        context = extract_text(file)
        chunks = chunk_text(context, chunk_size=1500)
        answers = []

        for idx, chunk in enumerate(chunks):
            prompt = f"""Use the following part of study material to answer the question:

Material:
\"\"\"
{chunk}
\"\"\"

Question: {question}
Answer:"""
            answer = query_groq(prompt)
            answers.append(f"Part {idx + 1}:\n{answer}")

        return "\n\n".join(answers)

    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=300)