import streamlit as st
from langchain_groq import ChatGroq
import os
from reportlab.lib.pagesizes import A4
from reportlab.pdfgen import canvas
import tempfile

# Load API key
def getting_api_key():
    return os.getenv("GROQ_KEY")

api_key = getting_api_key()

# Initialize LLM
llm = ChatGroq(
    api_key=api_key,
    model="llama-3.1-8b-instant"
)

st.title("📘 Study Question Generator")

# Input notes
context = st.text_area("Enter your notes:")

# Options
question_types = st.multiselect(
    "Select question types:",
    ["Multiple Choice", "Short Answer", "True/False", "Essay / Analytical"],
    default=["Short Answer"]
)

difficulty_level = st.selectbox(
    "Select difficulty level:",
    [
        "Beginner (basic understanding)",
        "Intermediate (application & explanation)",
        "Advanced (critical thinking & between-the-lines)"
    ]
)

num_questions = st.slider("Number of questions:", 3, 20, 5)

# Session state
if "questions" not in st.session_state:
    st.session_state.questions = ""

if "answers" not in st.session_state:
    st.session_state.answers = ""

# Generate Questions
if st.button("Generate Questions"):
    if context.strip() and question_types:
        prompt = f"""
Generate {num_questions} study questions from the material below.

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

Question types: {", ".join(question_types)}
Difficulty: {difficulty_level}

Do NOT include answers.
"""
        response = llm.invoke(prompt)
        st.session_state.questions = response.content
        st.session_state.answers = ""

# Show Questions
if st.session_state.questions:
    st.subheader("📝 Questions")
    st.write(st.session_state.questions)

    # Generate Answers
    if st.button("✅ Generate Answers & Explanations"):
        answer_prompt = f"""
You are an expert teacher.

Below are study questions. Provide:
- The correct answer
- A brief explanation for each question

Questions:
\"\"\"{st.session_state.questions}\"\"\"

Format clearly as:
Question:
Answer:
Explanation:
"""
        answer_response = llm.invoke(answer_prompt)
        st.session_state.answers = answer_response.content

# Show Answers
if st.session_state.answers:
    st.subheader("📖 Answers & Explanations")
    st.write(st.session_state.answers)

    # PDF Export
    if st.button("📄 Download Questions + Answers as PDF"):
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            c = canvas.Canvas(tmp.name, pagesize=A4)
            width, height = A4
            text = c.beginText(40, height - 50)
            text.setFont("Helvetica", 11)

            content = (
                "STUDY QUESTIONS\n\n"
                + st.session_state.questions
                + "\n\nANSWERS & EXPLANATIONS\n\n"
                + st.session_state.answers
            )

            for line in content.split("\n"):
                text.textLine(line)

            c.drawText(text)
            c.save()

            with open(tmp.name, "rb") as f:
                st.download_button(
                    "⬇️ Download PDF",
                    f,
                    file_name="study_questions_with_answers.pdf",
                    mime="application/pdf"
                )