ObsiGate/backend/bookslm.py

"""BooksLM — Context collection and caching for directory-scoped AI chat.

Collects markdown files from an Obsidian vault directory, applies secret
redaction, builds a system prompt with file contents, and caches results
for repeated queries.
"""

import base64
import hashlib
import json
import logging
import mimetypes
import os
import time
from pathlib import Path
from typing import Any

from backend.secret_redactor import redact_file_content

logger = logging.getLogger("obsigate.bookslm")

# ── Configuration limits ──
BOOKSLM_MAX_FILES = int(os.getenv("BOOKSLM_MAX_FILES", "200"))
BOOKSLM_MAX_TOTAL_CHARS = int(os.getenv("BOOKSLM_MAX_TOTAL_CHARS", "200000"))
BOOKSLM_MAX_FILE_CHARS = int(os.getenv("BOOKSLM_MAX_FILE_CHARS", "30000"))
# Maximum size of an image sent to a vision model (bytes, before base64).
BOOKSLM_MAX_IMAGE_BYTES = int(os.getenv("BOOKSLM_MAX_IMAGE_BYTES", "10000000"))

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".avif"}

# ── Cache ──
_cache: dict[str, dict[str, Any]] = {}
_CACHE_TTL = 300  # seconds


def _cache_key(vault_path: Path, directory: str, file_mtimes: list[tuple[str, float]]) -> str:
    """Build a SHA-256 cache key from vault+directory+file modification times."""
    raw = json.dumps({
        "vault": str(vault_path),
        "dir": directory,
        "mtimes": sorted(file_mtimes),
    }, sort_keys=True)
    return hashlib.sha256(raw.encode()).hexdigest()


def _should_skip(name: str) -> bool:
    """Return True if this file/directory name should be skipped."""
    skip_prefixes = (".",)
    skip_names = {"_attachments", "node_modules", ".git", ".obsidian", "__pycache__"}
    if name in skip_names:
        return True
    return bool(any(name.startswith(p) for p in skip_prefixes))


def _file_priority(path: Path) -> tuple[int, float]:
    """Sort key: README/index first, then by modification time descending.

    Returns (priority_group, -mtime) so that:
    - Group 0: README* and index* files (come first)
    - Group 1: all other files (come after)
    Within each group, newer files come first.
    """
    name_lower = path.stem.lower()
    if name_lower.startswith("readme") or name_lower.startswith("index"):
        group = 0
    else:
        group = 1
    try:
        mtime = path.stat().st_mtime
    except OSError:
        mtime = 0.0
    return (group, -mtime)


def collect_directory_context(vault_path: Path, directory: str) -> dict[str, Any]:
    """Walk a directory recursively, collect .md files with content.

    Args:
        vault_path: Absolute path to the vault root.
        directory: Relative directory path within the vault (empty = root).

    Returns:
        Dict with keys: files, total_chars, file_count, directory_tree.
    """
    target_dir = (vault_path / directory).resolve() if directory else vault_path.resolve()
    vault_resolved = vault_path.resolve()

    # Safety: ensure target is within vault
    try:
        target_dir.relative_to(vault_resolved)
    except ValueError:
        logger.warning(f"Directory outside vault: {target_dir}")
        return {"files": [], "total_chars": 0, "file_count": 0, "directory_tree": ""}

    if not target_dir.exists() or not target_dir.is_dir():
        return {"files": [], "total_chars": 0, "file_count": 0, "directory_tree": ""}

    # Check cache
    file_mtimes: list[tuple[str, float]] = []
    md_files: list[Path] = []
    try:
        for p in target_dir.rglob("*"):
            # Skip hidden dirs/files and special dirs
            parts = p.relative_to(target_dir).parts
            if any(_should_skip(part) for part in parts):
                continue
            if p.is_file() and p.suffix.lower() == ".md":
                md_files.append(p)
                try:
                    file_mtimes.append((str(p.relative_to(target_dir)), p.stat().st_mtime))
                except OSError:
                    file_mtimes.append((str(p.relative_to(target_dir)), 0.0))
    except PermissionError:
        logger.warning(f"Permission denied scanning {target_dir}")
        return {"files": [], "total_chars": 0, "file_count": 0, "directory_tree": ""}

    key = _cache_key(vault_resolved, directory, file_mtimes)
    if key in _cache:
        cached = _cache[key]
        if time.time() - cached.get("_ts", 0) < _CACHE_TTL:
            logger.debug(f"Cache hit for {directory}")
            return {k: v for k, v in cached.items() if k != "_ts"}

    # Sort by priority: README/index first, then by mtime descending
    md_files.sort(key=_file_priority)

    # Apply limits
    collected: list[dict[str, Any]] = []
    total_chars = 0

    for p in md_files:
        if len(collected) >= BOOKSLM_MAX_FILES:
            break
        if total_chars >= BOOKSLM_MAX_TOTAL_CHARS:
            break

        rel_path = str(p.relative_to(vault_resolved)).replace("\\", "/")
        try:
            content = p.read_text(encoding="utf-8", errors="replace")
        except Exception as e:
            logger.warning(f"Cannot read {rel_path}: {e}")
            continue

        # Redact secrets
        content = redact_file_content(content, rel_path)

        # Truncate if too long
        if len(content) > BOOKSLM_MAX_FILE_CHARS:
            content = content[:BOOKSLM_MAX_FILE_CHARS] + "\n\n[... tronqué]"

        remaining = BOOKSLM_MAX_TOTAL_CHARS - total_chars
        if len(content) > remaining:
            content = content[:remaining] + "\n\n[... tronqué]"

        title = p.stem.replace("-", " ").replace("_", " ").title()
        file_type = "markdown"

        collected.append({
            "path": rel_path,
            "title": title,
            "content": content,
            "type": file_type,
        })
        total_chars += len(content)

    # Build directory tree
    dir_tree = _build_directory_tree(target_dir, vault_resolved)

    result = {
        "files": collected,
        "total_chars": total_chars,
        "file_count": len(collected),
        "directory_tree": dir_tree,
        "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS,
        "max_files": BOOKSLM_MAX_FILES,
        "scope": "directory",
    }

    # Store in cache
    _cache[key] = {**result, "_ts": time.time()}
    logger.info(f"Collected {len(collected)} files ({total_chars} chars) from {directory or '/'}")
    return result


def _file_entry(target: Path, rel_path: str, remaining: int) -> dict[str, Any] | None:
    """Read, redact and truncate a single file into a context entry."""
    suffix = target.suffix.lower()
    try:
        if suffix == ".pdf":
            from backend.pdf_reader import extract_pdf_text

            content = extract_pdf_text(target)
            file_type = "pdf"
        else:
            content = target.read_text(encoding="utf-8", errors="replace")
            file_type = "markdown"
    except Exception as e:
        logger.warning(f"Cannot read {rel_path}: {e}")
        return None

    content = redact_file_content(content, rel_path)
    if len(content) > BOOKSLM_MAX_FILE_CHARS:
        content = content[:BOOKSLM_MAX_FILE_CHARS] + "\n\n[... tronqué]"
    if len(content) > remaining:
        content = content[:remaining] + "\n\n[... tronqué]"

    return {
        "path": rel_path,
        "title": target.stem.replace("-", " ").replace("_", " ").title(),
        "content": content,
        "type": file_type,
    }


def collect_files_context(
    vault_path: Path,
    rel_paths: list[str],
    scope: str = "documents",
) -> dict[str, Any]:
    """Collect an explicit list of files (open documents) as AI context.

    Unlike :func:`collect_directory_context`, this reads only the requested
    relative paths (markdown or PDF), never the whole directory. Paths outside
    the vault are silently ignored.

    Args:
        vault_path: Absolute path to the vault root.
        rel_paths: Relative file paths within the vault (in display order).
        scope: Context label exposed to the UI ("documents").

    Returns:
        Same shape as :func:`collect_directory_context`.
    """
    vault_resolved = vault_path.resolve()
    collected: list[dict[str, Any]] = []
    total_chars = 0
    seen: set[str] = set()

    for rel in rel_paths:
        if len(collected) >= BOOKSLM_MAX_FILES or total_chars >= BOOKSLM_MAX_TOTAL_CHARS:
            break
        if not rel or rel in seen:
            continue
        seen.add(rel)
        try:
            target = (vault_resolved / rel).resolve()
            target.relative_to(vault_resolved)
        except (ValueError, OSError):
            continue
        if not target.is_file():
            continue
        entry = _file_entry(target, rel, BOOKSLM_MAX_TOTAL_CHARS - total_chars)
        if entry is None:
            continue
        collected.append(entry)
        total_chars += len(entry["content"])

    return {
        "files": collected,
        "total_chars": total_chars,
        "file_count": len(collected),
        "directory_tree": "",
        "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS,
        "max_files": BOOKSLM_MAX_FILES,
        "scope": scope,
    }


def collect_adhoc_context(
    vault_path: Path,
    files: list[str] | None = None,
    directories: list[str] | None = None,
    scope: str = "general",
) -> dict[str, Any]:
    """Collect an ad-hoc mix of explicit files and directories.

    Backs the assistant ``@`` command: the user can attach extra files and
    directories to the current context without changing the base mode. Paths
    outside the vault are ignored by the underlying collectors.

    Args:
        vault_path: Absolute path to the vault root.
        files: Relative file paths to include.
        directories: Relative directory paths to include (recursive).
        scope: Context label exposed to the UI.

    Returns:
        Same shape as :func:`collect_directory_context`.
    """
    vault_resolved = vault_path.resolve()
    collected: list[dict[str, Any]] = []
    seen: set[str] = set()
    total_chars = 0

    def _add(entries: list[dict[str, Any]]) -> None:
        nonlocal total_chars
        for entry in entries:
            if len(collected) >= BOOKSLM_MAX_FILES or total_chars >= BOOKSLM_MAX_TOTAL_CHARS:
                return
            path = entry.get("path")
            if not path or path in seen:
                continue
            seen.add(path)
            collected.append(entry)
            total_chars += len(entry.get("content", ""))

    if files:
        _add(collect_files_context(vault_resolved, files, scope=scope)["files"])
    for directory in directories or []:
        _add(collect_directory_context(vault_resolved, directory)["files"])

    return {
        "files": collected,
        "total_chars": total_chars,
        "file_count": len(collected),
        "directory_tree": "",
        "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS,
        "max_files": BOOKSLM_MAX_FILES,
        "scope": scope,
    }


def merge_contexts(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:
    """Merge two context payloads, de-duplicating files by path."""
    files: list[dict[str, Any]] = []
    seen: set[str] = set()
    total_chars = 0
    for entry in list(base.get("files", [])) + list(extra.get("files", [])):
        path = entry.get("path")
        if not path or path in seen:
            continue
        if len(files) >= BOOKSLM_MAX_FILES or total_chars >= BOOKSLM_MAX_TOTAL_CHARS:
            break
        seen.add(path)
        files.append(entry)
        total_chars += len(entry.get("content", ""))

    tree = base.get("directory_tree", "") or extra.get("directory_tree", "")
    scope = extra.get("scope") or base.get("scope", "general")
    return {
        "files": files,
        "total_chars": total_chars,
        "file_count": len(files),
        "directory_tree": tree,
        "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS,
        "max_files": BOOKSLM_MAX_FILES,
        "scope": scope,
    }


def is_image_path(path: str) -> bool:
    """True when the path has a supported image extension."""
    return Path(path or "").suffix.lower() in IMAGE_EXTENSIONS


def load_vault_image_data_url(vault_path: Path, rel_path: str) -> str | None:
    """Read a vault image and return it as a ``data:`` URL for vision models.

    Returns ``None`` when the path is outside the vault, missing, too large,
    or not a supported image.
    """
    if not rel_path or not is_image_path(rel_path):
        return None
    vault_resolved = vault_path.resolve()
    try:
        target = (vault_resolved / rel_path).resolve()
        target.relative_to(vault_resolved)
    except (ValueError, OSError):
        return None
    if not target.is_file():
        return None
    try:
        if target.stat().st_size > BOOKSLM_MAX_IMAGE_BYTES:
            logger.warning("Image too large to send: %s", rel_path)
            return None
        raw = target.read_bytes()
    except OSError as exc:
        logger.warning("Cannot read image %s: %s", rel_path, exc)
        return None
    mime = mimetypes.guess_type(str(target))[0] or "image/png"
    encoded = base64.b64encode(raw).decode("ascii")
    return f"data:{mime};base64,{encoded}"


def empty_context(scope: str = "general") -> dict[str, Any]:
    """Return an empty context payload (used by the General assistant)."""
    return {
        "files": [],
        "total_chars": 0,
        "file_count": 0,
        "directory_tree": "",
        "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS,
        "max_files": BOOKSLM_MAX_FILES,
        "scope": scope,
    }


def _build_directory_tree(target_dir: Path, vault_root: Path) -> str:
    """Build a text representation of the directory tree (dirs + .md files)."""
    lines: list[str] = []
    try:
        for p in sorted(target_dir.rglob("*")):
            parts = p.relative_to(target_dir).parts
            if any(_should_skip(part) for part in parts):
                continue
            if p.is_dir():
                depth = len(p.relative_to(target_dir).parts)
                lines.append(f"{'  ' * depth}{p.name}/")
            elif p.is_file() and p.suffix.lower() == ".md":
                depth = len(p.relative_to(target_dir).parts)
                lines.append(f"{'  ' * depth}{p.name}")
    except PermissionError:
        pass
    return "\n".join(lines)


def build_system_prompt(context: dict[str, Any], scope: str = "directory") -> str:
    """Build a system prompt for document-scoped AI chat.

    Args:
        context: Output of collect_directory_context()/collect_files_context().
        scope: "directory" (whole folder) or "documents" (open files).

    Returns:
        System prompt string with file contents.
    """
    files = context.get("files", [])
    file_count = context.get("file_count", 0)
    total_chars = context.get("total_chars", 0)

    # Rough token estimate (1 token ≈ 4 chars)
    est_tokens = total_chars // 4
    token_warning = ""
    if est_tokens > 100_000:
        token_warning = f"\n⚠️ Attention : le contexte est très volumineux (~{est_tokens:,} tokens estimés). Les réponses peuvent être moins précises.\n"

    if scope == "documents":
        role = (
            "Tu es un assistant de recherche documentaire intégré à ObsiGate. "
            "Tu réponds UNIQUEMENT en te basant sur les documents ouverts par l'utilisateur ci-dessous. "
            "Cite tes sources avec le nom du fichier quand tu utilises une information. "
            "Si l'information ne se trouve pas dans les documents, dis-le clairement."
        )
        context_label = "📄 Documents ouverts"
    else:
        role = (
            "Tu es un assistant de recherche documentaire. "
            "Tu réponds UNIQUEMENT en te basant sur les documents fournis ci-dessous. "
            "Cite tes sources avec le nom du fichier quand tu utilises une information. "
            "Si l'information ne se trouve pas dans les documents, dis-le clairement."
        )
        context_label = "📚 Contexte"

    prompt = (
        f"{role}"
        f"\n\n{context_label} : {file_count} fichier(s) ({total_chars:,} caractères)"
        f"{token_warning}\n"
    )

    # Directory tree
    tree = context.get("directory_tree", "")
    if tree:
        prompt += f"\n📂 Arborescence du dossier :\n```\n{tree}\n```\n"

    # File contents
    prompt += "\n---\n"
    for f in files:
        prompt += f"\n## 📄 {f['title']} (`{f['path']}`)\n\n{f['content']}\n\n---\n"

    prompt += "\nFin du contexte. Réponds à la question de l'utilisateur en te basant uniquement sur ces documents."

    return prompt


GENERAL_SYSTEM_PROMPT = """Tu es l'assistant intégré d'ObsiGate, une application web auto-hébergée pour consulter, rechercher et éditer des vaults Obsidian (Markdown).

Tes deux rôles :
1. **Aider sur l'application** : expliquer la navigation, la recherche (full-text, filtres `tag:`, `created:`, `path:`), l'éditeur (CodeMirror, autosave, raccourcis), les onglets et le split view, les sauvegardes et la restauration, le partage public, l'export (HTML/Markdown/ePub/PDF), Mermaid, Excalidraw, les plugins, les thèmes, le mode hors-ligne, le MFA, etc.
2. **Proposer des actions concrètes** : créer un fichier ou un dossier dans un vault.

Quand l'utilisateur demande explicitement de créer un fichier, inclus EXACTEMENT un bloc de ce type dans ta réponse (et rien d'autre à l'intérieur du bloc) :

```obsigate-action
{"action": "create_file", "vault": "<nom du vault>", "path": "<chemin/relatif.md>", "content": "<contenu markdown>"}
```

Pour créer un dossier :

```obsigate-action
{"action": "create_directory", "vault": "<nom du vault>", "path": "<chemin/relatif>"}
```

Règles :
- Ne propose une action que si l'utilisateur la demande explicitement.
- Explique en une phrase ce que fait l'action avant le bloc.
- Utilise un chemin relatif se terminant par `.md` pour un fichier.
- N'invente jamais un nom de vault : utilise l'un des vaults disponibles listés ci-dessous.
- Réponds dans la langue de l'utilisateur, de façon concise et structurée (Markdown).
"""


def _format_app_context(app_context: dict[str, Any] | None, recent_files: list[dict[str, Any]] | None) -> str:
    """Render the live application state for the General assistant prompt.

    The General assistant has no document context; without this block it only
    knows the app exists. Passing what the user currently sees (open documents,
    current directory, active search, recently modified files) lets it answer
    "résume ce que je fais / où j'en suis" style questions.
    """
    app_context = app_context or {}
    lines: list[str] = []

    vault = app_context.get("vault") or app_context.get("current_vault")
    if vault:
        lines.append(f"- Vault sélectionné : {vault}")
    directory = app_context.get("directory")
    if directory:
        lines.append(f"- Répertoire courant : {directory}")
    current_path = app_context.get("current_path")
    if current_path:
        lines.append(f"- Document affiché dans le viewer : {current_path}")

    docs = app_context.get("open_documents") or []
    rendered_docs = []
    for doc in docs:
        if not isinstance(doc, dict) or not doc.get("path"):
            continue
        rendered_docs.append(f"{doc['path']} (vault {doc['vault']})" if doc.get("vault") else str(doc["path"]))
    if rendered_docs:
        lines.append("- Documents ouverts dans les onglets/panneaux : " + ", ".join(rendered_docs))

    query = app_context.get("search_query")
    if query:
        total = app_context.get("search_total")
        suffix = f" ({total} résultat(s))" if isinstance(total, int) else ""
        lines.append(f"- Recherche en cours : « {query} »{suffix}")

    results = app_context.get("search_results") or []
    rendered_results = [str(r.get("path")) for r in results[:10] if isinstance(r, dict) and r.get("path")]
    if rendered_results:
        lines.append("- Premiers résultats affichés : " + ", ".join(rendered_results))

    if recent_files:
        rendered_recent = [
            f"{f.get('vault')}/{f.get('path')}" for f in recent_files[:10] if f.get("path")
        ]
        if rendered_recent:
            lines.append("- Fichiers récemment modifiés : " + ", ".join(rendered_recent))

    if not lines:
        return ""
    return (
        "## Contexte applicatif actuel\n"
        "Voici ce que l'utilisateur voit ou fait en ce moment dans ObsiGate. "
        "Sers-t'en pour comprendre sa demande ; ne le répète pas inutilement.\n"
        + "\n".join(lines)
        + "\n"
    )


def build_general_system_prompt(
    vaults: list[str] | None = None,
    app_context: dict[str, Any] | None = None,
    recent_files: list[dict[str, Any]] | None = None,
) -> str:
    """System prompt for the General assistant (app help + actions).

    ``app_context`` carries the live UI state (open documents, current
    directory, active search) and ``recent_files`` the last modified files, so
    the assistant knows what the user is doing rather than answering blind.
    """
    prompt = GENERAL_SYSTEM_PROMPT
    if vaults:
        prompt += "\nVaults disponibles : " + ", ".join(sorted(vaults)) + "\n"
    else:
        prompt += "\nAucun vault n'est actuellement configuré.\n"
    block = _format_app_context(app_context, recent_files)
    if block:
        prompt += "\n" + block
    return prompt


def invalidate_cache(vault_path: Path | None = None, directory: str | None = None) -> int:
    """Invalidate cache entries.

    Args:
        vault_path: If provided, only invalidate entries for this vault.
        directory: If provided, only invalidate entries for this directory.

    Returns:
        Number of cache entries removed.
    """
    if vault_path is None and directory is None:
        count = len(_cache)
        _cache.clear()
        return count

    to_remove = list(_cache)
    for k in to_remove:
        del _cache[k]
    return len(to_remove)