studyhub/app/service.py

from __future__ import annotations

from pathlib import Path
from typing import Optional

from ..core.answer import AnswerResult, answer, answer_in_thread
from ..core.citations import Citation
from ..core.llm import LLMClient
from ..core.loader import RetrievalIndex
from ..core.study import (flashcards_from_chunks, make_quiz, mindmap_from_chunks,
                          quiz_items_from_chunks, summarize_file)
from ..core.topic import answer_topic, answer_topic_in_thread
from ..core.topic import topic_summary as _topic_summary
from ..graph.models import Graph
from ..ingest.embedder import Embedder


def _source_type(path: str) -> str:
    p = path.lower()
    if p.startswith("_uploads/") or "/_uploads/" in p:
        return "upload"
    if "syllabus" in p:
        return "syllabus"
    if "transcript" in p:
        return "transcript"
    if "note" in p:
        return "notes"
    if "slide" in p:
        return "slides"
    return "material"


class StudyService:
    def __init__(self, index: RetrievalIndex, embedder: Embedder, llm: LLMClient, materials_dir: str,
                 reranker=None, graph: Graph | None = None):
        self.index = index
        self.embedder = embedder
        self.llm = llm
        self.materials_dir = Path(materials_dir)
        self.reranker = reranker
        self.graph = graph if graph is not None else Graph()

    def ask(self, query: str) -> AnswerResult:
        return answer(query, self.index, self.embedder, self.llm, reranker=self.reranker)

    def ask_thread(self, query: str, prior: list[dict]) -> AnswerResult:
        return answer_in_thread(query, prior, self.index, self.embedder, self.llm, reranker=self.reranker)

    def summarize(self, file: str) -> AnswerResult:
        return summarize_file(self.index, file, self.llm)

    def quiz(self, topic: str, n: int = 5) -> AnswerResult:
        return make_quiz(topic, self.index, self.embedder, self.llm, n=n, reranker=self.reranker)

    def list_files(self) -> list[str]:
        seen: list[str] = []
        for c in self.index.chunks:
            if c.file not in seen:
                seen.append(c.file)
        return seen

    def download_path(self, file: str) -> Optional[str]:
        """Resolve an original under materials_dir, guarding against path traversal,
        and only for files that are actually in the index (allowlist)."""
        if file not in set(self.list_files()):
            return None
        base = self.materials_dir.resolve()
        target = (base / file).resolve()
        if (base == target or base in target.parents) and target.is_file():
            return str(target)
        return None

    def cited_files(self, result) -> list[str]:
        """Distinct on-disk paths for the files an answer cites, so a user can open/verify
        the exact source right next to the answer."""
        seen: set[str] = set()
        out: list[str] = []
        for c in result.sources:
            if c.file in seen:
                continue
            seen.add(c.file)
            path = self.download_path(c.file)
            if path:
                out.append(path)
        return out

    # --- Topic Graph (v2) ---

    def modules(self) -> list[dict]:
        return [{"id": m.id, "title": m.title, "lecturer": m.lecturer,
                 "n_topics": len(m.topic_ids), "n_sources": len(m.source_ids)}
                for m in self.graph.modules]

    def topics(self, module_id: str) -> list[dict]:
        return [{"id": t.id, "title": t.title, "n_chunks": len(t.chunk_ids)}
                for t in self.graph.topics_for_module(module_id)]

    def topic_summary(self, topic_id: str) -> AnswerResult:
        t = self.graph.topic_by_id(topic_id)
        if t is None:
            return AnswerResult("Unknown topic.", [], "")
        return _topic_summary(self.index, t.chunk_ids, self.llm)

    def topic_ask(self, topic_id: str, query: str) -> AnswerResult:
        t = self.graph.topic_by_id(topic_id)
        if t is None:
            return AnswerResult("Unknown topic.", [], "")
        return answer_topic(query, self.index, t.chunk_ids, self.embedder, self.llm, reranker=self.reranker)

    def topic_ask_thread(self, topic_id: str, query: str, prior: list[dict]) -> AnswerResult:
        t = self.graph.topic_by_id(topic_id)
        if t is None:
            return AnswerResult("Unknown topic.", [], "")
        return answer_topic_in_thread(query, prior, self.index, t.chunk_ids, self.embedder, self.llm,
                                      reranker=self.reranker)

    def topic_sources(self, topic_id: str) -> list[dict]:
        t = self.graph.topic_by_id(topic_id)
        if t is None:
            return []
        counts: dict[str, int] = {}
        for cid in t.chunk_ids:
            chunk = self.index.by_id.get(cid)
            if chunk is not None:
                counts[chunk.file] = counts.get(chunk.file, 0) + 1
        return [{"file": f, "type": _source_type(f), "n": n} for f, n in counts.items()]

    # --- Structured study features (quiz / flashcards / mind map) ---

    _NO_MATERIAL = "No material for this topic yet."

    def _topic_chunks(self, topic_id: str):
        """Topic + its resolvable chunks (same resolution as topic_summary: skip missing ids)."""
        t = self.graph.topic_by_id(topic_id)
        if t is None:
            return None, []
        return t, [self.index.by_id[c] for c in t.chunk_ids if c in self.index.by_id]

    def _citations_for(self, sids: set[str], id_map) -> list[dict]:
        """Citation dicts (same fields /api answers expose via Citation.__dict__) for the
        sids a generator actually cited, in S-number order."""
        out: list[dict] = []
        for sid in sorted(sids, key=lambda s: int(s[1:])):
            a = id_map[sid].anchor
            out.append(Citation(sid=sid, label=a.label, file=a.file, kind=a.kind,
                                locator=a.locator, is_notes=a.is_notes).__dict__)
        return out

    def topic_quiz(self, topic_id: str, n: int = 5) -> dict:
        _, chunks = self._topic_chunks(topic_id)
        if not chunks:
            return {"questions": [], "sources": [], "error": self._NO_MATERIAL}
        items, id_map = quiz_items_from_chunks(chunks, self.llm, n=n)
        used = {s for it in items for s in it["sids"]}
        return {"questions": items, "sources": self._citations_for(used, id_map)}

    def topic_flashcards(self, topic_id: str, n: int = 10) -> dict:
        _, chunks = self._topic_chunks(topic_id)
        if not chunks:
            return {"cards": [], "sources": [], "error": self._NO_MATERIAL}
        cards, id_map = flashcards_from_chunks(chunks, self.llm, n=n)
        used = {s for c in cards for s in c["sids"]}
        return {"cards": cards, "sources": self._citations_for(used, id_map)}

    def topic_mindmap(self, topic_id: str) -> dict:
        """Mind map rooted at the topic title (the LLM's root is overridden so the map
        always matches the node the user clicked). Empty generation is flagged with
        "error" so callers don't cache a transient LLM failure as the topic's map."""
        t, chunks = self._topic_chunks(topic_id)
        if t is None:
            return {"root": "", "children": [], "error": self._NO_MATERIAL}
        if not chunks:
            return {"root": t.title, "children": [], "error": self._NO_MATERIAL}
        children = mindmap_from_chunks(chunks, self.llm).get("children") or []
        out = {"root": t.title, "children": children}
        if not children:
            out["error"] = "Could not build a mind map for this topic — try again."
        return out