Bolt-Pet/bolt_pet/screen_text.py

"""Reading the text that's actually on a monitor, via screenshot + OCR.

This is the "Bolt can see what's on screen" half of the screen features. It is
**pull, not push**: nothing here runs on its own. The server has to ask, by
relaying `petctl read`, and the recognised text goes back as that command's
output through the existing tool-result relay (see server_client.converse).
That's deliberate on two counts — OCR of a 4K screen costs a second or two,
which would be tacked onto every single utterance if it ran automatically, and
"screen contents leave this machine" should be a thing Bolt decides to do and
you can see in the log, not a silent constant.

Both halves are optional and soft-fail with a reason, the way hotkey.py does:
capture needs `mss`, recognition needs a Tesseract or RapidOCR install. With
neither, `petctl read` reports what's missing instead of raising, and the rest
of the pet carries on.

The pure parts (cleaning OCR output, formatting the reply, deciding which
engine to use given what's installed) are split out and unit tested; only
capture and the OCR call itself need a real screen.
"""

from __future__ import annotations

import re
import shutil
from typing import Callable, Optional

from .monitors import Monitor

DEFAULT_MAX_CHARS = 4000

INSTALL_HINT = (
    "install one of: `pip install mss pytesseract` + `sudo apt install "
    "tesseract-ocr` (fastest), or `pip install mss rapidocr-onnxruntime` "
    "(no system package needed)"
)

# Lines that are almost certainly OCR noise rather than text: window chrome
# fragments, isolated punctuation, single stray characters.
_MIN_MEANINGFUL = 2


def _module_available(name: str) -> bool:
    import importlib.util

    try:
        return importlib.util.find_spec(name) is not None
    except (ImportError, ValueError):
        return False


# ── pure helpers (unit tested; no screen, no OCR engine needed) ──────────────

def resolve_engine(
    has_module: Callable[[str], bool] = _module_available,
    which: Callable[[str], Optional[str]] = shutil.which,
) -> tuple[Optional[str], str]:
    """Pick an OCR engine from what's installed.

    Returns `(engine, reason)`. *engine* is None when nothing usable is
    present, and *reason* then explains what to install. Probes are injected
    so this is testable on a machine with a different set of things installed.
    """
    if has_module("pytesseract") and which("tesseract"):
        return "pytesseract", ""
    if has_module("rapidocr_onnxruntime"):
        return "rapidocr", ""
    if has_module("pytesseract") and not which("tesseract"):
        return None, (
            "pytesseract is installed but the tesseract binary isn't on PATH "
            "(try: sudo apt install tesseract-ocr)"
        )
    return None, f"no OCR engine available — {INSTALL_HINT}"


def capture_available(has_module: Callable[[str], bool] = _module_available) -> bool:
    return has_module("mss")


def clean_ocr_text(raw: str, max_chars: int = DEFAULT_MAX_CHARS) -> str:
    """Squeeze raw OCR output into something worth sending.

    Screen OCR produces a lot of junk — single stray glyphs off window
    borders, runs of blank lines, the same toolbar label recognised twice. All
    of that costs tokens and tells the model nothing, so it goes.
    """
    if not raw:
        return ""
    lines: list[str] = []
    for line in raw.splitlines():
        line = re.sub(r"[^\S\n]+", " ", line).strip()
        if not line:
            continue
        if len(re.sub(r"[^0-9A-Za-z]", "", line)) < _MIN_MEANINGFUL:
            continue
        if lines and line == lines[-1]:
            continue  # consecutive duplicate
        lines.append(line)
    text = "\n".join(lines)
    if max_chars and len(text) > max_chars:
        text = text[: max_chars - 1].rstrip() + "…"
        text += "\n[truncated]"
    return text


def format_reading(monitor: Optional[Monitor], text: str) -> str:
    """The tool output handed back for `petctl read`."""
    where = f"monitor {monitor.number} ({monitor.name})" if monitor else "screen"
    if not text.strip():
        return f"[pet] read {where}: no text recognised"
    return f"[pet] text on {where}:\n{text}"


# ── capture + recognition (needs a real screen) ──────────────────────────────

def capture(monitor: Monitor):
    """Grab *monitor* as a PIL image, or None if capture isn't available."""
    try:
        import mss
        from PIL import Image
    except ImportError:
        return None
    try:
        box = {
            "left": monitor.x,
            "top": monitor.y,
            "width": monitor.width,
            "height": monitor.height,
        }
        with mss.mss() as sct:
            shot = sct.grab(box)
        return Image.frombytes("RGB", shot.size, shot.bgra, "raw", "BGRX")
    except Exception:
        return None


def _ocr(image, engine: str) -> str:
    if engine == "pytesseract":
        import pytesseract

        # Grayscale first: tesseract is measurably better on it than on the
        # colour desktop, and it's a cheap conversion.
        return pytesseract.image_to_string(image.convert("L"))
    if engine == "rapidocr":
        import numpy as np
        from rapidocr_onnxruntime import RapidOCR

        result, _ = RapidOCR()(np.array(image))
        if not result:
            return ""
        return "\n".join(line[1] for line in result)
    return ""


def read_monitor(monitor: Monitor, max_chars: int = DEFAULT_MAX_CHARS) -> str:
    """OCR one screen and return the formatted tool output.

    Never raises: every failure path returns a sentence explaining itself,
    because the return value goes straight back to the server as the result of
    a command Bolt chose to run.
    """
    if not capture_available():
        return f"[pet] can't capture the screen — {INSTALL_HINT}"
    engine, reason = resolve_engine()
    if engine is None:
        return f"[pet] can't read the screen — {reason}"
    image = capture(monitor)
    if image is None:
        return (
            f"[pet] couldn't capture monitor {monitor.number} "
            "(is this a Wayland session? mss needs X11)"
        )
    try:
        raw = _ocr(image, engine)
    except Exception as exc:
        return f"[pet] OCR failed on monitor {monitor.number}: {exc}"
    return format_reading(monitor, clean_ocr_text(raw, max_chars))


def read_monitors(monitors: list[Monitor], max_chars: int = DEFAULT_MAX_CHARS) -> str:
    """OCR several screens, splitting the character budget between them."""
    if not monitors:
        return "[pet] no monitor information available"
    if len(monitors) == 1:
        return read_monitor(monitors[0], max_chars)
    share = max(400, max_chars // len(monitors))
    return "\n\n".join(read_monitor(m, share) for m in monitors)