document_notes.py

"""Bulk LLM documentation of repo files into path notes.

Uses Settings LLM (OpenAI-compatible / Ollama). Optional Granite Guardian
second-pass loads rubrics from GRANITE_GUIDELINES_DIR (outside this repo).
"""

from __future__ import annotations

import json
import os
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

import backends_config
import llm_chat as llm_chat_util
import notes_store
import workspace as workspace_mod

TOOL_DIR = Path(__file__).resolve().parent


def _data_path() -> Path:
    return workspace_mod.get().data_path


def _notes_path() -> Path:
    return workspace_mod.get().notes_path


def _notes_db_path() -> Path:
    return workspace_mod.get().notes_db_path

MAX_FILE_CHARS = 100_000
# Hard ceiling for reading from disk (source files; fixtures are filtered out earlier).
MAX_FILE_BYTES = 2_500_000
BINARY_EXTS = frozenset(
    {
        # images / fonts / icons
        "png", "jpg", "jpeg", "gif", "webp", "bmp", "ico", "icns", "svg", "avif",
        "woff", "woff2", "ttf", "otf", "eot", "ttc",
        # archives / binaries
        "pdf", "zip", "gz", "tar", "7z", "rar", "exe", "dll", "so", "dylib",
        "wasm", "jar", "war", "class", "o", "a", "lib", "bin", "dat", "pdb",
        "pyc", "pyo",
        # audio / video / capture
        "mp3", "mp4", "webm", "ogg", "wav", "aif", "aiff", "au", "mid", "midi",
        "avi", "mov", "mkv", "flac", "m4a", "h264", "h265", "pcap",
        # DB / opaque blobs / maps
        "db", "sqlite", "sqlite3", "fstmap",
    }
)
# Path fragments that are never useful as LLM documentation targets.
SKIP_PATH_PARTS = (
    "/.code-review-graph/",
    "/parsable_mime/",
    "/gradle/wrapper/",
    "/blink_perf_tests/",
    "/node_modules/",
    "/.git/",
    "/dejavu-fonts-",
)
SKIP_PATH_PREFIXES = (
    ".code-review-graph/",
)
SKIP_PATH_SUFFIXES = (
    "/hsts_preload.fstmap",
    "/perf-rainbow.html",
    "/perf-rainbow-hard.html",
    "/gradle-wrapper.jar",
)
SKIP_BASENAMES = frozenset(
    {
        "gradle-wrapper.jar",
        "hsts_preload.fstmap",
    }
)

_job_lock = threading.Lock()
_job: dict[str, Any] = {
    "running": False,
    "op": "document",
    "started_at": None,
    "finished_at": None,
    "ok": None,
    "error": None,
    "message": None,
    "total": 0,
    "done": 0,
    "written": 0,
    "skipped": 0,
    "failed": 0,
    "current_path": None,
    "log": [],
}


def _utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def job_snapshot() -> dict[str, Any]:
    with _job_lock:
        return dict(_job)


def _set_job(**kwargs: Any) -> None:
    with _job_lock:
        _job.update(kwargs)


def _log(msg: str) -> None:
    with _job_lock:
        log = list(_job.get("log") or [])
        log.append({"ts": _utc_now(), "message": msg})
        _job["log"] = log[-80:]
        _job["message"] = msg


def guidelines_dir() -> Path | None:
    cfg = backends_config.load_backends().get("document") or {}
    raw = (
        os.environ.get("GRANITE_GUIDELINES_DIR")
        or cfg.get("guidelines_dir")
        or ""
    ).strip()
    if not raw:
        return None
    p = Path(raw).expanduser()
    return p if p.is_dir() else None


def load_rubric() -> str:
    gdir = guidelines_dir()
    if not gdir:
        return ""
    parts: list[str] = []
    for name in ("document_note_rubric.md", "safety_policy.md"):
        f = gdir / name
        if f.is_file():
            try:
                parts.append(f.read_text(encoding="utf-8", errors="replace")[:12_000])
            except OSError:
                pass
    return "\n\n---\n\n".join(parts)


def _walk_tree(node: dict[str, Any], prefix: str, out: list[str]) -> None:
    kind = node.get("kind")
    name = str(node.get("name") or "")
    if kind == "file":
        rel = f"{prefix}/{name}".replace("\\", "/").strip("/") if prefix else name
        if rel:
            out.append(rel)
        return
    if kind != "dir":
        return
    base = prefix
    for ch in node.get("children") or []:
        cname = str(ch.get("name") or "")
        child_base = f"{base}/{cname}".replace("\\", "/").strip("/") if base else cname
        if ch.get("kind") == "dir":
            _walk_tree(ch, child_base, out)
        elif ch.get("kind") == "file":
            out.append(child_base)


def _path_ext(rel: str) -> str:
    name = Path(rel.replace("\\", "/")).name
    if "." not in name:
        return ""
    return name.rsplit(".", 1)[-1].lower()


def should_document_path(rel: str) -> bool:
    """False for binaries, media, DBs, and known non-source fixtures."""
    p = (rel or "").replace("\\", "/").lstrip("/")
    if not p:
        return False
    low = f"/{p.lower()}"
    for part in SKIP_PATH_PARTS:
        if part in low:
            return False
    for pref in SKIP_PATH_PREFIXES:
        if p.lower().startswith(pref):
            return False
    for suf in SKIP_PATH_SUFFIXES:
        if low.endswith(suf):
            return False
    name = Path(p).name.lower()
    if name in SKIP_BASENAMES:
        return False
    ext = _path_ext(p)
    if ext in BINARY_EXTS:
        return False
    # Extension-less blobs under test resource trees (e.g. gif/test87a, binary_file).
    if not ext and ("/tests/" in low or "/resources/" in low or "/support/" in low):
        return False
    return True


def list_file_paths_from_data(directory: str | None = None) -> list[str]:
    """List text file paths from data.json (full tree; optional directory prefix)."""
    if not _data_path().is_file():
        return []
    try:
        data = json.loads(_data_path().read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return []
    tree = data.get("tree")
    if not isinstance(tree, dict):
        return []
    raw: list[str] = []
    # Root is a directory; walk its children with empty prefix for top-level names
    if tree.get("kind") == "dir":
        for ch in tree.get("children") or []:
            cname = str(ch.get("name") or "")
            if ch.get("kind") == "dir":
                _walk_tree(ch, cname, raw)
            elif ch.get("kind") == "file":
                raw.append(cname)
    else:
        _walk_tree(tree, "", raw)

    prefix = (directory or "").replace("\\", "/").strip("/")
    seen: set[str] = set()
    files: list[str] = []
    for p in raw:
        p = (p or "").replace("\\", "/").lstrip("/")
        if not p or p in seen:
            continue
        if prefix and not (p == prefix or p.startswith(prefix + "/")):
            continue
        if not should_document_path(p):
            continue
        seen.add(p)
        files.append(p)
    return files


def filter_paths_by_directory(paths: list[str], directory: str | None) -> list[str]:
    """Keep paths under a relative directory prefix (or exact file match)."""
    prefix = (directory or "").replace("\\", "/").strip("/")
    if not prefix:
        return list(paths)
    out: list[str] = []
    for p in paths:
        rel = (p or "").replace("\\", "/").lstrip("/")
        if rel == prefix or rel.startswith(prefix + "/"):
            out.append(rel)
    return out


def _note_is_locked(note: dict[str, Any] | None) -> bool:
    meta = (note or {}).get("meta") if isinstance(note, dict) else None
    return bool(isinstance(meta, dict) and meta.get("locked"))


def _doc_system_prompt(*, update: bool) -> str:
    structure = (
        "Use these sections when relevant:\n"
        "- Purpose\n"
        "- Layout / structure\n"
        "- Key APIs / exports\n"
        "- Usage\n"
        "- Related paths\n"
    )
    if update:
        return (
            "You update an existing documentation note so it matches the current source file. "
            "Preserve useful structure and accurate claims; correct outdated details; add new "
            "exports/APIs; remove stale ones. Write concise Markdown. Do not invent APIs not in the file.\n"
            + structure
        )
    return (
        "You document source files for a codebase graph tool. "
        "Write concise Markdown. Do not invent APIs not in the file.\n"
        + structure
    )


def _read_file(repo: Path, rel: str) -> str | None:
    if not should_document_path(rel):
        return None
    cand = (repo / rel).resolve()
    try:
        cand.relative_to(repo.resolve())
    except ValueError:
        return None
    if not cand.is_file():
        return None
    try:
        size = cand.stat().st_size
        if size <= 0 or size > MAX_FILE_BYTES:
            return None
        text = cand.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return None
    if not text.strip():
        return None
    if "\0" in text[:4096]:
        return None
    if len(text) > MAX_FILE_CHARS:
        text = text[:MAX_FILE_CHARS] + "\n\n… [truncated]"
    return text


def _effective_doc_reasoning_effort(cfg: dict[str, Any], model: str) -> str:
    """Document job effort only — never chat Settings. Skip if generator is not a reasoning model."""
    if not llm_chat_util.model_supports_reasoning_effort(model):
        return ""
    raw = str(cfg.get("reasoning_effort") or "").strip().lower()
    if raw in ("inherit", "same"):
        return ""
    return llm_chat_util.normalize_reasoning_effort(raw)


def _without_reasoning(settings: dict[str, Any]) -> dict[str, Any]:
    out = dict(settings)
    out.pop("reasoningEffort", None)
    out.pop("reasoning_effort", None)
    return out


def _parse_note_response(content: str, path: str) -> tuple[str, str]:
    content = llm_chat_util.strip_reasoning_markup(content)
    title = Path(path).name
    body = content
    if content.startswith("#"):
        first, _, rest = content.partition("\n")
        title = first.lstrip("#").strip() or title
        body = rest.strip() or content
    return title[:200], body


def _guardian_pass(llm_chat: Callable[[dict], dict], draft_title: str, draft_body: str, path: str) -> tuple[bool, str]:
    cfg = backends_config.load_backends().get("document") or {}
    if not cfg.get("guardian_enabled"):
        return True, "guardian-disabled"
    rubric = load_rubric()
    if not rubric:
        return True, "no-rubric"
    guardian_model = (cfg.get("guardian_model") or cfg.get("generator_model") or "").strip()
    # Use same provider settings from llm payload passed by caller via closure
    messages = [
        {
            "role": "system",
            "content": (
                "You are Granite Guardian scoring a documentation note. "
                "Reply with exactly one line: PASS or FAIL, then a short reason.\n\n"
                f"Scoring guidelines:\n{rubric[:8000]}"
            ),
        },
        {
            "role": "user",
            "content": f"Path: {path}\nTitle: {draft_title}\n\nBody:\n{draft_body[:6000]}",
        },
    ]
    # Caller injects provider/baseUrl/apiKey via llm_chat wrapper
    try:
        res = llm_chat({"messages": messages, "model": guardian_model} if guardian_model else {"messages": messages})
        text = (res.get("content") or "").strip()
    except Exception as exc:  # noqa: BLE001
        return True, f"guardian-error:{exc}"
    upper = text.upper()
    if upper.startswith("FAIL") or "\nFAIL" in upper[:40]:
        return False, text[:240]
    return True, text[:240]


def _run_document_job(
    *,
    repo: Path,
    mode: str,
    paths: list[str] | None,
    directory: str | None,
    llm_settings: dict[str, Any],
    llm_chat: Callable[[dict], dict],
) -> None:
    try:
        existing = notes_store.load_notes_file(_notes_path())
        is_update = mode == "update"
        if mode == "selected" and paths:
            targets = [p.replace("\\", "/").lstrip("/") for p in paths if p]
            targets = filter_paths_by_directory(targets, directory)
        else:
            # Prefer directory-scoped listing so prefix jobs don't depend on
            # whole-repo walk order (serial LLM — no batch size window).
            targets = list_file_paths_from_data(directory=directory)
            if paths:
                allow = {p.replace("\\", "/").lstrip("/") for p in paths}
                targets = [p for p in targets if p in allow]

        if mode == "missing":
            targets = [
                p
                for p in targets
                if p not in existing or not (existing.get(p) or {}).get("body")
            ]
        elif is_update:
            targets = [
                p
                for p in targets
                if (existing.get(p) or {}).get("body") and not _note_is_locked(existing.get(p))
            ]
        else:
            # all / selected — never overwrite locked documentation notes
            targets = [p for p in targets if not _note_is_locked(existing.get(p))]

        targets = [p for p in targets if should_document_path(p)]

        verb = "Updating" if is_update else "Documenting"
        if not targets:
            scope = (directory or "").replace("\\", "/").strip("/") or "workspace"
            why = (
                f"No files matched scope “{scope}”"
                + (
                    " (check the folder path, or re-scan if the graph is stale)."
                    if directory
                    else " in the scanned graph."
                )
            )
            _set_job(
                running=False,
                total=0,
                done=0,
                written=0,
                skipped=0,
                failed=0,
                message=why,
                ok=True,
                error=None,
                finished_at=_utc_now(),
                current_path=None,
                op="update" if is_update else "document",
            )
            _log(why)
            return

        _set_job(
            total=len(targets),
            done=0,
            written=0,
            skipped=0,
            failed=0,
            message=f"{verb} {len(targets)} files…",
            op="update" if is_update else "document",
        )
        cfg = backends_config.load_backends().get("document") or {}
        gen_model = (cfg.get("generator_model") or llm_settings.get("model") or "").strip()
        gen_effort = _effective_doc_reasoning_effort(cfg, gen_model)
        notes = dict(existing)
        system = _doc_system_prompt(update=is_update)

        for rel in targets:
            _set_job(current_path=rel)
            _log(f"{verb} {rel}")
            if _note_is_locked(notes.get(rel)):
                _set_job(skipped=_job["skipped"] + 1, done=_job["done"] + 1)
                _log(f"Skip (locked): {rel}")
                continue
            text = _read_file(repo, rel)
            if text is None:
                _set_job(skipped=_job["skipped"] + 1, done=_job["done"] + 1)
                _log(f"Skip (unreadable): {rel}")
                continue
            prev = notes.get(rel) or {}
            if is_update:
                user = (
                    f"File path: {rel}\n\n"
                    f"## Existing note\n"
                    f"Title: {prev.get('title') or Path(rel).name}\n\n"
                    f"{(prev.get('body') or '')[:8000]}\n\n"
                    f"## Current source\n```\n{text}\n```"
                )
            else:
                user = f"File path: {rel}\n\n```\n{text}\n```"
            try:
                payload = _without_reasoning(llm_settings)
                payload["model"] = gen_model or llm_settings.get("model")
                payload["messages"] = [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ]
                if gen_effort:
                    payload["reasoningEffort"] = gen_effort
                res = llm_chat(payload)
                title, body = _parse_note_response(res.get("content") or "", rel)

                def guardian_chat(extra: dict) -> dict:
                    merged = {**_without_reasoning(llm_settings), **extra}
                    merged.pop("reasoningEffort", None)
                    merged.pop("reasoning_effort", None)
                    if cfg.get("guardian_model"):
                        merged["model"] = cfg["guardian_model"]
                    if cfg.get("guardian_url"):
                        # allow separate base for guardian
                        gurl = str(cfg["guardian_url"]).rstrip("/")
                        if gurl:
                            merged["baseUrl"] = gurl
                            if "ollama" in gurl or cfg.get("guardian_provider") == "ollama":
                                merged["provider"] = "ollama"
                    return llm_chat(merged)

                ok, gmsg = _guardian_pass(guardian_chat, title, body, rel)
                if not ok:
                    _set_job(skipped=_job["skipped"] + 1, done=_job["done"] + 1)
                    _log(f"Guardian reject {rel}: {gmsg}")
                    continue

                prev_meta = prev.get("meta") if isinstance(prev.get("meta"), dict) else {}
                notes[rel] = {
                    "title": title,
                    "body": body,
                    "meta": {
                        **prev_meta,
                        "source": "llm-document",
                        "kind": "documentation",
                        "locked": True,
                        "model": gen_model or llm_settings.get("model"),
                        "guardian": "pass" if cfg.get("guardian_enabled") else "disabled",
                        "guardian_detail": gmsg if cfg.get("guardian_enabled") else None,
                        "at": _utc_now(),
                        "updated": bool(is_update),
                    },
                }
                notes_store.save_notes_with_history(_notes_path(), notes, db_path=_notes_db_path())
                _set_job(written=_job["written"] + 1, done=_job["done"] + 1)
            except Exception as exc:  # noqa: BLE001
                _set_job(failed=_job["failed"] + 1, done=_job["done"] + 1)
                _log(f"Fail {rel}: {exc}")

        _set_job(
            running=False,
            ok=True,
            finished_at=_utc_now(),
            current_path=None,
            message=f"Done · wrote {_job['written']} · skipped {_job['skipped']} · failed {_job['failed']}",
        )
    except Exception as exc:  # noqa: BLE001
        _set_job(
            running=False,
            ok=False,
            error=str(exc),
            finished_at=_utc_now(),
            message=str(exc),
        )


def start_document_job(
    repo: Path,
    *,
    mode: str = "missing",
    paths: list[str] | None = None,
    directory: str | None = None,
    llm_settings: dict[str, Any] | None = None,
    llm_chat: Callable[[dict], dict] | None = None,
) -> dict[str, Any]:
    if llm_chat is None:
        raise ValueError("llm_chat required")
    settings = llm_settings or {}
    mode_norm = mode if mode in ("missing", "all", "selected", "update") else "missing"
    with _job_lock:
        if _job.get("running"):
            return {"ok": False, "error": "document job already running", "job": dict(_job)}
        _job.update(
            {
                "running": True,
                "op": "update" if mode_norm == "update" else "document",
                "started_at": _utc_now(),
                "finished_at": None,
                "ok": None,
                "error": None,
                "message": "Starting…",
                "total": 0,
                "done": 0,
                "written": 0,
                "skipped": 0,
                "failed": 0,
                "current_path": None,
                "log": [],
            }
        )

    thread = threading.Thread(
        target=_run_document_job,
        kwargs={
            "repo": repo,
            "mode": mode_norm,
            "paths": paths,
            "directory": directory,
            "llm_settings": settings,
            "llm_chat": llm_chat,
        },
        daemon=True,
    )
    thread.start()
    return {"ok": True, "started": True, "job": job_snapshot()}