studyhub/web/api.py

"""FastAPI backend for the custom (Docker-Space) frontend.

Reuses the whole engine via dependency injection: a StudyService (topics, ask,
summary, download), a UserStore (invite-gated registration + login) and a
HistoryStore (per-user activity). Auth is a signed session token, accepted from
an httponly cookie or an Authorization: Bearer header — the header path matters
because huggingface.co embeds Spaces in an iframe where browsers block
third-party cookies, so the frontend keeps the token in localStorage as well.
"""
import base64
import hashlib
import hmac
import json
import time
from typing import Optional

_MAX_AGE = 7 * 24 * 3600


def _sign(secret: str, username: str) -> str:
    payload = json.dumps({"u": username, "t": int(time.time())})
    sig = hmac.new(secret.encode(), payload.encode(), hashlib.sha256).hexdigest()
    return base64.urlsafe_b64encode(json.dumps({"p": payload, "s": sig}).encode()).decode()


def _verify(secret: str, token: str) -> Optional[str]:
    try:
        obj = json.loads(base64.urlsafe_b64decode(token.encode()))
        expected = hmac.new(secret.encode(), obj["p"].encode(), hashlib.sha256).hexdigest()
        if not hmac.compare_digest(obj["s"], expected):
            return None
        data = json.loads(obj["p"])
        if time.time() - data["t"] > _MAX_AGE:
            return None
        return data["u"]
    except Exception:
        return None


def create_app(service, users, history, *, session_secret: str, static_dir: Optional[str] = None,
               dataset_name: str = "", decks=None, notes=None, ingest=None):
    import re as _re
    import time as _time

    from fastapi import FastAPI, HTTPException, Request, UploadFile
    from fastapi.responses import FileResponse, JSONResponse

    from ..accounts.uploads import upload_path, validate_upload

    app = FastAPI()

    def current_user(request: Request) -> Optional[str]:
        tok = request.cookies.get("sh_session")
        if not tok:
            auth = request.headers.get("authorization", "")
            if auth.lower().startswith("bearer "):
                tok = auth[7:].strip()
        return _verify(session_secret, tok) if tok else None

    def require_user(request: Request) -> str:
        u = current_user(request)
        if not u:
            raise HTTPException(status_code=401, detail="Not logged in")
        return u

    def require_admin(request: Request) -> str:
        u = require_user(request)
        if users.role(u) != "admin":
            raise HTTPException(status_code=403, detail="Admins only")
        return u

    def _result_json(r) -> dict:
        return {"answer": r.answer, "sources_md": r.sources_md, "sources": [s.__dict__ for s in r.sources]}

    def _call_model(fn, *a, **k):
        """Generation can fail for reasons that are nobody's bug (provider daily caps,
        bad failover key, outage) — surface those as a clean 503 the UI can show,
        never a naked 500."""
        try:
            return fn(*a, **k)
        except Exception as e:
            raise HTTPException(
                status_code=503,
                detail="The AI model is unavailable right now (provider limit or outage). "
                       "Try again in a few minutes.") from e

    @app.post("/api/register")
    async def register(req: Request):
        b = await req.json()
        ok, msg = users.register(b.get("invite", ""), b.get("username", ""), b.get("password", ""))
        return JSONResponse({"ok": ok, "message": msg}, status_code=200 if ok else 400)

    @app.post("/api/login")
    async def login(req: Request):
        b = await req.json()
        username = b.get("username", "")
        if not users.verify(username, b.get("password", "")):
            raise HTTPException(status_code=401, detail="Invalid username or password")
        token = _sign(session_secret, username)
        resp = JSONResponse({"ok": True, "username": username, "role": users.role(username), "token": token})
        resp.set_cookie("sh_session", token, httponly=True, samesite="none", secure=True,
                        max_age=_MAX_AGE)
        # Partitioned (CHIPS) keeps the cookie alive inside the huggingface.co iframe.
        resp.headers["set-cookie"] = resp.headers["set-cookie"] + "; Partitioned"
        return resp

    @app.post("/api/logout")
    def logout():
        resp = JSONResponse({"ok": True})
        resp.delete_cookie("sh_session")
        return resp

    @app.get("/api/me")
    def me(req: Request):
        u = current_user(req)
        return {"username": u, "role": users.role(u) if u else None}

    @app.get("/api/modules")
    def modules(req: Request):
        require_user(req)
        return service.modules()

    @app.get("/api/modules/{mid}/topics")
    def topics(mid: str, req: Request):
        require_user(req)
        return service.topics(mid)

    @app.get("/api/topics/{tid}/sources")
    def topic_sources(tid: str, req: Request):
        require_user(req)
        return service.topic_sources(tid)

    @app.get("/api/topics/{tid}/summary")
    def topic_summary(tid: str, req: Request, force: str = ""):
        # Summaries are derived from shared material and are the most expensive
        # generation (map-reduce over the whole topic) — cache class-wide like
        # mind maps. ?force=1 regenerates.
        u = require_user(req)
        key = "app/summaries/" + _re.sub(r"[^A-Za-z0-9._-]+", "_", tid) + ".json"
        if not force:
            cached = users.store.read_text(key)
            if cached:
                try:
                    out = json.loads(cached)
                    history.append(u, {"kind": "topic_summary", "topic": tid, "answer": out["answer"]})
                    return out
                except Exception:
                    pass
        r = _call_model(service.topic_summary, tid)
        history.append(u, {"kind": "topic_summary", "topic": tid, "answer": r.answer})
        out = _result_json(r)
        if not r.answer.startswith(("Unknown topic", "No material")):
            users.store.write_text(key, json.dumps(out, ensure_ascii=False))
        return out

    @app.post("/api/topics/{tid}/ask")
    async def topic_ask(tid: str, req: Request):
        u = require_user(req)
        b = await req.json()
        q = b.get("query", "")
        thread_id = b.get("thread_id") or ""
        if thread_id:
            return _continue_thread(u, thread_id, q, lambda qq, prior: service.topic_ask_thread(tid, qq, prior))
        r = _call_model(service.topic_ask, tid, q)
        eid = history.append(u, {"kind": "topic_ask", "topic": tid, "query": q, "answer": r.answer,
                                 "exchanges": [{"q": q, "a": r.answer}]})
        return {**_result_json(r), "thread_id": eid}

    def _continue_thread(u: str, thread_id: str, q: str, runner) -> dict:
        entry = history.get(u, thread_id)
        if entry is None:
            raise HTTPException(status_code=404, detail="Unknown thread")
        prior = entry.get("exchanges") or (
            [{"q": entry["query"], "a": entry.get("answer", "")}] if entry.get("query") else [])
        r = _call_model(runner, q, prior)
        history.append_exchange(u, thread_id, q, r.answer)
        return {**_result_json(r), "thread_id": thread_id}

    @app.post("/api/ask")
    async def ask(req: Request):
        u = require_user(req)
        b = await req.json()
        q = b.get("query", "")
        thread_id = b.get("thread_id") or ""
        if thread_id:
            return _continue_thread(u, thread_id, q, service.ask_thread)
        r = _call_model(service.ask, q)
        eid = history.append(u, {"kind": "ask", "query": q, "answer": r.answer,
                                 "exchanges": [{"q": q, "a": r.answer}]})
        return {**_result_json(r), "thread_id": eid}

    @app.get("/api/threads/{thread_id}")
    def get_thread(thread_id: str, req: Request):
        u = require_user(req)
        entry = history.get(u, thread_id)
        if entry is None:
            raise HTTPException(status_code=404, detail="Unknown thread")
        return entry

    @app.post("/api/history/pin")
    async def history_pin(req: Request):
        u = require_user(req)
        b = await req.json()
        if not history.set_pinned(u, b.get("id", ""), bool(b.get("pinned", True))):
            raise HTTPException(status_code=404, detail="Unknown entry")
        return {"ok": True}

    @app.post("/api/history/delete")
    async def history_delete(req: Request):
        u = require_user(req)
        if not history.delete(u, (await req.json()).get("id", "")):
            raise HTTPException(status_code=404, detail="Unknown entry")
        return {"ok": True}

    @app.get("/api/history")
    def get_history(req: Request):
        return history.load(require_user(req))

    @app.get("/api/files")
    def files(req: Request):
        require_user(req)
        from ..app.service import _source_type
        return [{"file": f, "type": _source_type(f)} for f in service.list_files()]

    @app.post("/api/upload")
    async def upload(req: Request, file: UploadFile):
        u = require_user(req)
        data = await file.read()
        ok, name_or_err = validate_upload(file.filename or "", len(data))
        if not ok:
            raise HTTPException(status_code=400, detail=name_or_err)
        path = upload_path(u, name_or_err)
        users.store.save_bytes(path, data)
        history.append(u, {"kind": "upload", "file": name_or_err})
        if ingest is not None:
            ingest.enqueue(path, data)
            return {"ok": True, "message": "Uploaded — indexing now, searchable in a few minutes.",
                    "path": path}
        return {"ok": True, "message": "Uploaded — searchable after the next index rebuild.", "path": path}

    @app.get("/api/uploads/mine")
    def my_uploads(req: Request):
        u = require_user(req)
        prefix = f"uploads/{u}/"
        return [{"path": p, "name": p.split("/")[-1],
                 "status": ingest.status(p) if ingest is not None else "pending-rebuild"}
                for p in users.store.list_files(prefix)]

    @app.get("/api/admin/overview")
    def admin_overview(req: Request):
        require_admin(req)
        from ..app.service import _source_type
        uploads = [{"path": p, "user": p.split("/")[1] if p.count("/") >= 2 else "?", "name": p.split("/")[-1]}
                   for p in users.store.list_files("uploads/")]
        for up in uploads:
            up["status"] = ingest.status(up["path"]) if ingest is not None else "pending-rebuild"
        return {"invite": users.invite_code, "dataset": dataset_name,
                "users": users.users_with_roles(), "seed": sorted(users.seed),
                "invites": users.invites.rows() if users.invites else [],
                "originals": [{"file": f, "type": _source_type(f)} for f in service.list_files()],
                "uploads": uploads}

    @app.post("/api/admin/invites")
    async def admin_mint_invite(req: Request):
        require_admin(req)
        if not users.invites:
            raise HTTPException(status_code=503, detail="Invite links not configured")
        label = ((await req.json()).get("label") or "").strip()
        if not label:
            raise HTTPException(status_code=400, detail="Give the link a name (who is it for?)")
        token = users.invites.mint(label)
        return {"token": token, "path": f"/?invite={token}"}

    @app.post("/api/admin/invites/revoke")
    async def admin_revoke_invite(req: Request):
        require_admin(req)
        if not users.invites:
            raise HTTPException(status_code=503, detail="Invite links not configured")
        if not users.invites.revoke((await req.json()).get("token", "")):
            raise HTTPException(status_code=404, detail="Unknown or already-used link")
        return {"ok": True}

    @app.post("/api/admin/role")
    async def admin_role(req: Request):
        require_admin(req)
        b = await req.json()
        if not users.set_role(b.get("username", ""), b.get("role", "")):
            raise HTTPException(status_code=400, detail="Unknown user or invalid role")
        return {"ok": True}

    @app.post("/api/admin/remove")
    async def admin_remove(req: Request):
        require_admin(req)
        users.remove((await req.json()).get("username", ""))
        return {"ok": True}

    @app.post("/api/admin/uploads/delete")
    async def admin_delete_upload(req: Request):
        require_admin(req)
        path = (await req.json()).get("path", "")
        if not path.startswith("uploads/") or ".." in path:
            raise HTTPException(status_code=400, detail="Only files under uploads/ can be deleted here")
        users.store.delete(path)
        return {"ok": True}

    def _topic_title(tid: str) -> str:
        try:
            t = service.graph.topic_by_id(tid)
            return t.title if t else tid
        except AttributeError:
            return tid

    @app.post("/api/topics/{tid}/quiz")
    async def topic_quiz(tid: str, req: Request):
        u = require_user(req)
        n = int((await req.json()).get("n", 5) or 5)
        r = _call_model(service.topic_quiz, tid, n=max(1, min(n, 10)))
        history.append(u, {"kind": "quiz", "topic": tid})
        return r

    @app.post("/api/topics/{tid}/flashcards")
    async def topic_flashcards(tid: str, req: Request):
        u = require_user(req)
        if decks is None:
            raise HTTPException(status_code=503, detail="Flashcards not configured")
        n = int((await req.json()).get("n", 10) or 10)
        gen = _call_model(service.topic_flashcards, tid, n=max(1, min(n, 20)))
        now = _time.time()
        added = decks.add_cards(u, tid, _topic_title(tid), gen.get("cards") or [], now=now)
        history.append(u, {"kind": "flashcards", "topic": tid})
        out = {"added": added, **decks.stats(u, now)}
        if gen.get("error"):
            out["error"] = gen["error"]
        return out

    @app.get("/api/flashcards")
    def flashcards_due(req: Request):
        u = require_user(req)
        if decks is None:
            raise HTTPException(status_code=503, detail="Flashcards not configured")
        now = _time.time()
        due = [{"id": c["id"], "topic": c.get("topic", ""), "front": c["front"], "back": c["back"]}
               for c in decks.due_cards(u, now)]
        return {"due": due, **decks.stats(u, now)}

    @app.post("/api/flashcards/grade")
    async def flashcards_grade(req: Request):
        u = require_user(req)
        if decks is None:
            raise HTTPException(status_code=503, detail="Flashcards not configured")
        b = await req.json()
        d = decks.grade(u, b.get("id", ""), bool(b.get("good")), _time.time())
        if d < 0:
            raise HTTPException(status_code=404, detail="Unknown card")
        return {"ok": True, "next_due_days": d}

    @app.get("/api/topics/{tid}/mindmap")
    def topic_mindmap(tid: str, req: Request):
        require_user(req)
        # Mind maps are derived from shared material, so the first request renders
        # and caches for the whole class.
        key = "app/mindmaps/" + _re.sub(r"[^A-Za-z0-9._-]+", "_", tid) + ".json"
        cached = users.store.read_text(key)
        if cached:
            try:
                return json.loads(cached)
            except Exception:
                pass
        r = _call_model(service.topic_mindmap, tid)
        if not r.get("error"):
            users.store.write_text(key, json.dumps(r, ensure_ascii=False))
        return r

    @app.get("/api/notes")
    def notes_list(req: Request):
        u = require_user(req)
        if notes is None:
            raise HTTPException(status_code=503, detail="Notes not configured")
        return notes.list_notes(u)

    @app.post("/api/notes")
    async def notes_add(req: Request):
        u = require_user(req)
        if notes is None:
            raise HTTPException(status_code=503, detail="Notes not configured")
        b = await req.json()
        title, body = (b.get("title") or "").strip(), (b.get("body") or "").strip()
        if not title and not body:
            raise HTTPException(status_code=400, detail="Empty note")
        return {"id": notes.add(u, title, body, b.get("topic_id") or "", now=_time.time())}

    @app.post("/api/notes/update")
    async def notes_update(req: Request):
        u = require_user(req)
        if notes is None:
            raise HTTPException(status_code=503, detail="Notes not configured")
        b = await req.json()
        if not notes.update(u, b.get("id", ""), (b.get("title") or "").strip(),
                            (b.get("body") or "").strip(), now=_time.time()):
            raise HTTPException(status_code=404, detail="Unknown note")
        return {"ok": True}

    @app.post("/api/notes/delete")
    async def notes_delete(req: Request):
        u = require_user(req)
        if notes is None:
            raise HTTPException(status_code=503, detail="Notes not configured")
        if not notes.delete(u, (await req.json()).get("id", "")):
            raise HTTPException(status_code=404, detail="Unknown note")
        return {"ok": True}

    @app.get("/api/download")
    def download(file: str, req: Request, t: str = ""):
        # New-tab navigation can't carry the Authorization header, so the citation
        # links pass the session token as ?t= instead.
        u = current_user(req) or (_verify(session_secret, t) if t else None)
        if not u:
            raise HTTPException(status_code=401, detail="Not logged in")
        path = service.download_path(file)
        if not path:
            raise HTTPException(status_code=404, detail="Not found")
        return FileResponse(path)

    if static_dir:
        from fastapi.staticfiles import StaticFiles
        app.mount("/static", StaticFiles(directory=static_dir), name="static")

        @app.get("/")
        def index():
            return FileResponse(f"{static_dir}/index.html")

    return app