chemenu/tools/chemenu/commands/search.py

"""`wikitool search` - find pages without reading `kb/index.md`.

This command exists to make retrieval cheap. Before it, the documented way to
find a page was to read the whole generated index; at a few hundred pages that
is tens of thousands of tokens spent to learn three filenames. A search returns
the same pointers for a fraction of it.

Two halves, deliberately kept separate:

- Text search is answered by a pluggable backend (`rg` today) - see
  `chemenu/search/`.
- Frontmatter predicates (`--field`) are evaluated here, in-process, on the
  structured YAML rather than on its rendering. With no text at all this is a
  pure structured query, which is how "systems below 0.6 confidence, oldest
  first" is asked without a second command.

Scope is `kb/` only. `instructions/` is discovered through
`wikitool instructions list`, because a procedure is found by what it is *for*
(its description), not by keywords in its body.
"""
from __future__ import annotations

import json
from pathlib import Path

import typer

from chemenu import config
from chemenu.commands._util import fail, today_iso
from chemenu.frontmatter_io import read_page
from chemenu.kb_scan import iter_kb_pages
from chemenu.page import Page
from chemenu.search import filters
from chemenu.search.base import page_key
from chemenu.search.filters import PredicateError
from chemenu.search.fuse import reciprocal_rank_fusion
from chemenu.search.registry import UnknownBackend, resolve
from chemenu.search.ripgrep import RipgrepFailed, RipgrepMissing, build_hit
from chemenu.search.types import Predicate, SearchHit, SearchQuery

TITLE_WIDTH = 34
SUMMARY_WIDTH = 84


def load_pages_by_path(kb_dir: Path | None = None, root: Path | None = None) -> dict[str, Page]:
    """Every page under `kb/`, keyed by repo-relative path.

    Path-keyed rather than title-keyed on purpose: `load_kb_pages()` drops one
    of two pages sharing a stem, and search should still find both - a
    duplicate title is a lint finding, not a reason to hide a page.
    """
    kb_dir = kb_dir or config.KB_DIR
    root = root or config.ROOT
    pages: dict[str, Page] = {}
    for path in iter_kb_pages(kb_dir):
        frontmatter, body = read_page(path)
        pages[page_key(path, root)] = Page(path=path, frontmatter=frontmatter, body=body)
    return pages


def _sort_key(hit: SearchHit, field: str):
    value = hit.as_dict().get(field)
    if value is None:
        # Missing values sort last in either direction rather than crashing on
        # a None comparison.
        return (1, "")
    if isinstance(value, (int, float)):
        return (0, value)
    return (0, str(value).lower())


def sort_hits(hits: list[SearchHit], sort: str | None) -> list[SearchHit]:
    """Sort by a hit field. A leading `-` reverses, e.g. `--sort -confidence`."""
    if not sort:
        return hits
    descending = sort.startswith("-")
    field = sort.lstrip("-")
    ordered = sorted(hits, key=lambda h: _sort_key(h, field), reverse=descending)
    return ordered


def run_search(
    query: SearchQuery,
    pages: dict[str, Page],
    backends: list,
    kb_dir: Path | None = None,
) -> list[SearchHit]:
    """Answer a query. Pure: no I/O beyond whatever a backend does."""
    filters.validate_fields(query.predicates, pages)

    if query.text:
        rankings = [backend.search(query, pages) for backend in backends]
        hits = rankings[0] if len(rankings) == 1 else reciprocal_rank_fusion(rankings)
        allowed = filters.apply_predicates(pages, query.predicates, kb_dir)
        hits = [hit for hit in hits if hit.path in allowed]
    else:
        selected = filters.apply_predicates(pages, query.predicates, kb_dir)
        hits = [
            build_hit(page, key, [], query, backend="frontmatter", kb_dir=kb_dir)
            for key, page in selected.items()
        ]
        hits.sort(key=lambda h: h.title.lower())

    hits = sort_hits(hits, query.sort)
    return hits[: query.limit] if query.limit else hits


def _truncate(text: str, width: int) -> str:
    text = " ".join(text.split())
    return text if len(text) <= width else text[: width - 1] + "\u2026"


def render_table(hits: list[SearchHit], show_matches: bool) -> str:
    if not hits:
        return "No matches."
    lines = []
    for hit in hits:
        kind = hit.kind or "?"
        if hit.subtype:
            kind = f"{kind}/{hit.subtype}"
        lines.append(
            f"{hit.score:6.1f}  {_truncate(hit.title, TITLE_WIDTH):<{TITLE_WIDTH}}  "
            f"{kind:<18}  {_truncate(hit.summary, SUMMARY_WIDTH)}"
        )
        if show_matches:
            for match in hit.matches:
                lines.append(f"          {hit.path}:{match.line}: {_truncate(match.text, 100)}")
    lines.append("")
    lines.append(f"{len(hits)} result(s).")
    return "\n".join(lines)


def search_command(
    text: str = typer.Argument(
        None,
        help="Text to search for. Omit it to run a pure frontmatter query.",
    ),
    field: list[str] = typer.Option(
        None,
        "--field",
        "-f",
        help="Frontmatter predicate, repeatable (AND). Forms: field=value, "
        "field~substring, 'field>=value', 'field:*' (present), '!field' (absent).",
    ),
    kind: str = typer.Option(None, "--kind", help="Shorthand for --field kind=<value>."),
    subtype: str = typer.Option(None, "--subtype", help="Shorthand for --field subtype=<value>."),
    collection: str = typer.Option(
        None, "--collection", help="Shorthand for --field collection=<value>."
    ),
    tag: str = typer.Option(None, "--tag", help="Shorthand for --field tags=<value>."),
    regex: bool = typer.Option(
        False, "--regex", help="Treat the query as a regex. Off by default: terms are literal."
    ),
    limit: int = typer.Option(20, "--limit", help="Maximum number of results. 0 for no limit."),
    sort: str = typer.Option(
        None, "--sort", help="Sort by a result field; prefix with '-' to reverse, e.g. -confidence."
    ),
    backend: str = typer.Option(
        None,
        "--backend",
        help="Search backend(s), comma-separated. Default 'rg' (or $WIKITOOL_SEARCH_BACKEND).",
    ),
    show_matches: bool = typer.Option(
        False, "--matches", help="Print the matching lines under each result."
    ),
    json_out: bool = typer.Option(False, "--json", help="Print the results as JSON."),
):
    """Search kb/ by text, by frontmatter, or by both."""
    raw_predicates = list(field or [])
    for value, name in ((kind, "kind"), (subtype, "subtype"), (collection, "collection")):
        if value:
            raw_predicates.append(f"{name}={value}")
    if tag:
        raw_predicates.append(f"tags={tag}")

    if not text and not raw_predicates:
        fail("Nothing to search for: give a query, or at least one --field predicate.")

    try:
        predicates: tuple[Predicate, ...] = tuple(
            filters.parse_predicate(raw) for raw in raw_predicates
        )
    except PredicateError as exc:
        fail(str(exc))

    try:
        backends = resolve(backend)
    except UnknownBackend as exc:
        fail(str(exc))

    query = SearchQuery(
        text=text,
        predicates=predicates,
        regex=regex,
        limit=limit,
        sort=sort,
    )

    pages = load_pages_by_path()
    try:
        hits = run_search(query, pages, backends)
    except PredicateError as exc:
        fail(str(exc))
    except RipgrepMissing as exc:
        fail(str(exc))
    except RipgrepFailed as exc:
        fail(str(exc))

    if json_out:
        payload = {
            "generated": today_iso(),
            "query": text,
            "predicates": [p.render() for p in predicates],
            "backend": ",".join(b.name for b in backends),
            "count": len(hits),
            "results": [hit.as_dict() for hit in hits],
        }
        typer.echo(json.dumps(payload, indent=2))
        return

    typer.echo(render_table(hits, show_matches))