import os
import streamlit as st
from groq import Groq

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

import pandas as pd
import docx
from pypdf import PdfReader

from reportlab.platypus import SimpleDocTemplate, Paragraph
from reportlab.lib.styles import getSampleStyleSheet
import tempfile

# -----------------------------
# GROQ CLIENT
# -----------------------------
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# -----------------------------
# PAGE CONFIG
# -----------------------------
st.set_page_config(
    page_title="AI Study Assistant 📚",
    page_icon="🎓",
    layout="wide"
)

# -----------------------------
# SIDEBAR
# -----------------------------
st.sidebar.title("📌 Settings")

education_level = st.sidebar.selectbox(
    "Select Education Level",
    [
        "Primary School",
        "Middle School",
        "Secondary School",
        "High School",
        "Undergraduate",
        "Graduate"
    ]
)

st.sidebar.markdown("---")
st.sidebar.write("Developed by **Ahmad Bilal** | Fiverr Portfolio Demo")

# -----------------------------
# HEADER
# -----------------------------
st.markdown(
    """
    <div style='text-align:center; padding:10px; background-color:#f0f2f6; border-radius:10px'>
        <h1 style='color:#0f4c81'>🎓 AI Study Assistant</h1>
        <p style='font-size:18px'>Upload study materials and ask questions instantly!</p>
    </div>
    """,
    unsafe_allow_html=True
)

# -----------------------------
# FILE UPLOADER
# -----------------------------
uploaded_files = st.file_uploader(
    "Upload Study Documents",
    type=["pdf","docx","txt","csv","xlsx"],
    accept_multiple_files=True
)

valid_files = []

if uploaded_files:

    MAX_FILE_SIZE = 20 * 1024 * 1024

    for file in uploaded_files:
        if file.size > MAX_FILE_SIZE:
            st.error(f"{file.name} is too large. Upload files under 20MB.")
        else:
            valid_files.append(file)

    st.success(f"{len(valid_files)} file(s) ready for processing")

# -----------------------------
# FILE LOADERS
# -----------------------------
def load_pdf(file):
    reader = PdfReader(file)
    text = ""
    for page in reader.pages:
        if page.extract_text():
            text += page.extract_text()
    return text


def load_docx(file):
    doc = docx.Document(file)
    return "\n".join([p.text for p in doc.paragraphs])


def load_csv(file):
    df = pd.read_csv(file)
    return df.to_string()


def load_xlsx(file):
    df = pd.read_excel(file)
    return df.to_string()


def load_txt(file):
    return file.read().decode("utf-8")


# -----------------------------
# DOCUMENT PROCESSING
# -----------------------------
def process_docs(files):

    text = ""

    for file in files:

        if file.type == "application/pdf":
            text += load_pdf(file)

        elif file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
            text += load_docx(file)

        elif file.type == "text/csv":
            text += load_csv(file)

        elif file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
            text += load_xlsx(file)

        else:
            text += load_txt(file)

    return text


# -----------------------------
# VECTOR STORE
# -----------------------------
@st.cache_resource
def create_vectorstore(text):

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=100
    )

    chunks = splitter.split_text(text)

    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2"
    )

    vectorstore = FAISS.from_texts(chunks, embeddings)

    return vectorstore


# -----------------------------
# PROMPT BUILDER
# -----------------------------
def build_prompt(context, question, level):

    style = {
        "Primary School": "Explain like teaching a 5 year old using fun examples.",
        "Middle School": "Explain with easy examples.",
        "Secondary School": "Explain clearly using simple ideas.",
        "High School": "Explain with reasoning and examples.",
        "Undergraduate": "Explain in academic but clear language.",
        "Graduate": "Provide detailed academic explanation."
    }

    prompt = f"""
Use the study material below to answer the question.

Study Material:
{context}

Question:
{question}

Explanation Style:
{style[level]}
"""

    return prompt


# -----------------------------
# GROQ LLM
# -----------------------------
def ask_llm(prompt):

    chat_completion = client.chat.completions.create(
        messages=[{"role":"user","content":prompt}],
        model="llama-3.3-70b-versatile"
    )

    return chat_completion.choices[0].message.content


# -----------------------------
# SUMMARY
# -----------------------------
def generate_summary(text):

    prompt = f"""
Create a short and simple summary of this study material.

{text}
"""

    return ask_llm(prompt)


# -----------------------------
# PDF GENERATOR
# -----------------------------
def create_pdf(text):

    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")

    styles = getSampleStyleSheet()

    story = [Paragraph(text, styles["Normal"])]

    doc = SimpleDocTemplate(temp_file.name)

    doc.build(story)

    return temp_file.name


# -----------------------------
# MAIN LOGIC
# -----------------------------
if valid_files:

    raw_text = process_docs(valid_files)

    vectorstore = create_vectorstore(raw_text)

    st.markdown("---")

    st.subheader("❓ Ask a Question")

    question = st.text_input("Type your question")

    if question:

        col1, col2 = st.columns([2,1])

        docs = vectorstore.similarity_search(question, k=3)

        context = "\n".join([doc.page_content for doc in docs])

        prompt = build_prompt(context, question, education_level)

        answer = ask_llm(prompt)

        with col1:

            st.markdown("### 📖 Answer")

            st.success(answer)

        with col2:

            st.markdown("### 📝 Summary")

            if st.button("Generate Summary"):

                summary = generate_summary(context)

                st.info(summary)

                st.download_button(
                    "Download Markdown",
                    summary,
                    file_name="summary.md"
                )

                pdf_file = create_pdf(summary)

                with open(pdf_file, "rb") as f:

                    st.download_button(
                        "Download PDF",
                        f,
                        file_name="summary.pdf"
                    )

else:

    st.info("📂 Upload at least one study document to start.")