app.py

import os
from dataclasses import dataclass
from hashlib import sha1
from pathlib import Path
from typing import Any, Callable

import streamlit as st
from dotenv import load_dotenv
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langchain_google_genai import ChatGoogleGenerativeAI
from pypdf import PdfReader
from pptx import Presentation
from docx import Document as DocxDocument


def load_css(path):
    with open(path) as f:
        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

load_dotenv()
load_css("assets/style.css")
APP_ROOT = Path(__file__).resolve().parent
DEFAULT_DOWNLOADS_DIR = APP_ROOT / "downloads"
DEFAULT_SYSTEM_PROMPT = (
    "You are a helpful teaching assistant who answers questions using only the "
    "class context provided earlier. Cite the filename when possible and say you do "
    "not know if the answer is missing."
)
CHEERFUL_SYSTEM_PROMPT = (
    "You are a cheerful and upbeat teaching assistant who answers questions using only the "
    "class context provided earlier. Keep your tone warm but concise. Cite the filename when possible and say you do "
    "not know if the answer is missing."
)


@dataclass
class FileContext:
    """Holds raw text and metadata for a single file."""

    filename: str
    course: str
    module: str
    content: str

    def serialize(self) -> str:
        body = self.content.strip()
        if not body:
            return ""
        return (
            f"### File: {self.filename}\n"
            f"Course: {self.course}\n"
            f"Module: {self.module}\n"
            "---\n"
            f"{body}"
        )


@dataclass(frozen=True)
class FileSnapshot:
    """Lightweight descriptor for cache-friendly file snapshots."""

    absolute_path: str
    relative_path: str
    suffix: str
    size: int
    modified_ns: int

    def path(self) -> Path:
        return Path(self.absolute_path)


@dataclass(frozen=True)
class CourseIndex:
    """Immutable container with a course's file manifest and fingerprint."""

    snapshots: tuple[FileSnapshot, ...]
    fingerprint: str


@dataclass(frozen=True)
class CourseContextPayload:
    """Holds serialized context text and companion metadata."""

    text: str
    filenames: tuple[str, ...]
    size_bytes: int


SUPPORTED_EXTENSIONS: dict[str, str] = {
    ".pdf": "pdf",
    ".ppt": "ppt",
    ".pptx": "ppt",
    ".txt": "text",
    ".md": "text",
    ".rtf": "text",
    ".docx": "docx",
    ".py": "text",
    ".js": "text",
    ".ts": "text",
    ".tsx": "text",
    ".java": "text",
    ".c": "text",
    ".cpp": "text",
    ".cs": "text",
    ".json": "text",
    ".yml": "text",
    ".yaml": "text",
    ".csv": "text",
    ".html": "text" # NEW
}
MAX_FILE_SIZE_BYTES = 100 * 1024 * 1024  # 100 MB per file


def read_pdf(path: Path) -> str:
    with path.open("rb") as handle:
        reader = PdfReader(handle)
        parts: list[str] = []
        for page in reader.pages:
            text = page.extract_text() or ""
            if text:
                parts.append(text)
    return "\n".join(parts)


def read_ppt(path: Path) -> str:
    with path.open("rb") as handle:
        presentation = Presentation(handle)
    slide_text: list[str] = []
    for slide_index, slide in enumerate(presentation.slides, start=1):
        buffer: list[str] = []
        for shape in slide.shapes:
            if not getattr(shape, "has_text_frame", False):
                continue
            text_frame = getattr(shape, "text_frame", None)
            if text_frame is None:
                continue
            paragraph_text = [para.text.strip() for para in text_frame.paragraphs if para.text]
            snippet = "\n".join(filter(None, paragraph_text))
            if snippet:
                buffer.append(snippet)
        if buffer:
            slide_text.append(f"Slide {slide_index}:\n" + "\n".join(buffer))
    return "\n\n".join(slide_text)


def read_docx(path: Path) -> str:
    with path.open("rb") as handle:
        document = DocxDocument(handle)
        lines = [para.text for para in document.paragraphs if para.text]
    return "\n".join(lines)


def read_text_file(path: Path) -> str:
    with path.open("r", encoding="utf-8", errors="ignore") as handle:
        return handle.read()


LOADER_MAP: dict[str, Callable[[Path], str]] = {
    "pdf": read_pdf,
    "ppt": read_ppt,
    "docx": read_docx,
    "text": read_text_file,
}


def init_session_state() -> None:
    defaults = {
        "messages": [],
        "current_context": None,
        "current_class": None,
        "current_file_list": [],
        "context_fingerprint": None,
        "api_key": None,
        "system_prompt": DEFAULT_SYSTEM_PROMPT,
    }
    for key, value in defaults.items():
        st.session_state.setdefault(key, value)


def list_course_directories(downloads_root: Path) -> list[Path]:
    if not downloads_root.exists():
        return []
    return sorted([path for path in downloads_root.iterdir() if path.is_dir()])


def discover_files(root: Path) -> list[tuple[Path, os.stat_result]]:
    records: list[tuple[Path, os.stat_result]] = []
    for item in sorted(root.rglob("*")):
        if not item.is_file():
            continue
        ext = item.suffix.lower()
        if ext not in SUPPORTED_EXTENSIONS:
            continue
        try:
            stats = item.stat()
        except OSError:
            continue
        if stats.st_size == 0 or stats.st_size > MAX_FILE_SIZE_BYTES:
            continue
        records.append((item, stats))
    return records


@st.fragment
def render_chat_history() -> None:
    for payload in st.session_state.get("messages", []):
        role = payload["role"]
        with st.chat_message("user" if role == "user" else "assistant",
                     avatar="assets/BEAN.png" if role=="assistant" else "assets/USER.png"):
            st.markdown(payload["content"])


def get_file_metadata(file_path: Path, downloads_dir: Path) -> dict[str, str]:
    try:
        relative_path = file_path.relative_to(downloads_dir)
        parts = relative_path.parts
        course = parts[0] if len(parts) >= 1 else "Unknown Course"
        module = parts[2] if len(parts) >= 3 else (parts[1] if len(parts) >= 2 else "Unknown Module")
    except ValueError:
        course = "Unknown Course"
        module = "Unknown Module"
    return {"course": course, "module": module, "filename": file_path.name}


def collect_file_snapshots(course_path: Path, downloads_dir: Path) -> CourseIndex:
    inventories = discover_files(course_path)
    if not inventories:
        raise ValueError("No supported files found in this class.")

    snapshots: list[FileSnapshot] = []
    hasher = sha1()
    for file_path, stats in inventories:
        try:
            relative_path = file_path.relative_to(downloads_dir)
        except ValueError:
            relative_path = Path(file_path.name)
        relative_str = relative_path.as_posix()
        hasher.update(relative_str.encode("utf-8"))
        hasher.update(stats.st_size.to_bytes(8, "big", signed=False))
        hasher.update(stats.st_mtime_ns.to_bytes(8, "big", signed=False))
        snapshots.append(
            FileSnapshot(
                absolute_path=str(file_path),
                relative_path=relative_str,
                suffix=file_path.suffix.lower(),
                size=stats.st_size,
                modified_ns=stats.st_mtime_ns,
            )
        )

    return CourseIndex(snapshots=tuple(snapshots), fingerprint=hasher.hexdigest())


@st.cache_data(show_spinner=False, max_entries=1024)
def load_snapshot_text(snapshot: FileSnapshot) -> str:
    """Read and normalize a single file snapshot with memoization via st.cache_data."""

    loader_key = SUPPORTED_EXTENSIONS.get(snapshot.suffix)
    if not loader_key:
        return ""
    loader = LOADER_MAP.get(loader_key)
    if not loader:
        return ""
    return loader(snapshot.path()).strip()


@st.cache_resource(show_spinner="Building class context...", max_entries=16)
def assemble_course_context(
    *,
    fingerprint: str,
    snapshots: tuple[FileSnapshot, ...],
    downloads_dir_str: str,
) -> CourseContextPayload:
    """Aggregate cached file reads into a single shared context blob using st.cache_resource."""

    _ = fingerprint  # Included so the cache key tracks the manifest fingerprint.
    downloads_dir = Path(downloads_dir_str)
    context_chunks: list[str] = []
    filenames: list[str] = []

    for snapshot in snapshots:
        try:
            raw_text = load_snapshot_text(snapshot)
        except Exception as exc:  # pragma: no cover - surface to sidebar for visibility
            st.sidebar.warning(f"Failed to read {Path(snapshot.absolute_path).name}: {exc}")
            continue
        if not raw_text:
            continue
        metadata = get_file_metadata(snapshot.path(), downloads_dir)
        chunk = FileContext(
            filename=metadata["filename"],
            course=metadata["course"],
            module=metadata["module"],
            content=raw_text,
        ).serialize()
        if chunk:
            context_chunks.append(chunk)
            filenames.append(metadata["filename"])

    if not context_chunks:
        raise ValueError("Unable to read any supported files for this class.")

    text = "\n\n".join(context_chunks)
    return CourseContextPayload(
        text=text,
        filenames=tuple(filenames),
        size_bytes=len(text.encode("utf-8")),
    )


def answer_question(
    *,
    question: str,
    api_key: str,
    chat_history: list[dict[str, Any]],
    context_text: str,
    system_prompt: str,
) -> str:
    llm = ChatGoogleGenerativeAI(
        model="gemini-2.5-flash-lite-preview-09-2025",
        temperature=0.2,
        api_key=api_key,
    )

    context_message = SystemMessage(
        content=(
            f"### Class Context\n"
            f"The following text contains *all* material for the current class.\n"
            f"---\n{context_text}\n---"
        )
    )
    system_message = SystemMessage(content=system_prompt)

    messages: list[BaseMessage] = [context_message, system_message]

    for msg in chat_history:
        if msg["role"] == "user":
            messages.append(HumanMessage(content=msg["content"]))
        elif msg["role"] == "assistant":
            messages.append(AIMessage(content=msg["content"]))

    messages.append(HumanMessage(content=question))

    try:
        response = llm.invoke(messages)
        if isinstance(response.content, str):
            return response.content
        if isinstance(response.content, list):
            fragments: list[str] = []
            for fragment in response.content:
                if isinstance(fragment, str):
                    fragments.append(fragment)
                else:
                    fragments.append(str(fragment))
            return "\n".join(fragments)
        return str(response.content)
    except Exception as exc:
        if "context_length_exceeded" in str(exc) or "input_size_exceeded" in str(exc):
            return (
                "Error: The combined text for this class is too large for the model's context window. "
                "Remove some files or split the class into smaller subsets."
            )
        return f"An error occurred while contacting Google Gemini: {exc}"


def resolve_downloads_dir() -> Path:
    override = os.getenv("DOWNLOADS_ROOT")
    if override:
        return Path(override).expanduser().resolve()
    return DEFAULT_DOWNLOADS_DIR


def start_app() -> None:
    st.set_page_config(page_title="Chat With Your Class", page_icon="assets/BEAN_TRANS.png", layout="wide")
    with st.sidebar:
        st.image("assets/SB_LOGO_TRANS.png", width=375)
        st.title("⭐ Welcome to StudyBean! ⭐")

    col1, col2 = st.columns([1, 15])
    with col1:
        st.image("assets/BEAN_TRANS.png", width=100)
    with col2:
        st.title("StudyBean: Chat With Your Class")
    downloads_dir = resolve_downloads_dir()
    if not downloads_dir.exists():
        st.error(
            "Downloads folder not found. Set the DOWNLOADS_ROOT env var or place your Canvas exports under "
            f"{downloads_dir}."
        )
        st.stop()

    init_session_state()

    api_key = os.getenv("GOOGLE_API_KEY", "")

    if not api_key:
        st.warning("Enter your Google API key in the sidebar to enable chatting.")

    # system_prompt_input = st.sidebar.text_area(
    #     "System prompt",
    #     value=st.session_state.get("system_prompt", DEFAULT_SYSTEM_PROMPT),
    #     height=140,
    # )
    # sanitized_prompt = (system_prompt_input or "").strip()
    # st.session_state["system_prompt"] = sanitized_prompt or DEFAULT_SYSTEM_PROMPT


    # NEW : Tone Selector #
    tone = st.sidebar.radio(
        "Study Vibe",
        ["📘 Neutral", "✨ Cheerful"],
        help="Pick the tone of the bot's responses",
        index=0,
        horizontal=True,
    )
    if tone.startswith("📘"):
        st.session_state["system_prompt"] = DEFAULT_SYSTEM_PROMPT
    else:
        st.session_state["system_prompt"] = CHEERFUL_SYSTEM_PROMPT

    course_dirs = list_course_directories(downloads_dir)
    if not course_dirs:
        st.error("No class folders found inside downloads/.")
        st.stop()

    course_names = [course.name for course in course_dirs]
    default_index = 0
    if st.session_state.get("current_class") in course_names:
        default_index = course_names.index(st.session_state["current_class"])

    selected_course_name = st.sidebar.selectbox(
        "🚀 Class",
        options=course_names,
        index=default_index,
        help="Select which class context to load into the LLM prompt",
    )

    force_reload = st.sidebar.button("Force rebuild context", use_container_width=True)
    if st.sidebar.button("Clear chat history", use_container_width=True):
        st.session_state["messages"] = []

    course_path = next(course for course in course_dirs if course.name == selected_course_name)
    try:
        course_index = collect_file_snapshots(course_path, downloads_dir)
    except ValueError as exc:
        st.session_state["current_context"] = None
        st.session_state["current_file_list"] = []
        st.session_state["current_class"] = None
        st.session_state["context_fingerprint"] = None
        st.sidebar.error(str(exc))
        st.info("Add supported files to this class to start chatting.")
        return

    snapshots = course_index.snapshots
    snapshot_hash = course_index.fingerprint

    needs_load = (
        force_reload
        or st.session_state.get("current_context") is None
        or st.session_state.get("current_class") != selected_course_name
        or st.session_state.get("context_fingerprint") != snapshot_hash
    )

    if needs_load:
        if force_reload:
            load_snapshot_text.clear()
            assemble_course_context.clear()
        try:
            payload = assemble_course_context(
                fingerprint=snapshot_hash,
                snapshots=snapshots,
                downloads_dir_str=str(downloads_dir),
            )
        except ValueError as exc:
            st.session_state["current_context"] = None
            st.session_state["current_file_list"] = []
            st.session_state["current_class"] = None
            st.session_state["context_fingerprint"] = None
            st.sidebar.error(str(exc))
        else:
            st.session_state["current_context"] = payload.text
            st.session_state["current_file_list"] = list(payload.filenames)
            st.session_state["current_class"] = selected_course_name
            st.session_state["context_fingerprint"] = snapshot_hash
            st.sidebar.success(f"Loaded {len(payload.filenames)} files for {selected_course_name}.")
            sidebar_size_mb = payload.size_bytes / (1024 * 1024)
            st.sidebar.caption(f"Context size: {sidebar_size_mb:.2f} MB")

    context_text = st.session_state.get("current_context")

    if not context_text:
        st.info("Select a class to automatically load the full context before chatting.")
        return

    st.subheader(f"Selected Class: {st.session_state.get('current_class', 'Unknown class')}")
    with st.expander(f"Files in context ({len(st.session_state.get('current_file_list', []))})"):
        st.write(st.session_state.get("current_file_list", []))

    render_chat_history()
    user_prompt = st.chat_input("Ask something about this class...")

    if user_prompt:
        if not api_key:
            st.warning("Add your Google API key in the sidebar before chatting.")
            return
        st.session_state["messages"].append({"role": "user", "content": user_prompt})
        with st.spinner("Thinking..."):
            response_text = answer_question(
                question=user_prompt,
                api_key=api_key,
                chat_history=st.session_state["messages"][:-1],
                context_text=context_text,
                system_prompt=st.session_state["system_prompt"],
            )
        st.session_state["messages"].append({"role": "assistant", "content": response_text})
        render_chat_history()


if __name__ == "__main__":
    start_app()