bdutils.py

"""Shared helpers for beads-utils scripts.

This module lives alongside the scripts in the repo root. Python places the
script's directory on sys.path[0] when executing, so `from bdutils import ...`
resolves the sibling module without any prelude.
"""
from __future__ import annotations

import contextlib
import functools
import json
import os
import re
import shlex
import shutil
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import NoReturn

__version__ = "0.4.0"

ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")


def add_version_arg(parser) -> None:
    """Add a standard '--version' flag that prints '<prog> <__version__>'.

    Keeps the version string identical across every beads-utils script so the
    whole collection versions as one unit (see __version__ above)."""
    parser.add_argument(
        "--version", action="version", version=f"%(prog)s {__version__}"
    )


def error(msg: str) -> NoReturn:
    """Exit with a lowercase 'error: ...' message on stderr."""
    sys.exit(f"error: {msg}")


def warn(msg: str) -> None:
    """Write a 'warning: ...' line to stderr without exiting."""
    sys.stderr.write(f"warning: {msg}\n")


def format_ts(iso: str) -> str:
    """Format an ISO-8601 (UTC, 'Z'-suffixed) timestamp as 'YYYY-MM-DD HH:MM' in
    local time. Empty input returns ''; unparseable input is returned as-is."""
    if not iso:
        return ""
    try:
        dt = datetime.fromisoformat(iso.replace("Z", "+00:00"))
    except ValueError:
        return iso
    return dt.astimezone().strftime("%Y-%m-%d %H:%M")


def format_priority(p) -> str:
    """Render a beads priority value as 'P0'..'P4'. None/'' becomes 'P?'."""
    if p is None or p == "":
        return "P?"
    try:
        return f"P{int(p)}"
    except (TypeError, ValueError):
        return str(p)


def resolve_project_path(arg: str) -> Path:
    """Resolve a user-supplied project path and verify it has a .beads/ dir.

    Exits with error() if the path is missing the .beads/ marker.
    """
    project_path = Path(arg).expanduser().resolve()
    if not (project_path / ".beads").is_dir():
        error(f"no .beads/ directory at {project_path}")
    return project_path


# --- Color output --------------------------------------------------------
#
# Plain SGR escapes; no dependency. Scripts that colorize take the flag via
# add_color_arg() and resolve it once with want_color(), so the spelling and
# the auto-detection rules stay identical across the collection — the same
# reason add_version_arg() exists.

COLOR_MODES = ("auto", "always", "never")

COLORS = {
    "red": "\033[31m",
    "green": "\033[32m",
    "yellow": "\033[33m",
    "blue": "\033[34m",
    "magenta": "\033[35m",
    "cyan": "\033[36m",
    "bold": "\033[1m",
    "dim": "\033[2m",
}
RESET = "\033[0m"


def add_color_arg(parser) -> None:
    """Add a standard '--color=auto|always|never' flag (default: auto)."""
    parser.add_argument(
        "--color", choices=COLOR_MODES, default="auto", metavar="WHEN",
        help="Colorize output: 'auto' (default: only when stdout is a tty and "
             "NO_COLOR is unset), 'always' (keeps color through a pipe), or "
             "'never'.",
    )


def want_color(mode: str = "auto") -> bool:
    """Whether to emit ANSI color, given a --color mode.

    'auto' requires an interactive stdout, an unset NO_COLOR (no-color.org:
    any non-empty value disables color), and a TERM that isn't 'dumb'.

    Callers that page must resolve this *before* entering paged_output():
    that context manager yields the pager's stdin pipe, so testing the
    yielded stream reports "not a tty" in exactly the case — a user at a
    terminal — where color is wanted.
    """
    if mode == "always":
        return True
    if mode == "never":
        return False
    if os.environ.get("NO_COLOR", ""):
        return False
    if os.environ.get("TERM", "") == "dumb":
        return False
    return sys.stdout.isatty()


def paint(text: str, color: str, enabled: bool = True) -> str:
    """Wrap `text` in `color` when `enabled`; return it untouched otherwise.

    `color` is one or more space-separated names from COLORS, so attributes
    compose: 'bold blue' emits weight and hue together.

    Two ways to "just raise the contrast" that don't work, both verified
    against a real Solarized Light terminal:
      - the bright slots (90-97): Solarized repurposes them as greys, so
        bright green is base01, not a brighter green;
      - 'bold': terminals with "draw bold text in bright colors" enabled
        route it to those same grey slots, so bolding a hue can *remove* it.
    Prefer picking a better-separated hue over reaching for either.

    An unknown name is skipped rather than raising: a miscolored line should
    never take down an otherwise working command.
    """
    if not enabled or not text:
        return text
    codes = [COLORS[name] for name in color.split() if name in COLORS]
    return f"{''.join(codes)}{text}{RESET}" if codes else text


# --- Dolt / git plumbing -------------------------------------------------
#
# Shared by the Dolt-aware scripts (bd-dolt-check, bd-dolt-diff). Every
# helper here is read-only and degrades to a None/empty result rather than
# raising, so callers can report "not verifiable" instead of crashing when
# the dolt CLI is absent or a ref doesn't resolve.


def strip_ansi(s: str) -> str:
    """Remove ANSI SGR escapes (the dolt CLI colorizes even when piped)."""
    return ANSI_RE.sub("", s)


def run(cmd: list[str], cwd: Path, check: bool = False) -> subprocess.CompletedProcess[str]:
    """Capture-output subprocess run rooted at `cwd` (never os.chdir)."""
    return subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, check=check)


@functools.cache
def have_dolt() -> bool:
    """True if the `dolt` CLI is on PATH."""
    return shutil.which("dolt") is not None


def read_metadata(beads_dir: Path) -> tuple[str, str]:
    """Return (dolt_database, dolt_mode) from .beads/metadata.json."""
    path = beads_dir / "metadata.json"
    if not path.is_file():
        error(f"no metadata.json in {beads_dir}")
    try:
        data = json.loads(path.read_text())
    except (OSError, json.JSONDecodeError) as exc:
        error(f"could not read {path}: {exc}")
    return data.get("dolt_database", "unknown"), data.get("dolt_mode", "unknown")


def locate_dolt_db(beads_dir: Path, db_name: str) -> Path | None:
    """Find the Dolt database dir, covering both embedded and server layouts."""
    for candidate in (beads_dir / "embeddeddolt" / db_name, beads_dir / "dolt" / db_name):
        if (candidate / ".dolt").is_dir():
            return candidate
    return None


def read_repo_state(dolt_db_dir: Path) -> dict:
    """Parse .dolt/repo_state.json, or {} if unreadable."""
    path = dolt_db_dir / ".dolt" / "repo_state.json"
    if not path.is_file():
        return {}
    try:
        return json.loads(path.read_text())
    except (OSError, json.JSONDecodeError):
        return {}


def get_head_branch(dolt_db_dir: Path) -> str:
    """Current Dolt branch name from repo_state.json (default 'main')."""
    head = read_repo_state(dolt_db_dir).get("head", "")
    prefix = "refs/heads/"
    return head[len(prefix):] if head.startswith(prefix) else "main"


def read_dolt_remotes(dolt_db_dir: Path) -> dict[str, str]:
    """Map of Dolt remote name -> url from repo_state.json."""
    remotes = read_repo_state(dolt_db_dir).get("remotes", {}) or {}
    return {name: r.get("url", "?") for name, r in remotes.items()}


def get_git_remote_url(project_path: Path) -> str:
    """URL of git remote 'origin', or '(no remote)'."""
    url = run(["git", "remote", "get-url", "origin"], cwd=project_path).stdout.strip()
    return url if url else "(no remote)"


def pick_dolt_remote(dolt_remotes: dict[str, str]) -> str | None:
    """The Dolt remote `bd dolt push` would use: 'origin', else the first."""
    if "origin" in dolt_remotes:
        return "origin"
    return next(iter(dolt_remotes), None)


def read_sync_remote(beads_dir: Path) -> str | None:
    """`sync.remote` from .beads/config.yaml, or None.

    bd writes it in both spellings seen on this machine: a flat
    `sync.remote: "..."` line, and a nested `sync:` block holding
    `remote: "..."`. A line-level read covers both without a YAML parser
    (bdutils stays stdlib-only); commented-out lines are skipped, since bd's
    template ships a commented `# sync.remote:` example.
    """
    try:
        lines = (beads_dir / "config.yaml").read_text().splitlines()
    except OSError:
        return None
    in_sync = False
    for line in lines:
        if not line.strip() or line.lstrip().startswith("#"):
            continue
        key, sep, value = line.strip().partition(":")
        if not sep:
            continue
        if line[0] not in " \t":
            in_sync = key == "sync" and not value.strip()
            if key == "sync.remote":
                return value.strip().strip("'\"") or None
        elif in_sync and key == "remote":
            return value.strip().strip("'\"") or None
    return None


def git_url_for_dolt_remote(url: str) -> str | None:
    """The URL `git ls-remote` can read for a Dolt remote, or None.

    Dolt's git-backed remotes are spelled `git+ssh://` / `git+https://`;
    dropping the `git+` gives the plain git URL. A remote already in git's own
    spelling (scp-style `git@host:path`, `ssh://`, `https://`) passes through.
    Anything else (az://, gs://, DoltHub) is not a git repository at all.

    `./` path segments are dropped, as Dolt drops them: bd stored several of
    this machine's remotes as `git+ssh://git@github.com/./deg/<repo>.git`
    (from an scp-style `git@github.com:deg/<repo>`), which `dolt fetch`
    reads fine and GitHub refuses ("./deg/miluim is not a valid repository
    name") when git passes it through literally.
    """
    if url.startswith(("git+ssh://", "git+https://", "git+http://")):
        scheme, _, rest = url[len("git+"):].partition("://")
        host, _, path = rest.partition("/")
        return f"{scheme}://{host}/{'/'.join(p for p in path.split('/') if p != '.')}"
    if re.match(r"^[\w.-]+@[^:/]+:", url) or url.startswith(("ssh://", "https://", "http://")):
        return url
    return None


def same_repo_url(a: str, b: str) -> bool:
    """True when two remote URLs name the same repository, whatever their
    spelling (`git+ssh://git@h/o/r.git`, `git@h:o/r`, `https://h/o/r`)."""
    def norm(u: str) -> str:
        u = re.sub(r"^(git\+)?[a-z]+://", "", u)
        u = re.sub(r"^[^@/]+@", "", u)
        u = u.replace(":", "/", 1).replace("/./", "/")
        return re.sub(r"(\.git)?/*$", "", u).lower()
    return norm(a) == norm(b)


def resolve_beads_remote(project_path: Path, beads_dir: Path,
                         dolt_remotes: dict[str, str]) -> tuple[str, str] | None:
    """(url, source) of the remote a beads repo's Dolt data is pushed to.

    In the order bd itself trusts them: the Dolt database's own remote (what
    `bd dolt push` actually pushes to), then `sync.remote` from config.yaml,
    and only then the git origin. Going straight to the git origin was
    beads-utils-tj0: ~/Documents/degel/website pushes its beads to a private
    repo while its git origin is public, so the check reported "no off-machine
    copy" about a repo that had one.
    """
    name = pick_dolt_remote(dolt_remotes)
    if name is not None:
        return dolt_remotes[name], f"Dolt remote '{name}'"
    sync_remote = read_sync_remote(beads_dir)
    if sync_remote:
        return sync_remote, "sync.remote"
    git_url = get_git_remote_url(project_path)
    if git_url != "(no remote)":
        return git_url, "git origin"
    return None


# Long enough for a slow link, short enough that one dead git remote does
# not stall a --global run of bd-verify-backup. (`dolt fetch` has no timeout.)
GIT_NETWORK_TIMEOUT = 60


def git_ls_remote_ref(project_path: Path, url: str, ref: str) -> tuple[str | None, str | None]:
    """(hash, error) for `ref` on the git repository at `url`.

    Exactly one is set when git could ask: the hash if the ref exists,
    neither if it does not. `error` (git's last stderr line) means git could
    not ask at all -- offline, refused auth, a repository GitHub rejects --
    which used to read as "never pushed".
    """
    env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
    try:
        result = subprocess.run(["git", "ls-remote", url, ref], cwd=project_path,
                                capture_output=True, text=True, env=env,
                                timeout=GIT_NETWORK_TIMEOUT)
    except subprocess.TimeoutExpired:
        return None, f"timed out after {GIT_NETWORK_TIMEOUT}s"
    if result.returncode != 0:
        lines = [ln.strip() for ln in result.stderr.splitlines() if ln.strip()]
        return None, (lines[-1] if lines else f"git exited {result.returncode}")
    line = result.stdout.strip()
    return (line.split()[0] if line else None), None


@dataclass
class DoltSyncState:
    """Everything bd-dolt-check concludes about one beads repo's Dolt data.

    Computed once by dolt_sync_state() and rendered by each caller:
    bd-dolt-check verbosely, bd-verify-backup as one cell of a table.

    `status` is one of:
      NOT FOUND     the remote has no refs/dolt/data (never pushed)
      UNREACHABLE   git could not ask the remote (`ls_error` says why)
      UNVERIFIABLE  the remote has data, but the commit delta is unknown
                    (`detail` says why)
      IN SYNC / OUT OF SYNC / BEHIND / DIVERGED / DIFFERS
                    the commit comparison; DIFFERS = heads differ but
                    neither range counts
    The working set (`dirty`) is a separate axis: None = could not look,
    [] = clean.
    """
    db_name: str
    dolt_mode: str
    dolt_db_dir: Path | None
    dolt_remotes: dict[str, str]
    beads_remote: tuple[str, str] | None
    dirty: list[dict] | None
    status: str = ""
    detail: str = ""
    ls_url: str | None = None
    remote_ref: str | None = None
    ls_error: str | None = None
    branch: str = "main"
    remote_name: str = "origin"
    tracking_ref: str = ""
    fetch_failed: bool = False
    local_head: str | None = None
    tracking_head: str | None = None
    ahead: int = 0
    behind: int = 0

    @property
    def n_dirty(self) -> int:
        return len(self.dirty) if self.dirty else 0


def dolt_sync_state(project_path: Path, fetch: bool = True) -> DoltSyncState:
    """Compare a beads repo's Dolt data with its remote. Never runs `bd`."""
    beads_dir = project_path / ".beads"
    db_name, dolt_mode = read_metadata(beads_dir)
    dolt_db_dir = locate_dolt_db(beads_dir, db_name)
    dolt_remotes = read_dolt_remotes(dolt_db_dir) if dolt_db_dir else {}
    st = DoltSyncState(
        db_name=db_name, dolt_mode=dolt_mode, dolt_db_dir=dolt_db_dir,
        dolt_remotes=dolt_remotes,
        beads_remote=resolve_beads_remote(project_path, beads_dir, dolt_remotes),
        dirty=dolt_status(dolt_db_dir) if dolt_db_dir is not None else None,
    )

    st.ls_url = git_url_for_dolt_remote(st.beads_remote[0]) if st.beads_remote else None
    if st.beads_remote is None:
        st.status = "NOT FOUND"
        return st
    # Asking the remote is network, so fetch=False skips it and lets the
    # last-fetched tracking ref decide; a repo that never had one ends up
    # UNVERIFIABLE below. A remote that is not a git repository (az://,
    # DoltHub) has no refs/dolt/data to look for at all.
    if fetch and st.ls_url:
        st.remote_ref, st.ls_error = git_ls_remote_ref(project_path, st.ls_url, "refs/dolt/data")
        if st.ls_error:
            st.status = "UNREACHABLE"
            return st
        if st.remote_ref is None:
            st.status = "NOT FOUND"
            return st

    if dolt_db_dir is None or not have_dolt():
        st.status = "UNVERIFIABLE"
        st.detail = "run from the project dir with the dolt CLI installed for a full check"
        return st

    st.branch = get_head_branch(dolt_db_dir)
    st.remote_name = pick_dolt_remote(dolt_remotes) or "origin"
    st.tracking_ref = f"remotes/{st.remote_name}/{st.branch}"
    if fetch:
        st.fetch_failed = not dolt_fetch(dolt_db_dir, st.remote_name)

    st.local_head = dolt_rev(dolt_db_dir, st.branch)
    st.tracking_head = dolt_rev(dolt_db_dir, st.tracking_ref)
    if st.local_head is None:
        st.status = "UNVERIFIABLE"
        st.detail = f"could not read local Dolt HEAD for branch '{st.branch}'"
        return st
    if st.tracking_head is None:
        st.status = "UNVERIFIABLE"
        st.detail = f"no local remote-tracking ref '{st.tracking_ref}'; run 'bd dolt push' once"
        return st
    if st.local_head == st.tracking_head:
        st.status = "IN SYNC"
        return st

    st.ahead = dolt_count_range(dolt_db_dir, f"{st.tracking_ref}..{st.branch}") or 0
    st.behind = dolt_count_range(dolt_db_dir, f"{st.branch}..{st.tracking_ref}") or 0
    if st.ahead and st.behind:
        st.status = "DIVERGED"
    elif st.ahead:
        st.status = "OUT OF SYNC"
    elif st.behind:
        st.status = "BEHIND"
    else:
        st.status = "DIFFERS"
    return st


def dolt_fetch(dolt_db_dir: Path, remote: str) -> bool:
    """Best-effort `dolt fetch <remote>` to refresh remote-tracking refs.

    Returns True on success. Used so a remote that moved ahead (e.g. pushed
    from another machine) is reflected in remotes/<remote>/<branch> before we
    compare against it.
    """
    if not have_dolt():
        return False
    return run(["dolt", "fetch", remote], cwd=dolt_db_dir).returncode == 0


def dolt_rev(dolt_db_dir: Path, ref: str) -> str | None:
    """Commit hash at `ref`, or None if dolt is missing or `ref` doesn't resolve."""
    if not have_dolt():
        return None
    result = run(["dolt", "log", "--oneline", "-n", "1", ref], cwd=dolt_db_dir)
    if result.returncode != 0:
        return None
    lines = strip_ansi(result.stdout).strip().splitlines()
    return lines[0].split()[0] if lines and lines[0] else None


def dolt_count_range(dolt_db_dir: Path, rev_range: str) -> int | None:
    """Number of commits in a Dolt `A..B` range; None if dolt is missing or it errors."""
    if not have_dolt():
        return None
    result = run(["dolt", "log", "--oneline", rev_range], cwd=dolt_db_dir)
    if result.returncode != 0:
        return None
    return sum(1 for ln in result.stdout.splitlines() if ln.strip())


def dolt_log_range(dolt_db_dir: Path, rev_range: str) -> list[str]:
    """One-line commit summaries for a Dolt `A..B` range, oldest last."""
    if not have_dolt():
        return []
    result = run(["dolt", "log", "--oneline", rev_range], cwd=dolt_db_dir)
    if result.returncode != 0:
        return []
    return [strip_ansi(ln) for ln in result.stdout.splitlines() if ln.strip()]


def dolt_sql_json(dolt_db_dir: Path, query: str) -> list[dict] | None:
    """Run a read-only query and return its rows, or None on any failure.

    Dolt emits `{"rows": [...]}`. Note that JSON output omits NULL columns
    entirely, so callers must use `.get()` rather than indexing.
    """
    if not have_dolt():
        return None
    result = run(["dolt", "sql", "-r", "json", "-q", query], cwd=dolt_db_dir)
    if result.returncode != 0:
        return None
    try:
        return json.loads(result.stdout).get("rows", [])
    except (json.JSONDecodeError, AttributeError):
        return None


def dolt_table_columns(dolt_db_dir: Path, table: str, rev: str | None = None) -> list[str]:
    """Column names of `table`, optionally at revision `rev`.

    Needed because a schema migration can add columns, so two revisions of
    the same table may not share a column set — selecting a column that only
    exists on one side is a hard error in Dolt. Callers that compare across
    revisions should intersect the two results.
    """
    if rev:
        # `'` is legal in ref names and `dolt sql -q` runs `;`-separated
        # statements, so an unescaped rev could break out of the literal.
        safe_rev = rev.replace("'", "''")
        target = f"{table} as of '{safe_rev}'"
    else:
        target = table
    rows = dolt_sql_json(dolt_db_dir, f"describe {target};")
    if rows is None:
        return []
    return [r["Field"] for r in rows if "Field" in r]


def dolt_status(dolt_db_dir: Path) -> list[dict] | None:
    """Rows of the `dolt_status` system table: every table with uncommitted
    changes in the working set.

    None means dolt could not be asked (CLI missing, or the query failed);
    `[]` means it was asked and the working set is clean — a clean repo emits
    `{}`, which dolt_sql_json turns into an empty row list.

    Rows are returned raw. Dolt owns this vocabulary (`modified`, `new table`,
    `deleted`, `renamed`, `conflict`, ...), so callers report whatever comes
    back rather than matching against a list that would rot. `staged` is not
    filtered on either: a staged table is still uncommitted.
    """
    return dolt_sql_json(
        dolt_db_dir,
        "select table_name, staged, status from dolt_status order by table_name;",
    )


def _open_pager() -> subprocess.Popen[str] | None:
    if not sys.stdout.isatty():
        return None
    pager_env = os.environ.get("PAGER", "").strip()
    if pager_env:
        cmd = shlex.split(pager_env)
    elif shutil.which("less"):
        cmd = ["less", "-FRX"]
    else:
        return None
    env = os.environ.copy()
    env.setdefault("LESS", "FRX")
    try:
        return subprocess.Popen(cmd, stdin=subprocess.PIPE, text=True, env=env)
    except OSError:
        return None


@contextlib.contextmanager
def paged_output(no_pager: bool = False):
    """Yield a writable stream piped through $PAGER (or 'less -FRX') when
    stdout is a tty; otherwise yield sys.stdout.

    '-F' makes less exit immediately when content fits one screen, so short
    output is indistinguishable from direct-to-stdout. BrokenPipeError is
    swallowed so Ctrl-C / quitting the pager is clean — and on the direct
    stdout path, sys.stdout is redirected to /dev/null afterwards so Python's
    atexit flush can't re-raise when the consumer (e.g. `… | head -1`)
    closed the pipe before we finished writing.
    """
    pager = None if no_pager else _open_pager()
    out = pager.stdin if pager else sys.stdout
    try:
        yield out
    except BrokenPipeError:
        if pager is None:
            try:
                devnull = os.open(os.devnull, os.O_WRONLY)
                os.dup2(devnull, sys.stdout.fileno())
            except OSError:
                pass
    finally:
        if pager is not None:
            with contextlib.suppress(BrokenPipeError):
                pager.stdin.close()
            pager.wait()