LLeMbas/src/lembas/api/pages.py

"""Full-page routes: the chat shell and the user's own settings."""

from __future__ import annotations

from fastapi import APIRouter, HTTPException, Request, status
from fastapi.responses import RedirectResponse
from sqlalchemy import select
from sqlalchemy.orm import Session as DBSession

from lembas.api.deps import Db, RequiredUser
from lembas.db.models import Chat, Folder, Message, User
from lembas.security import permissions
from lembas.services import chat as chat_service
from lembas.services.markdown import render_markdown
from lembas.web.templating import render

router = APIRouter(tags=["pages"])


def _chat_context(db: DBSession, user: User, chat: Chat | None) -> dict:
    """Model lists and permissions every chat page needs.

    Pinned and unpinned are split here rather than in the template so the
    picker's optgroups stay a plain loop.
    """
    models = chat_service.available_models(db, user)
    current = next((m for m in models if m.model_id == chat.model_id), None) if chat else None
    pinned = [m for m in models if m.pinned]
    return {
        "models": models,
        "pinned_models": pinned,
        # Excludes the pinned ones: they already have their own optgroup, and
        # listing a model twice gives the <select> two options with the same
        # value, both marked selected.
        "other_models": [m for m in models if not m.pinned] if pinned else models,
        "current_model": current,
    }


def _sidebar_context(db: DBSession, user: User) -> dict:
    """Folder tree plus the chats that belong to no folder.

    Only root folders are queried; children come through the relationship and
    render recursively in the template.
    """
    folders = list(
        db.scalars(
            select(Folder)
            .where(Folder.user_id == user.id, Folder.parent_id.is_(None))
            .order_by(Folder.position, Folder.name)
        )
    )
    unfiled = list(
        db.scalars(
            select(Chat)
            .where(
                Chat.user_id == user.id,
                Chat.folder_id.is_(None),
                Chat.archived.is_(False),
            )
            .order_by(Chat.pinned.desc(), Chat.updated_at.desc())
        )
    )
    return {
        "folders": folders,
        "unfiled_chats": unfiled,
        "can": permissions.resolve(db, user),
    }


@router.get("/")
async def home(user: RequiredUser):
    return RedirectResponse("/chat", status_code=status.HTTP_303_SEE_OTHER)


@router.get("/chat")
async def chat_index(request: Request, db: Db, user: RequiredUser):
    return render(
        request,
        "chat/index.html",
        {
            "chat": None,
            "messages": [],
            **_chat_context(db, user, None),
            **_sidebar_context(db, user),
        },
    )


@router.get("/chat/{chat_id}")
async def chat_detail(request: Request, db: Db, user: RequiredUser, chat_id: str):
    chat = db.get(Chat, chat_id)
    if chat is None or chat.user_id != user.id:
        raise HTTPException(status.HTTP_404_NOT_FOUND, "That chat no longer exists.")

    messages = list(
        db.scalars(
            select(Message).where(Message.chat_id == chat.id).order_by(Message.created_at)
        )
    )

    # Markdown is rendered once here rather than in the template so the same
    # helper produces the page and the streamed final frame -- one code path,
    # no chance of the two disagreeing.
    bodies = {
        message.id: render_markdown(message.content)
        for message in messages
        if message.role == "assistant" and message.content
    }

    return render(
        request,
        "chat/index.html",
        {
            "chat": chat,
            "messages": messages,
            "bodies": bodies,
            **_chat_context(db, user, chat),
            **_sidebar_context(db, user),
        },
    )


@router.get("/settings")
async def settings_page(
    request: Request,
    db: Db,
    user: RequiredUser,
    error: str = "",
    saved: str = "",
):
    # error/saved arrive as query parameters because the password form redirects
    # back here: a POST that re-rendered in place would re-submit on refresh.
    return render(
        request,
        "settings.html",
        {
            "chat": None,
            "error": error,
            "saved": saved,
            **_chat_context(db, user, None),
            **_sidebar_context(db, user),
        },
    )