ObsiGate/backend/bookslm_routes.py

"""BooksLM API routes — directory-scoped AI chat for Obsidian vaults."""

import json
import logging
from pathlib import Path
from typing import Any

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from backend.agent.loop import run_agent
from backend.ai_chat import chat_completion, stream_completion
from backend.auth.middleware import check_vault_access, require_auth
from backend.bookslm import (
    build_general_system_prompt,
    build_system_prompt,
    collect_adhoc_context,
    collect_directory_context,
    collect_files_context,
    empty_context,
    load_vault_image_data_url,
    merge_contexts,
)
from backend.indexer import get_vault_data, index
from backend.model_capabilities import model_supports_vision
from backend.schemas import BooksLMContextResponse
from backend.skills import get_skill_prompt
from backend.tools.api import ToolContext, ToolMode

logger = logging.getLogger("obsigate.bookslm_routes")
router = APIRouter(prefix="/api/ai/bookslm", tags=["BooksLM"])

VALID_MODES = {"directory", "documents", "general"}


# ── Request models ──


class BooksLMContextRequest(BaseModel):
    vault: str | None = Field(default=None, description="Vault name (required for directory/documents modes)")
    directory: str = Field(default="", description="Relative directory path within the vault")
    mode: str = Field(default="directory", description="Context mode: 'directory', 'documents' or 'general'")
    context_files: list[str] = Field(
        default_factory=list,
        description="Relative file paths to use as context (documents mode)",
    )
    extra_files: list[str] = Field(
        default_factory=list,
        description="Ad-hoc files added with the '@' command (any mode)",
    )
    extra_directories: list[str] = Field(
        default_factory=list,
        description="Ad-hoc directories added with the '@' command (any mode)",
    )
    app_context: dict[str, Any] | None = Field(
        default=None,
        description="Live client UI state for the General assistant: open_documents, "
                    "current_path, directory, vault, search_query, search_total, search_results.",
    )


class BooksLMChatRequest(BaseModel):
    vault: str | None = Field(default=None, description="Vault name (required for directory/documents modes)")
    directory: str = Field(default="", description="Relative directory path within the vault")
    mode: str = Field(default="directory", description="Context mode: 'directory', 'documents' or 'general'")
    context_files: list[str] = Field(
        default_factory=list,
        description="Relative file paths to use as context (documents mode)",
    )
    extra_files: list[str] = Field(
        default_factory=list,
        description="Ad-hoc files added with the '@' command (any mode)",
    )
    extra_directories: list[str] = Field(
        default_factory=list,
        description="Ad-hoc directories added with the '@' command (any mode)",
    )
    message: str = Field(description="User message")
    images: list[dict[str, Any]] = Field(
        default_factory=list,
        description="Images for vision models. Each item: {data, mime_type} (pasted) "
                    "or {path} (vault-relative file).",
    )
    skill: str | None = Field(
        default=None,
        description="Skill id selected with the '/' command; its prompt is added to the system prompt.",
    )
    conversation_history: list[dict[str, str]] = Field(
        default_factory=list,
        description="Previous conversation turns [{role, content}]",
    )
    provider: str | None = Field(
        default=None,
        description="AI provider override (e.g. 'deepseek', 'openrouter', 'gemini', 'nvidia', 'xiaomi', 'mistral', 'qwencloud'). "
                    "If not set, uses DEFAULT_PROVIDER.",
    )
    model: str | None = Field(
        default=None,
        description="Model name to use for this request. If not set, uses the provider's default model.",
    )
    confirm: dict[str, Any] | None = Field(
        default=None,
        description="Pending tool confirmation to apply (two-step propose/apply). "
                    "Shape: the ``error`` object of a previous ``confirmation`` event.",
    )
    confirm_messages: list[dict[str, Any]] | None = Field(
        default=None,
        description="Conversation snapshot returned alongside a ``confirmation`` event, "
                    "echoed back to resume the agent run.",
    )
    app_context: dict[str, Any] | None = Field(
        default=None,
        description="Live client UI state for the General assistant: open_documents, "
                    "current_path, directory, vault, search_query, search_total, search_results.",
    )


def _normalize_mode(mode: str | None) -> str:
    mode = (mode or "directory").lower()
    return mode if mode in VALID_MODES else "directory"


def _resolve_vault_path(vault: str | None, current_user):
    """Validate vault access and return (vault_name, vault_path)."""
    if not vault:
        raise HTTPException(status_code=400, detail="Champ 'vault' requis pour ce contexte")
    if not check_vault_access(vault, current_user):
        raise HTTPException(status_code=403, detail=f"Accès refusé à la vault '{vault}'")
    vault_data = get_vault_data(vault)
    if not vault_data:
        raise HTTPException(status_code=404, detail=f"Vault '{vault}' not found")
    return vault, Path(vault_data["path"])


def _build_context(
    mode: str,
    vault_path: Path | None,
    directory: str,
    context_files: list[str],
    extra_files: list[str] | None = None,
    extra_directories: list[str] | None = None,
):
    """Collect the context payload for the requested mode.

    Ad-hoc files/directories (``@`` command) are merged on top of the base
    context. In General mode, attaching ad-hoc files promotes the effective
    scope to ``documents`` so the prompt actually includes their content.
    """
    if mode == "general":
        base = empty_context("general")
    elif mode == "documents":
        ctx = collect_files_context(vault_path, context_files, scope="documents")  # type: ignore[arg-type]
        # No open document could be read → degrade gracefully to General.
        base = ctx if ctx["file_count"] else empty_context("general")
    else:
        base = collect_directory_context(vault_path, directory)  # type: ignore[arg-type]

    if (extra_files or extra_directories) and vault_path is not None:
        adhoc_scope = "documents" if base.get("scope") == "general" else base.get("scope", "directory")
        adhoc = collect_adhoc_context(vault_path, extra_files, extra_directories, scope=adhoc_scope)
        if adhoc["file_count"]:
            base = merge_contexts(base, adhoc)
    return base


def _submitted_app_context(req) -> dict[str, Any] | None:
    """Return the client-submitted live app state (may be None)."""
    ctx = getattr(req, "app_context", None)
    return ctx if isinstance(ctx, dict) else None


def _recent_files_for_prompt(current_user, limit: int = 10) -> list[dict[str, Any]]:
    """Best-effort list of the user's most recently modified files.

    Used to give the General assistant a sense of what the user has been
    working on. Never raises: any failure yields an empty list.
    """
    try:
        from backend.services.recent import list_recent

        username = current_user.get("username") if isinstance(current_user, dict) else None
        user_vaults = (
            current_user.get("_token_vaults") or current_user.get("vaults", [])
            if isinstance(current_user, dict)
            else []
        )
        data = list_recent(username, user_vaults, limit=limit, mode="modified")
        return list(data.get("files", []))
    except Exception:  # pragma: no cover - defensive, prompt enrichment only
        logger.debug("Could not gather recent files for assistant prompt", exc_info=True)
        return []


def _resolve_system_prompt(req, current_user) -> str:
    """Resolve the vault access and build the assistant system prompt.

    Shared by the classic chat endpoint and the tool-calling agent endpoint.
    """
    mode = _normalize_mode(req.mode)
    vault_path: Path | None = None
    if mode != "general":
        _, vault_path = _resolve_vault_path(req.vault, current_user)
    elif getattr(req, "extra_files", None) or getattr(req, "extra_directories", None):
        # General mode has no base context, but ad-hoc files/directories added
        # with `@` still need a vault to be read from.
        vault_path = _resolve_optional_vault_path(req, current_user)

    context = _build_context(
        mode,
        vault_path,
        req.directory,
        req.context_files,
        getattr(req, "extra_files", None),
        getattr(req, "extra_directories", None),
    )
    effective_mode = context.get("scope", mode)

    if effective_mode == "general":
        prompt = build_general_system_prompt(
            list(index.keys()),
            app_context=_submitted_app_context(req),
            recent_files=_recent_files_for_prompt(current_user),
        )
    elif effective_mode == "documents":
        prompt = build_system_prompt(context, scope="documents", vault_name=req.vault)
    elif context["file_count"] == 0:
        # Empty (or unreadable) directory: don't block the request. Answer as
        # the General assistant would, telling the model the folder is empty so
        # it can still help (create a file, explain the app, etc.).
        prompt = build_general_system_prompt(
            list(index.keys()),
            app_context=_submitted_app_context(req),
            recent_files=_recent_files_for_prompt(current_user),
        )
        prompt += (
            f"\n## Dossier vide\nLe dossier « {req.directory or '/'} » "
            f"(vault {req.vault}) ne contient aucun fichier markdown exploitable. "
            "Réponds quand même à la demande de l'utilisateur sans contexte "
            "documentaire, et propose une action de création si c'est pertinent.\n"
        )
    else:
        prompt = build_system_prompt(context, scope="directory", vault_name=req.vault)

    skill_id = getattr(req, "skill", None)
    if skill_id:
        skill_prompt = get_skill_prompt(skill_id, current_user)
        if skill_prompt:
            prompt += "\n\n## Skill actif\n" + skill_prompt
    return prompt


def _resolve_optional_vault_path(req, current_user) -> Path | None:
    """Best-effort vault path resolution (never raises)."""
    if not getattr(req, "vault", None):
        return None
    try:
        _, vault_path = _resolve_vault_path(req.vault, current_user)
        return vault_path
    except HTTPException:
        return None


def _build_user_content(req, vault_path: Path | None):
    """Build the user message content, attaching images when present.

    Returns a plain string when there is no image, otherwise an OpenAI-style
    multimodal content array (which ``ai_chat`` adapts for Gemini).
    """
    images = getattr(req, "images", None) or []
    if not images:
        return req.message

    parts: list[dict[str, Any]] = [{"type": "text", "text": req.message}]
    for image in images:
        if not isinstance(image, dict):
            continue
        data_url: str | None = None
        if image.get("data"):
            mime = image.get("mime_type") or "image/png"
            data_url = f"data:{mime};base64,{image['data']}"
        elif image.get("path") and vault_path is not None:
            data_url = load_vault_image_data_url(vault_path, str(image["path"]))
        if data_url:
            parts.append({"type": "image_url", "image_url": {"url": data_url}})
    return parts


def _validate_vision_support(req) -> None:
    """Reject image requests when the selected model cannot analyse images."""
    if not (getattr(req, "images", None)):
        return
    provider = _resolve_provider_name(req.provider)
    if not provider:
        return
    from backend.ai import PROVIDERS

    model = req.model or PROVIDERS.get(provider, {}).get("model", "")
    if not model_supports_vision(provider, model):
        raise HTTPException(
            status_code=400,
            detail=f"Le modèle '{model or provider}' ne supporte pas l'analyse d'images. "
                   "Choisissez un modèle compatible vision.",
        )


def _resolve_provider_name(requested: str | None) -> str | None:
    """Pick the provider to use: explicit override, else first available."""
    from backend.ai import DEFAULT_PROVIDER, PROVIDERS

    cfg_name = (requested or DEFAULT_PROVIDER).lower()
    if cfg_name in PROVIDERS and PROVIDERS[cfg_name].get("api_key"):
        return cfg_name
    for pname, pcfg in PROVIDERS.items():
        if pcfg.get("api_key") and pname != "gemini":
            return pname
    return None


def _effective_model(provider: str | None, requested: str | None) -> str:
    """Model actually used for a request.

    The client may leave `model` empty (provider default) — reporting the raw
    request would show nothing in the "provider · model" tag, so the provider's
    configured default is returned instead.
    """
    if requested:
        return requested
    if not provider:
        return ""
    from backend.ai import PROVIDERS

    return PROVIDERS.get(provider, {}).get("model", "") or ""


def _tool_sources(rec) -> list[dict[str, str]]:
    """Compact web sources of a tool result (rendered as links in the UI).

    Only the web tools produce sources: ``web_search`` returns ranked results,
    ``fetch_url`` a single page. Everything else yields an empty list so the
    SSE payload stays small.
    """
    data = rec.result if isinstance(rec.result, dict) else {}
    name = rec.name or ""
    sources: list[dict[str, str]] = []
    if name == "web_search":
        for item in (data.get("results") or [])[:8]:
            if not isinstance(item, dict):
                continue
            url = item.get("url") or ""
            if not url:
                continue
            sources.append({"title": item.get("title") or url, "url": url})
    elif name == "fetch_url":
        url = data.get("url") or ""
        if url:
            sources.append({"title": data.get("title") or url, "url": url})
    return sources


def _tool_event_sse(rec) -> str:
    """Serialize one executed tool call as an SSE ``tool`` event."""
    payload = json.dumps(
        {
            "name": rec.name,
            "ok": rec.ok,
            "arguments": rec.arguments,
            "step": rec.step,
            "sources": _tool_sources(rec),
        },
        ensure_ascii=False,
    )
    return f"event: tool\ndata: {payload}\n\n"


def _thought_event_sse(note: dict) -> str:
    """Serialize one intermediate reasoning note as an SSE ``step`` event."""
    payload = json.dumps({"step": note}, ensure_ascii=False)
    return f"event: step\ndata: {payload}\n\n"


# ── Endpoints ──


@router.post("/context", response_model=BooksLMContextResponse)
async def api_bookslm_context(
    req: BooksLMContextRequest,
    current_user=Depends(require_auth),
):
    """Collect context for the AI assistant.

    Three modes are supported:

    * ``directory`` — every supported file under a vault directory;
    * ``documents`` — only the explicitly listed open documents;
    * ``general`` — no document context (assistant for the app itself).
    """
    mode = _normalize_mode(req.mode)
    if mode == "general":
        if not (req.extra_files or req.extra_directories):
            return empty_context("general")
        vault_path = _resolve_optional_vault_path(req, current_user)
    else:
        _, vault_path = _resolve_vault_path(req.vault, current_user)
    return _build_context(
        mode, vault_path, req.directory, req.context_files,
        req.extra_files, req.extra_directories,
    )


@router.post(
    "/chat",
    response_class=StreamingResponse,
    responses={200: {"content": {"text/event-stream": {}}, "description": "SSE token stream"}},
)
async def api_bookslm_chat(
    req: BooksLMChatRequest,
    current_user=Depends(require_auth),
):
    """Chat with the AI assistant.

    Builds a system prompt from the selected context (directory, open
    documents or general app knowledge), then streams the provider's answer
    as Server-Sent Events.
    """
    _validate_vision_support(req)
    system_prompt = _resolve_system_prompt(req, current_user)
    vault_path = _resolve_optional_vault_path(req, current_user)

    messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]
    for turn in req.conversation_history:
        role = turn.get("role")
        content = turn.get("content", "")
        if role in ("user", "assistant") and content:
            messages.append({"role": role, "content": content})
    messages.append({"role": "user", "content": _build_user_content(req, vault_path)})

    async def generate_sse():
        try:
            # Resolve provider: explicit override wins, else first available.
            cfg_name = _resolve_provider_name(req.provider)
            if cfg_name is None:
                err = "Aucun fournisseur AI configuré (clés API manquantes)"
                error_data = json.dumps({"error": err}, ensure_ascii=False)
                yield f"event: error\ndata: {error_data}\n\n"
                return

            # Stream token deltas as they arrive from the provider.
            async for token in stream_completion(
                messages,
                provider=cfg_name,
                model=req.model,
                temperature=0.3,
                max_tokens=4096,
            ):
                data = json.dumps(
                    {
                        "token": token,
                        "provider": cfg_name,
                        "model": _effective_model(cfg_name, req.model),
                    },
                    ensure_ascii=False,
                )
                yield f"event: message\ndata: {data}\n\n"
            yield "event: done\ndata: {}\n\n"
        except Exception as e:
            logger.error(f"BooksLM chat error: {e}")
            error_data = json.dumps({"error": str(e)}, ensure_ascii=False)
            yield f"event: error\ndata: {error_data}\n\n"

    return StreamingResponse(
        generate_sse(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )


@router.post(
    "/agent",
    response_class=StreamingResponse,
    responses={200: {"content": {"text/event-stream": {}}, "description": "SSE tool/agent stream"}},
)
async def api_bookslm_agent(
    req: BooksLMChatRequest,
    current_user=Depends(require_auth),
):
    """Chat with the tool-calling agent.

    Same context as ``/chat`` but the model may call tools (read/search the
    vault) through the shared tool layer. Emits one ``tool`` event per executed
    tool call, then a final ``message`` event. Mutating tools pause the run with
    a ``confirmation`` event (two-step propose/apply) carrying the pending call
    and the conversation snapshot; the client resumes by echoing them back in
    ``confirm`` / ``confirm_messages``.
    """
    _validate_vision_support(req)
    system_prompt = _resolve_system_prompt(req, current_user)
    vault_path = _resolve_optional_vault_path(req, current_user)

    messages: list[dict] = [{"role": "system", "content": system_prompt}]
    for turn in req.conversation_history:
        role = turn.get("role")
        content = turn.get("content", "")
        if role in ("user", "assistant") and content:
            messages.append({"role": role, "content": content})
    messages.append({"role": "user", "content": _build_user_content(req, vault_path)})

    ctx = ToolContext(user=current_user, mode=ToolMode.IN_APP)

    async def _llm(msgs, tool_schemas):
        return await chat_completion(
            msgs,
            tools=tool_schemas,
            provider=req.provider,
            model=req.model,
            temperature=0.3,
            max_tokens=4096,
        )

    async def generate_sse():
        import asyncio

        try:
            cfg_name = _resolve_provider_name(req.provider)
            if cfg_name is None:
                error_data = json.dumps(
                    {"error": "Aucun fournisseur AI configuré (clés API manquantes)"},
                    ensure_ascii=False,
                )
                yield f"event: error\ndata: {error_data}\n\n"
                return

            # Stream tool events live: each executed step is pushed on the
            # queue by the loop callback and emitted as soon as it happens,
            # so the UI can grow its « N steps » block while thinking.
            queue: asyncio.Queue = asyncio.Queue()

            def _on_tool(rec) -> None:
                queue.put_nowait(("tool", rec))

            run_task = asyncio.create_task(run_agent(
                messages,
                ctx=ctx,
                llm=_llm,
                resume_messages=req.confirm_messages,
                confirm_pending=req.confirm,
                on_tool_call=_on_tool,
                on_thought=lambda note: queue.put_nowait(("thought", note)),
            ))

            # Drain every completed step as soon as it lands, while the agent
            # keeps running in the background. If the client disconnects, the
            # generator is cancelled: release the run so it cannot orphan.
            try:
                while True:
                    try:
                        kind, item = await asyncio.wait_for(queue.get(), timeout=0.25)
                    except asyncio.TimeoutError:
                        if run_task.done():
                            break
                        continue
                    yield _tool_event_sse(item) if kind == "tool" else _thought_event_sse(item)
                while not queue.empty():
                    kind, item = queue.get_nowait()
                    yield _tool_event_sse(item) if kind == "tool" else _thought_event_sse(item)
                result = run_task.result()
            finally:
                if not run_task.done():
                    run_task.cancel()

            if result.stopped == "confirmation_required":
                pending = json.dumps(
                    {"pending": result.pending or {}, "messages": result.messages},
                    ensure_ascii=False,
                )
                yield f"event: confirmation\ndata: {pending}\n\n"
            else:
                data = json.dumps(
                    {
                        "token": result.content,
                        "provider": cfg_name,
                        "model": _effective_model(cfg_name, req.model),
                        "iterations": result.iterations,
                        "stopped": result.stopped,
                    },
                    ensure_ascii=False,
                )
                yield f"event: message\ndata: {data}\n\n"
            yield "event: done\ndata: {}\n\n"
        except Exception as e:
            logger.error(f"BooksLM agent error: {e}")
            error_data = json.dumps({"error": str(e)}, ensure_ascii=False)
            yield f"event: error\ndata: {error_data}\n\n"

    return StreamingResponse(
        generate_sse(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )