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

# ✅ Streamlit Page Config
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()

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

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

# ✅ Study Plan Generation
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 detailed {days}-day study plan from the syllabus below.
You should assume the user can study {hours_per_day} hours per day.
Output should be in {language}.

Syllabus:
\"\"\"
{content}
\"\"\"
"""
        return query_groq(prompt)
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

# ✅ Question Answering
def ask_question(file, question):
    try:
        context = extract_text(file)
        prompt = f"""Based on the following study material, answer the question below:

Study Material:
\"\"\"
{context}
\"\"\"

Question: {question}
Answer:"""
        return query_groq(prompt)
    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)