arborist/aborist/cli.py

"""Aborist CLI: ingest / search / verify / stats."""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

from aborist import __version__
from aborist.ingest import ingest_source, verify_random_sample
from aborist.progress import Progress
from aborist.search import FTS5Backend
from aborist.sources import WikipediaCurDump
from aborist.store import DEFAULT_DB_PATH, connect, connect_query, stats


def _cmd_ingest(args: argparse.Namespace) -> int:
    if args.source in ("wikipedia_cur", "wikipedia_old"):
        if not args.path:
            print(f"--path is required for {args.source}", file=sys.stderr)
            return 2
        from aborist.sources import WikipediaSqlDump

        table = "cur" if args.source == "wikipedia_cur" else "old"
        shard = None
        if args.shard:
            rank_str, total_str = args.shard.split("/", 1)
            shard = (int(rank_str), int(total_str))
        src = WikipediaSqlDump(path=args.path, table=table, shard=shard)
    elif args.source == "html":  # noqa: SIM114 — keep branch shape
        try:
            from aborist.sources import HtmlPageSource
        except ImportError:
            print(
                "html source requires extras: pip install 'aborist[html]'",
                file=sys.stderr,
            )
            return 2
        urls: list[str] = list(args.url or [])
        if args.urls_from:
            urls.extend(
                line.strip()
                for line in Path(args.urls_from).read_text(encoding="utf-8").splitlines()
                if line.strip() and not line.lstrip().startswith("#")
            )
        if not urls:
            print("html source needs --url or --urls-from", file=sys.stderr)
            return 2
        src = HtmlPageSource(urls, respect_robots=not args.no_robots)
    elif args.source in ("grok_export", "grok_media"):
        if not args.path:
            print(f"--path is required for {args.source}", file=sys.stderr)
            return 2
        from aborist.sources import GrokExportSource, GrokMediaPostsSource

        cls = GrokExportSource if args.source == "grok_export" else GrokMediaPostsSource
        src = cls(path=args.path)
    elif args.source in ("wikipedia_xml", "wikipedia_xml_history", "wikipedia_abstract"):
        if not args.path:
            print(f"--path is required for {args.source}", file=sys.stderr)
            return 2
        from aborist.sources import WikipediaAbstractDump, WikipediaXmlDump

        if args.source == "wikipedia_abstract":
            src = WikipediaAbstractDump(path=args.path)
        else:
            shard = None
            if args.shard:
                rank_str, total_str = args.shard.split("/", 1)
                shard = (int(rank_str), int(total_str))
            src = WikipediaXmlDump(
                path=args.path,
                shard=shard,
                multi_revision=(args.source == "wikipedia_xml_history"),
            )
    elif args.source in ("git_repo", "hg_repo"):
        if not args.path:
            print(f"--path is required for {args.source}", file=sys.stderr)
            return 2
        from aborist.sources import GitRepoSource, MercurialRepoSource

        cls = GitRepoSource if args.source == "git_repo" else MercurialRepoSource
        src = cls(repo_path=args.path)
    else:
        print(f"unknown source: {args.source}", file=sys.stderr)
        return 2

    # Resolve target DB: if --shards-dir is set with --shard, write to a
    # per-shard file. Each shard owns its own SQLite file, so N parallel
    # ingests have ZERO writer-lock contention.
    target_db = args.db
    if args.shards_dir:
        if not args.shard:
            print(
                "--shards-dir requires --shard rank/total",
                file=sys.stderr,
            )
            return 2
        rank_str, total_str = args.shard.split("/", 1)
        rank = int(rank_str)
        total = int(total_str)
        shards_dir = Path(args.shards_dir)
        shards_dir.mkdir(parents=True, exist_ok=True)
        digits = max(3, len(str(total - 1)))
        target_db = shards_dir / f"{rank:0{digits}d}.db"

    progress: Progress | None = None
    if not args.quiet:
        prefix = ""
        if args.shard:
            prefix = f"[shard {args.shard}] "
        progress = Progress(
            interval=args.progress_interval,
            total_estimate=args.total_estimate,
            prefix=prefix,
        )

    conn = connect(target_db)
    try:
        result = ingest_source(
            conn,
            src,
            chunker_name=args.chunker,
            limit=args.limit,
            batch_size=args.batch_size,
            resume=args.resume,
            progress=progress,
        )
    finally:
        conn.close()
    print(json.dumps(result.__dict__, indent=2))
    return 0


def _cmd_search(args: argparse.Namespace) -> int:
    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        backend = FTS5Backend(conn)
        hits = backend.search(args.query, limit=args.limit)
    finally:
        conn.close()
    if args.json:
        print(
            json.dumps(
                [
                    {
                        "document_root": h.document_root,
                        "document_uri": h.document_uri,
                        "chunk_idx": h.chunk_idx,
                        "snippet": h.snippet,
                        "score": h.score,
                        "audit_mode": h.audit_mode.value,
                        "title": h.title,
                    }
                    for h in hits
                ],
                indent=2,
            )
        )
    else:
        for h in hits:
            print(f"[{h.audit_mode.value}] {h.score:7.3f}  {h.title or h.document_uri}")
            print(f"    chunk {h.chunk_idx}  root={h.document_root[:16]}…")
            print(f"    {h.snippet}")
            print()
    return 0


def _cmd_verify(args: argparse.Namespace) -> int:
    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        result = verify_random_sample(conn, n=args.n)
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0 if result["failed"] == 0 else 1


def _cmd_distill(args: argparse.Namespace) -> int:
    from aborist.distill import get_distiller
    from aborist.distill.runner import distill_existing
    from aborist.store import discover_shards

    try:
        distiller = get_distiller(args.process)
    except ValueError as e:
        print(str(e), file=sys.stderr)
        return 2

    # Sharded mode: iterate over each shard's DB and distill in place.
    # Cores stay in their source shard so the per-shard audit/derivation
    # chains remain self-contained.
    if args.global_shards_dir:
        shard_paths = discover_shards(args.global_shards_dir)
        if not shard_paths:
            print(f"no shards in {args.global_shards_dir}", file=sys.stderr)
            return 2

        per_shard: list[dict] = []
        totals = {
            "scanned": 0,
            "distilled": 0,
            "skipped_existing": 0,
            "skipped_cold": 0,
            "skipped_empty": 0,
        }
        for sp in shard_paths:
            conn = connect(sp)
            try:
                r = distill_existing(
                    conn,
                    distiller,
                    kind=args.kind,
                    source_type=args.source_type,
                    limit=args.limit,
                    chunker_name=args.chunker,
                    batch_size=args.batch_size,
                )
            finally:
                conn.close()
            per_shard.append({"shard": sp.name, **r})
            for k in totals:
                totals[k] += r[k]
        print(json.dumps({**totals, "shards": per_shard}, indent=2))
        return 0

    conn = connect(args.db)
    try:
        result = distill_existing(
            conn,
            distiller,
            kind=args.kind,
            source_type=args.source_type,
            limit=args.limit,
            chunker_name=args.chunker,
            batch_size=args.batch_size,
        )
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0


def _cmd_ask(args: argparse.Namespace) -> int:
    import os

    from aborist.qa import ask
    from aborist.qa.client import OpenAICompatibleClient, StubClient

    base_url = args.endpoint or os.environ.get(
        "ABORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1"
    )
    model = args.model or os.environ.get(
        "ABORIST_LLM_MODEL",
        "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic",
    )
    revision = os.environ.get("ABORIST_LLM_REVISION", "")
    quantization = os.environ.get("ABORIST_LLM_QUANTIZATION", "fp8-dynamic")
    api_key = os.environ.get("ABORIST_LLM_API_KEY")

    client: object
    if args.dry_run:
        client = StubClient(
            answer=f"[STUB] would have answered '{args.question}' against root {args.document_root[:16]}…"
        )
    else:
        client = OpenAICompatibleClient(base_url=base_url, api_key=api_key)

    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        result = ask(
            conn,
            document_root=args.document_root,
            question=args.question,
            client=client,
            model_id=model,
            revision=revision,
            quantization=quantization,
        )
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0 if result.get("status") in ("cache_hit", "cache_miss_then_written") else 1


def _cmd_query(args: argparse.Namespace) -> int:
    """Multi-source RAG: question -> top-K corpus docs -> Hermes -> cache."""
    import os

    from aborist.qa.client import OpenAICompatibleClient, StubClient
    from aborist.qa.query import query

    base_url = args.endpoint or os.environ.get(
        "ABORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1"
    )
    model = args.model or os.environ.get(
        "ABORIST_LLM_MODEL",
        "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic",
    )
    revision = os.environ.get("ABORIST_LLM_REVISION", "")
    quantization = os.environ.get("ABORIST_LLM_QUANTIZATION", "fp8-dynamic")
    api_key = os.environ.get("ABORIST_LLM_API_KEY")

    client: object
    if args.dry_run:
        client = StubClient(
            answer="[STUB] dry-run: would have asked Hermes-3 with the assembled context."
        )
    else:
        client = OpenAICompatibleClient(base_url=base_url, api_key=api_key)

    qa_db = args.qa_db
    if qa_db is None:
        if args.global_shards_dir:
            qa_db = Path(args.global_shards_dir) / "qa.db"
        else:
            qa_db = Path.home() / ".aborist" / "qa.db"
    qa_db = Path(qa_db)

    shards_dir = (
        Path(args.global_shards_dir) if args.global_shards_dir else None
    )
    single_db = None if shards_dir else args.db

    result = query(
        question=args.question,
        qa_db=qa_db,
        chat_client=client,
        model_id=model,
        revision=revision,
        quantization=quantization,
        shards_dir=shards_dir,
        single_db=single_db,
        top_k=args.top_k,
        over_fetch=args.over_fetch,
        max_context_chars=args.max_context_chars,
    )

    print(json.dumps(result, indent=2))
    return (
        0
        if result.get("status") in ("cache_hit", "cache_miss_then_written")
        else 1
    )


def _falsify_cache_key(
    cache_key_value: str,
    *,
    state: str,
    reason: str,
    by_actor: str,
    shards_dir: Path | None,
    db_path: Path | None,
) -> dict:
    """Mark a providence_cache record as failed/stale/quarantined across shards.

    Searches every shard for the cache_key (it lives in exactly one).
    Updates the row's falsification_state, appends a falsification log
    entry, and writes a 'falsify' audit event so the chain records the act.
    """
    import time as _time

    from aborist.store import append_audit, discover_shards, transaction

    if state not in ("failed", "stale", "quarantined"):
        return {"status": "invalid_state", "value": state}

    paths: list[Path] = (
        discover_shards(shards_dir) if shards_dir else [Path(db_path)]
    )

    for sp in paths:
        c = connect(sp)
        try:
            row = c.execute(
                "SELECT cache_key, falsification_state FROM providence_cache "
                "WHERE cache_key = ?",
                (cache_key_value,),
            ).fetchone()
            if row is None:
                continue
            now = int(_time.time())
            with transaction(c):
                event_hash = append_audit(
                    c,
                    event_type="falsify",
                    subject_root=cache_key_value,
                    body={
                        "cache_key": cache_key_value,
                        "from_state": row["falsification_state"],
                        "to_state": state,
                        "reason": reason,
                        "by_actor": by_actor,
                    },
                    ts=now,
                )
                c.execute(
                    "UPDATE providence_cache "
                    "SET falsification_state = ?, audit_event_hash = ? "
                    "WHERE cache_key = ?",
                    (state, event_hash, cache_key_value),
                )
                c.execute(
                    "INSERT INTO falsifications "
                    "(cache_key, state, reason, by_actor, at, audit_event_hash) "
                    "VALUES (?, ?, ?, ?, ?, ?)",
                    (cache_key_value, state, reason, by_actor, now, event_hash),
                )
            return {
                "status": "falsified",
                "cache_key": cache_key_value,
                "shard": sp.name,
                "from_state": row["falsification_state"],
                "to_state": state,
                "reason": reason,
                "by_actor": by_actor,
                "audit_event_hash": event_hash,
                "ts": now,
            }
        finally:
            c.close()

    return {"status": "not_found", "cache_key": cache_key_value}


def _cmd_providence(args: argparse.Namespace) -> int:
    """List providence_cache records or falsify one by cache_key."""
    if getattr(args, "falsify", None):
        result = _falsify_cache_key(
            args.falsify,
            state=args.state,
            reason=args.reason or "",
            by_actor=args.by_actor or os.environ.get("USER", "unknown"),
            shards_dir=Path(args.global_shards_dir) if args.global_shards_dir else None,
            db_path=Path(args.db) if args.db else None,
        )
        print(json.dumps(result, indent=2))
        return 0 if result.get("status") == "falsified" else 1
    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        if args.document_uri:
            rows = conn.execute(
                "SELECT cache_key, question_text, answer_text, falsification_state, "
                " hit_count, created_at FROM providence_cache "
                "WHERE document_uri = ? ORDER BY created_at DESC",
                (args.document_uri,),
            ).fetchall()
        elif args.source_root:
            rows = conn.execute(
                "SELECT cache_key, question_text, answer_text, falsification_state, "
                " hit_count, created_at FROM providence_cache "
                "WHERE source_root = ? ORDER BY created_at DESC",
                (args.source_root,),
            ).fetchall()
        else:
            rows = conn.execute(
                "SELECT cache_key, question_text, answer_text, falsification_state, "
                " hit_count, created_at FROM providence_cache "
                "ORDER BY created_at DESC LIMIT ?",
                (args.limit,),
            ).fetchall()
    finally:
        conn.close()
    out = [
        {
            "cache_key": r["cache_key"],
            "state": r["falsification_state"],
            "hit_count": r["hit_count"],
            "question": r["question_text"],
            "answer": r["answer_text"],
            "created_at": r["created_at"],
        }
        for r in rows
    ]
    print(json.dumps(out, indent=2))
    return 0


def _cmd_evict(args: argparse.Namespace) -> int:
    from aborist.evict import evict_to_cold

    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        result = evict_to_cold(
            conn,
            source_type=args.source_type,
            older_than_days=args.older_than_days,
            document_roots=args.document_root or None,
        )
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0


def _cmd_rehydrate(args: argparse.Namespace) -> int:
    from aborist.evict import rehydrate

    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        if args.all_cold:
            roots = [
                r["document_root"]
                for r in conn.execute(
                    "SELECT DISTINCT document_root FROM chunks WHERE tier = 'cold'"
                ).fetchall()
            ]
        else:
            roots = list(args.document_root or [])
        if not roots:
            print(
                "rehydrate needs --document-root R or --all-cold",
                file=sys.stderr,
            )
            return 2
        results = []
        for r in roots:
            res = rehydrate(conn, r)
            res["document_root"] = r
            results.append(res)
    finally:
        conn.close()
    print(json.dumps(results, indent=2))
    drift = sum(1 for r in results if r.get("status") == "drift_detected")
    return 1 if drift else 0


def _cmd_activity(args: argparse.Namespace) -> int:
    """Recent activity: Q&A records + freshly cached docs across all shards.

    Designed for an agent to read before deciding the next action — what was
    just asked, what was just integrated, what's the corpus state.
    """
    import time as _time

    from aborist.store import discover_shards

    shard_paths: list[Path] = []
    if args.global_shards_dir:
        shard_paths = discover_shards(args.global_shards_dir)
    else:
        shard_paths = [Path(args.db)]

    cutoff_ts = 0
    if args.since_seconds:
        cutoff_ts = int(_time.time()) - args.since_seconds

    qa_records: list[dict] = []
    ingest_events: list[dict] = []
    derive_events: list[dict] = []
    falsifications: list[dict] = []
    corpus = {
        "documents_total": 0,
        "documents_surface": 0,
        "documents_core": 0,
        "providence_total": 0,
        "providence_live": 0,
        "providence_stale": 0,
        "providence_failed": 0,
        "audit_events_total": 0,
    }

    for sp in shard_paths:
        c = connect(sp)
        try:
            corpus["documents_total"] += c.execute(
                "SELECT COUNT(*) FROM documents"
            ).fetchone()[0]
            corpus["documents_surface"] += c.execute(
                "SELECT COUNT(*) FROM documents WHERE kind='surface'"
            ).fetchone()[0]
            corpus["documents_core"] += c.execute(
                "SELECT COUNT(*) FROM documents WHERE kind='core'"
            ).fetchone()[0]
            corpus["providence_total"] += c.execute(
                "SELECT COUNT(*) FROM providence_cache"
            ).fetchone()[0]
            corpus["providence_live"] += c.execute(
                "SELECT COUNT(*) FROM providence_cache WHERE falsification_state='live'"
            ).fetchone()[0]
            corpus["providence_stale"] += c.execute(
                "SELECT COUNT(*) FROM providence_cache WHERE falsification_state='stale'"
            ).fetchone()[0]
            corpus["providence_failed"] += c.execute(
                "SELECT COUNT(*) FROM providence_cache WHERE falsification_state='failed'"
            ).fetchone()[0]
            corpus["audit_events_total"] += c.execute(
                "SELECT COUNT(*) FROM audit_events"
            ).fetchone()[0]

            # Q&A records
            for r in c.execute(
                "SELECT cache_key, question_text, answer_text, "
                " falsification_state, hit_count, created_at, last_hit_at, "
                " document_uri FROM providence_cache "
                "WHERE created_at >= ? ORDER BY created_at DESC LIMIT ?",
                (cutoff_ts, args.limit),
            ).fetchall():
                ans = r["answer_text"] or ""
                qa_records.append(
                    {
                        "ts": r["created_at"],
                        "shard": sp.name,
                        "cache_key": r["cache_key"],
                        "question": r["question_text"],
                        "answer_preview": (
                            ans if len(ans) <= args.preview_chars
                            else ans[: args.preview_chars] + "…"
                        ),
                        "sources_uri": r["document_uri"],
                        "state": r["falsification_state"],
                        "hit_count": r["hit_count"],
                        "last_hit_at": r["last_hit_at"],
                    }
                )

            # Recent ingest events
            for r in c.execute(
                "SELECT subject_root, body, ts FROM audit_events "
                "WHERE event_type = 'ingest' AND ts >= ? "
                "ORDER BY ts DESC LIMIT ?",
                (cutoff_ts, args.limit),
            ).fetchall():
                body = json.loads(r["body"]) if r["body"] else {}
                ingest_events.append(
                    {
                        "ts": r["ts"],
                        "shard": sp.name,
                        "document_root": r["subject_root"],
                        "document_uri": body.get("document_uri"),
                        "source_type": body.get("source_type"),
                        "chunks": body.get("chunks"),
                        "supersedes": body.get("supersedes"),
                    }
                )

            # Recent derive events (distillations)
            for r in c.execute(
                "SELECT subject_root, body, ts FROM audit_events "
                "WHERE event_type = 'derive' AND ts >= ? "
                "ORDER BY ts DESC LIMIT ?",
                (cutoff_ts, args.limit),
            ).fetchall():
                body = json.loads(r["body"]) if r["body"] else {}
                derive_events.append(
                    {
                        "ts": r["ts"],
                        "shard": sp.name,
                        "core_root": r["subject_root"],
                        "src_root": body.get("src_root"),
                        "process_id": body.get("process_id"),
                        "compression_ratio": body.get("compression_ratio"),
                        "compression_depth": body.get("compression_depth"),
                    }
                )

            # Recent falsifications
            for r in c.execute(
                "SELECT cache_key, state, reason, by_actor, at FROM falsifications "
                "WHERE at >= ? ORDER BY at DESC LIMIT ?",
                (cutoff_ts, args.limit),
            ).fetchall():
                falsifications.append(
                    {
                        "ts": r["at"],
                        "shard": sp.name,
                        "cache_key": r["cache_key"],
                        "state": r["state"],
                        "reason": r["reason"],
                        "by_actor": r["by_actor"],
                    }
                )
        finally:
            c.close()

    qa_records.sort(key=lambda x: -x["ts"])
    ingest_events.sort(key=lambda x: -x["ts"])
    derive_events.sort(key=lambda x: -x["ts"])
    falsifications.sort(key=lambda x: -x["ts"])

    print(
        json.dumps(
            {
                "as_of": int(_time.time()),
                "shards": [str(p) for p in shard_paths],
                "corpus": corpus,
                "recent_qa": qa_records[: args.limit],
                "recent_ingests": ingest_events[: args.limit],
                "recent_derives": derive_events[: args.limit],
                "recent_falsifications": falsifications[: args.limit],
            },
            indent=2,
        )
    )
    return 0


def _cmd_stats(args: argparse.Namespace) -> int:
    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        result = stats(conn)
    finally:
        conn.close()
    print(json.dumps(result, indent=2))
    return 0


def _check_audit_chain(conn: sqlite3.Connection) -> tuple[int, int]:
    """Return (events_checked, breaks) for one audit chain in `conn`."""
    import hashlib

    rows = conn.execute(
        "SELECT seq, event_hash, prev_event_hash, body FROM audit_events ORDER BY seq"
    ).fetchall()
    prev = None
    breaks = 0
    for r in rows:
        h = hashlib.sha256()
        if r["prev_event_hash"]:
            h.update(bytes.fromhex(r["prev_event_hash"]))
        h.update(r["body"].encode("utf-8"))
        if h.hexdigest() != r["event_hash"]:
            breaks += 1
        if r["prev_event_hash"] != prev:
            breaks += 1
        prev = r["event_hash"]
    return len(rows), breaks


def _cmd_analyze(args: argparse.Namespace) -> int:
    """Compression spectrum, depth distribution, audit chain integrity."""
    from aborist.store import discover_shards

    # In sharded mode, audit chains live per-shard (each shard has its own
    # genesis -> latest). Check each independently and aggregate.
    audit_summary: dict | None = None
    if args.global_shards_dir:
        per_shard_chain = []
        total_events = 0
        total_breaks = 0
        for sp in discover_shards(args.global_shards_dir):
            sc = connect(sp)
            try:
                ev, br = _check_audit_chain(sc)
            finally:
                sc.close()
            per_shard_chain.append({"shard": sp.name, "events": ev, "breaks": br})
            total_events += ev
            total_breaks += br
        audit_summary = {
            "events": total_events,
            "breaks": total_breaks,
            "shards": per_shard_chain,
        }

    conn = (
        connect_query(args.db, shards_dir=args.global_shards_dir)
        if args.global_shards_dir
        else connect(args.db)
    )
    try:
        # Depth distribution.
        depth = conn.execute(
            "SELECT compression_depth, COUNT(*) AS n "
            "FROM documents GROUP BY compression_depth ORDER BY 1"
        ).fetchall()

        # Per-process compression ratios.
        procs = conn.execute(
            "SELECT json_extract(body, '$.process_id') AS process_id, "
            "       json_extract(body, '$.src_kind') AS src_kind, "
            "       AVG(CAST(json_extract(body, '$.compression_ratio') AS REAL)) AS mean_ratio, "
            "       MIN(CAST(json_extract(body, '$.compression_ratio') AS REAL)) AS min_ratio, "
            "       MAX(CAST(json_extract(body, '$.compression_ratio') AS REAL)) AS max_ratio, "
            "       COUNT(*) AS n_events "
            "FROM audit_events WHERE event_type='derive' "
            "GROUP BY process_id, src_kind"
        ).fetchall()

        # Source/kind crosstab.
        kinds = conn.execute(
            "SELECT source_type, kind, COUNT(*) AS n "
            "FROM documents GROUP BY source_type, kind ORDER BY 3 DESC"
        ).fetchall()

        # Tier distribution.
        tiers = conn.execute(
            "SELECT tier, COUNT(*) AS n FROM chunks GROUP BY tier"
        ).fetchall()

        # Top inbound link targets (the 'gravity wells' of the corpus).
        gravity = conn.execute(
            "SELECT dst_uri, COUNT(*) AS inbound FROM edges "
            "WHERE edge_type='wikilink' AND dst_uri != '' "
            "GROUP BY dst_uri ORDER BY inbound DESC LIMIT ?",
            (args.gravity_top,),
        ).fetchall()

        # Audit chain integrity (per-shard if sharded; single chain otherwise).
        if audit_summary is None:
            ev, br = _check_audit_chain(conn)
            audit_summary = {"events": ev, "breaks": br}

        report = {
            "compression_depth_histogram": [
                {"depth": r["compression_depth"], "count": r["n"]} for r in depth
            ],
            "distillers": [
                {
                    "process_id": r["process_id"],
                    "src_kind": r["src_kind"],
                    "n_events": r["n_events"],
                    "compression_ratio": {
                        "mean": (
                            round(r["mean_ratio"], 4)
                            if r["mean_ratio"] is not None
                            else None
                        ),
                        "min": (
                            round(r["min_ratio"], 4)
                            if r["min_ratio"] is not None
                            else None
                        ),
                        "max": (
                            round(r["max_ratio"], 4)
                            if r["max_ratio"] is not None
                            else None
                        ),
                    },
                }
                for r in procs
            ],
            "documents_by_source_kind": [
                {"source_type": r["source_type"], "kind": r["kind"], "count": r["n"]}
                for r in kinds
            ],
            "chunks_by_tier": {r["tier"]: r["n"] for r in tiers},
            "audit_chain": audit_summary,
            "gravity_top_inbound": [
                {"uri": r["dst_uri"], "inbound": r["inbound"]} for r in gravity
            ],
        }
    finally:
        conn.close()
    print(json.dumps(report, indent=2))
    return 0


def _cmd_mesh_status(args: argparse.Namespace) -> int:
    from aborist.mesh import current_epoch, is_enabled, load_identity
    from aborist.mesh.state import roster_at

    conn = connect(args.db)
    try:
        ident = load_identity(conn)
        epoch = current_epoch(conn)
        roster = roster_at(conn, epoch) if epoch is not None else []
        out = {
            "enabled": is_enabled(conn),
            "identity": (
                {
                    "member_id": ident.member_id,
                    "group_name": ident.group_name,
                    "sign_pub_hex": ident.sign_pub.hex(),
                    "dh_pub_hex": ident.dh_pub.hex(),
                    "created_at": ident.created_at,
                }
                if ident
                else None
            ),
            "current_epoch": epoch,
            "roster": [
                {
                    "member_id": m.member_id,
                    "role": m.role,
                    "sign_pub_hex": m.sign_pub.hex(),
                    "dh_pub_hex": m.dh_pub.hex(),
                }
                for m in roster
            ],
        }
    finally:
        conn.close()
    print(json.dumps(out, indent=2))
    return 0


def _cmd_mesh_init(args: argparse.Namespace) -> int:
    from aborist.mesh import init_identity

    conn = connect(args.db)
    try:
        ident = init_identity(conn, group_name=args.group, member_id=args.member_id)
    except RuntimeError as e:
        print(f"error: {e}", file=sys.stderr)
        conn.close()
        return 2
    finally:
        conn.close()
    print(
        json.dumps(
            {
                "member_id": ident.member_id,
                "group_name": ident.group_name,
                "sign_pub_hex": ident.sign_pub.hex(),
                "dh_pub_hex": ident.dh_pub.hex(),
                "note": "share sign_pub_hex + dh_pub_hex with the founder of any group "
                "you want to join. Run 'mesh enable' to flip the gating flag on.",
            },
            indent=2,
        )
    )
    return 0


def _cmd_mesh_enable(args: argparse.Namespace) -> int:
    from aborist.mesh import set_enabled

    conn = connect(args.db)
    try:
        set_enabled(conn, True)
    finally:
        conn.close()
    print(json.dumps({"enabled": True}, indent=2))
    return 0


def _cmd_mesh_disable(args: argparse.Namespace) -> int:
    from aborist.mesh import set_enabled

    conn = connect(args.db)
    try:
        set_enabled(conn, False)
    finally:
        conn.close()
    print(json.dumps({"enabled": False}, indent=2))
    return 0


def _cmd_mesh_members(args: argparse.Namespace) -> int:
    from aborist.mesh import current_epoch
    from aborist.mesh.state import roster_at

    conn = connect(args.db)
    try:
        epoch = current_epoch(conn)
        if epoch is None:
            print(json.dumps({"error": "mesh not initialized"}, indent=2))
            return 2
        roster = roster_at(conn, epoch)
    finally:
        conn.close()
    print(
        json.dumps(
            {
                "epoch": epoch,
                "members": [
                    {
                        "member_id": m.member_id,
                        "role": m.role,
                        "sign_pub_hex": m.sign_pub.hex(),
                        "dh_pub_hex": m.dh_pub.hex(),
                    }
                    for m in roster
                ],
            },
            indent=2,
        )
    )
    return 0


def _cmd_mesh_add(args: argparse.Namespace) -> int:
    from aborist.mesh.members import add_member

    try:
        sign_pub = bytes.fromhex(args.sign_pub)
        dh_pub = bytes.fromhex(args.dh_pub)
    except ValueError:
        print("error: --sign-pub and --dh-pub must be hex-encoded 32-byte keys", file=sys.stderr)
        return 2
    if len(sign_pub) != 32 or len(dh_pub) != 32:
        print("error: keys must decode to exactly 32 bytes", file=sys.stderr)
        return 2

    conn = connect(args.db)
    try:
        epoch = add_member(
            conn,
            member_id=args.member_id,
            sign_pub=sign_pub,
            dh_pub=dh_pub,
            role=args.role,
        )
    except (PermissionError, RuntimeError, ValueError) as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    finally:
        conn.close()
    print(json.dumps({"new_epoch": epoch, "added": args.member_id}, indent=2))
    return 0


def _cmd_mesh_kick(args: argparse.Namespace) -> int:
    from aborist.mesh.members import kick_member

    conn = connect(args.db)
    try:
        epoch = kick_member(conn, member_id=args.member_id, reason=args.reason)
    except (PermissionError, RuntimeError, ValueError) as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    finally:
        conn.close()
    print(
        json.dumps(
            {"new_epoch": epoch, "kicked": args.member_id, "reason": args.reason},
            indent=2,
        )
    )
    return 0


def _cmd_mesh_rotate(args: argparse.Namespace) -> int:
    from aborist.mesh.members import scheduled_rotate

    conn = connect(args.db)
    try:
        epoch = scheduled_rotate(conn, reason=args.reason)
    except (PermissionError, RuntimeError) as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    finally:
        conn.close()
    print(json.dumps({"new_epoch": epoch, "reason": args.reason}, indent=2))
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="aborist",
        description="An arborist for trees and forests of cross-linked information.",
    )
    p.add_argument("--version", action="version", version=f"aborist {__version__}")
    p.add_argument(
        "--db",
        type=Path,
        default=DEFAULT_DB_PATH,
        help=f"path to aborist SQLite db (default: {DEFAULT_DB_PATH})",
    )
    p.add_argument(
        "--shards-dir",
        dest="global_shards_dir",
        default=None,
        help=(
            "for read commands: attach all shards in this directory and "
            "expose them as UNION views over the standard tables"
        ),
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    ingest = sub.add_parser("ingest", help="ingest documents from a source")
    ingest.add_argument(
        "--source",
        required=True,
        choices=[
            "wikipedia_cur",
            "wikipedia_old",
            "wikipedia_xml",
            "wikipedia_xml_history",
            "wikipedia_abstract",
            "html",
            "grok_export",
            "grok_media",
            "git_repo",
            "hg_repo",
        ],
        help="source type",
    )
    ingest.add_argument(
        "--path",
        help=(
            "path to dump file (wikipedia) or to xAI export root / "
            "prod-grok-backend.json (grok_export, grok_media)"
        ),
    )
    ingest.add_argument(
        "--url", action="append", help="URL to ingest (html source; repeatable)"
    )
    ingest.add_argument(
        "--urls-from",
        dest="urls_from",
        help="file with one URL per line (html source)",
    )
    ingest.add_argument(
        "--no-robots",
        dest="no_robots",
        action="store_true",
        help="do not consult robots.txt (use only for explicitly authorized sites)",
    )
    ingest.add_argument(
        "--chunker", default=None, help="chunker name (default: tok-512-v1)"
    )
    ingest.add_argument(
        "--limit", type=int, default=None, help="cap number of documents"
    )
    ingest.add_argument(
        "--batch-size",
        dest="batch_size",
        type=int,
        default=200,
        help="documents per SQLite transaction (default 200)",
    )
    ingest.add_argument(
        "--shard",
        default=None,
        help=(
            "rank/total — yield only every N-th doc for parallel ingest. "
            "spawn N processes, each with --shard 0/N, 1/N, ... they "
            "parallelize parser CPU and serialize writes via WAL"
        ),
    )
    ingest.add_argument(
        "--shards-dir",
        dest="shards_dir",
        default=None,
        help=(
            "directory for attach-forever sharding. With --shard rank/total, "
            "writes to shards-dir/<rank>.db instead of --db, removing the "
            "WAL writer-lock contention entirely. Reads via aborist --shards-dir"
        ),
    )
    ingest.add_argument(
        "--resume",
        action="store_true",
        help=(
            "rsync-style: read each source's last high-water mark from this "
            "DB's meta table and skip rows whose id is <= it. Safe to kill "
            "and restart at any time"
        ),
    )
    ingest.add_argument(
        "--quiet",
        action="store_true",
        help="suppress periodic stderr progress output",
    )
    ingest.add_argument(
        "--progress-interval",
        dest="progress_interval",
        type=float,
        default=2.0,
        help="seconds between stderr progress lines (default 2.0)",
    )
    ingest.add_argument(
        "--total-estimate",
        dest="total_estimate",
        type=int,
        default=None,
        help=(
            "estimated total docs the source will yield. enables percent "
            "+ ETA in progress output"
        ),
    )
    ingest.set_defaults(func=_cmd_ingest)

    search = sub.add_parser("search", help="keyword search (VISUAL audit mode)")
    search.add_argument("query", help="query string")
    search.add_argument("--limit", type=int, default=20)
    search.add_argument("--json", action="store_true", help="output JSON")
    search.set_defaults(func=_cmd_search)

    verify = sub.add_parser(
        "verify", help="round-trip Merkle proofs for N random documents"
    )
    verify.add_argument("-n", type=int, default=10)
    verify.set_defaults(func=_cmd_verify)

    distill = sub.add_parser(
        "distill",
        help="compress docs into Merkle-signed cores (surface->core or core->core)",
    )
    distill.add_argument(
        "--process", default="first-sentence-v1", help="distiller name"
    )
    distill.add_argument(
        "--kind",
        choices=["surface", "core"],
        default="surface",
        help="source kind to scan; 'core' runs recursive distillation",
    )
    distill.add_argument(
        "--source-type",
        dest="source_type",
        default=None,
        help="restrict to one source_type",
    )
    distill.add_argument(
        "--chunker", default=None, help="chunker for the core doc"
    )
    distill.add_argument(
        "--limit", type=int, default=None, help="cap number of docs scanned"
    )
    distill.add_argument(
        "--batch-size",
        dest="batch_size",
        type=int,
        default=200,
        help="cores written per SQLite transaction (default 200)",
    )
    distill.set_defaults(func=_cmd_distill)

    ask_cmd = sub.add_parser(
        "ask",
        help="answer a question about a document (cache-first, STRICT)",
    )
    ask_cmd.add_argument(
        "--document-root",
        dest="document_root",
        required=True,
        help="document_root to ask about",
    )
    ask_cmd.add_argument(
        "--question", required=True, help="question text"
    )
    ask_cmd.add_argument(
        "--model",
        default=None,
        help="model_id (default $ABORIST_LLM_MODEL or hermes-3)",
    )
    ask_cmd.add_argument(
        "--endpoint",
        default=None,
        help="OpenAI-compatible base URL (default $ABORIST_LLM_ENDPOINT)",
    )
    ask_cmd.add_argument(
        "--dry-run",
        dest="dry_run",
        action="store_true",
        help="use StubClient — no network call",
    )
    ask_cmd.set_defaults(func=_cmd_ask)

    query_cmd = sub.add_parser(
        "query",
        help="multi-source RAG: question -> top-K corpus docs -> Hermes -> cache",
    )
    query_cmd.add_argument("question", help="the question to ask")
    query_cmd.add_argument(
        "--top-k", dest="top_k", type=int, default=8,
        help="max distinct source documents in context (default 8)",
    )
    query_cmd.add_argument(
        "--over-fetch", dest="over_fetch", type=int, default=32,
        help="FTS5 hits to fetch per shard before dedup (default 32)",
    )
    query_cmd.add_argument(
        "--max-context-chars", dest="max_context_chars", type=int, default=60000,
        help="cap on assembled context bytes (default 60000)",
    )
    query_cmd.add_argument(
        "--qa-db", dest="qa_db", default=None,
        help=(
            "providence_cache target DB. default: <shards-dir>/qa.db, or "
            "~/.aborist/qa.db when no shards-dir"
        ),
    )
    query_cmd.add_argument(
        "--model", default=None,
        help="model_id (default $ABORIST_LLM_MODEL or hermes-3)",
    )
    query_cmd.add_argument(
        "--endpoint", default=None,
        help="OpenAI-compatible base URL (default $ABORIST_LLM_ENDPOINT)",
    )
    query_cmd.add_argument(
        "--dry-run", dest="dry_run", action="store_true",
        help="use StubClient — assembles context but skips the LLM call",
    )
    query_cmd.set_defaults(func=_cmd_query)

    prov_cmd = sub.add_parser(
        "providence",
        help="list or falsify providence_cache records",
    )
    prov_cmd.add_argument("--document-uri", dest="document_uri", default=None)
    prov_cmd.add_argument("--source-root", dest="source_root", default=None)
    prov_cmd.add_argument("--limit", type=int, default=20)
    prov_cmd.add_argument(
        "--falsify",
        default=None,
        metavar="CACHE_KEY",
        help=(
            "mark a providence_cache record as failed/stale/quarantined. "
            "Lookups will skip it. Audit chain records the act"
        ),
    )
    prov_cmd.add_argument(
        "--state",
        default="failed",
        choices=["failed", "stale", "quarantined"],
        help="falsification state to set (default: failed)",
    )
    prov_cmd.add_argument(
        "--reason",
        default=None,
        help="reason text stored in falsifications log",
    )
    prov_cmd.add_argument(
        "--by-actor",
        dest="by_actor",
        default=None,
        help="who is falsifying (default: $USER)",
    )
    prov_cmd.set_defaults(func=_cmd_providence)

    evict_cmd = sub.add_parser(
        "evict",
        help="demote surface chunks hot→cold (NULL content, retain leaf_hash)",
    )
    evict_cmd.add_argument(
        "--source-type",
        dest="source_type",
        default=None,
        help="restrict to one source_type",
    )
    evict_cmd.add_argument(
        "--older-than-days",
        dest="older_than_days",
        type=int,
        default=None,
        help="only evict docs older than N days",
    )
    evict_cmd.add_argument(
        "--document-root",
        action="append",
        default=None,
        help="explicit document_root(s) to evict; repeatable",
    )
    evict_cmd.set_defaults(func=_cmd_evict)

    rehydrate_cmd = sub.add_parser(
        "rehydrate",
        help="refetch URI, verify leaves, restore cold content if root matches",
    )
    rehydrate_cmd.add_argument(
        "--document-root",
        action="append",
        default=None,
        help="explicit document_root(s) to rehydrate; repeatable",
    )
    rehydrate_cmd.add_argument(
        "--all-cold",
        dest="all_cold",
        action="store_true",
        help="rehydrate every document with cold chunks",
    )
    rehydrate_cmd.set_defaults(func=_cmd_rehydrate)

    activity_cmd = sub.add_parser(
        "activity",
        help="recent Q&A + freshly cached docs (agent-readable timeline)",
    )
    activity_cmd.add_argument(
        "--limit", type=int, default=10,
        help="max items per category (default 10)",
    )
    activity_cmd.add_argument(
        "--since-seconds",
        dest="since_seconds",
        type=int,
        default=0,
        help="only events newer than this many seconds (0 = all time, default)",
    )
    activity_cmd.add_argument(
        "--preview-chars",
        dest="preview_chars",
        type=int,
        default=240,
        help="answer preview length (default 240 chars)",
    )
    activity_cmd.set_defaults(func=_cmd_activity)

    stats_cmd = sub.add_parser("stats", help="counts: docs, chunks, edges, audit")
    stats_cmd.set_defaults(func=_cmd_stats)

    analyze_cmd = sub.add_parser(
        "analyze",
        help="compression spectrum, depth distribution, audit chain integrity",
    )
    analyze_cmd.add_argument(
        "--gravity-top",
        dest="gravity_top",
        type=int,
        default=10,
        help="N top inbound-linked URIs to report (default 10)",
    )
    analyze_cmd.set_defaults(func=_cmd_analyze)

    # ----- mesh subcommands (off by default) ---------------------------------
    mesh_cmd = sub.add_parser(
        "mesh",
        help="federation/gossip layer (off by default; opt-in via 'mesh enable')",
    )
    mesh_sub = mesh_cmd.add_subparsers(dest="mesh_op", required=True)

    mesh_status = mesh_sub.add_parser("status", help="show enabled flag, identity, current epoch + roster")
    mesh_status.set_defaults(func=_cmd_mesh_status)

    mesh_init = mesh_sub.add_parser("init", help="generate this peer's keys; create epoch 0")
    mesh_init.add_argument("--group", required=True, help="group name")
    mesh_init.add_argument("--member-id", dest="member_id", default=None, help="optional fixed member id (default: random 8-hex)")
    mesh_init.set_defaults(func=_cmd_mesh_init)

    mesh_enable = mesh_sub.add_parser("enable", help="flip the mesh.enabled flag on")
    mesh_enable.set_defaults(func=_cmd_mesh_enable)

    mesh_disable = mesh_sub.add_parser("disable", help="flip the mesh.enabled flag off")
    mesh_disable.set_defaults(func=_cmd_mesh_disable)

    mesh_members = mesh_sub.add_parser("members", help="list current epoch's roster")
    mesh_members.set_defaults(func=_cmd_mesh_members)

    mesh_add = mesh_sub.add_parser("add", help="admin-only: add a peer to the roster (bumps epoch)")
    mesh_add.add_argument("--member-id", dest="member_id", required=True)
    mesh_add.add_argument("--sign-pub", dest="sign_pub", required=True, help="hex Ed25519 pubkey (32 bytes / 64 hex chars)")
    mesh_add.add_argument("--dh-pub", dest="dh_pub", required=True, help="hex X25519 pubkey")
    mesh_add.add_argument("--role", choices=["admin", "member"], default="member")
    mesh_add.set_defaults(func=_cmd_mesh_add)

    mesh_kick = mesh_sub.add_parser("kick", help="admin-only: evict a peer (bumps epoch; old signatures stay valid, new gossip is opaque to them)")
    mesh_kick.add_argument("--member-id", dest="member_id", required=True)
    mesh_kick.add_argument("--reason", required=True)
    mesh_kick.set_defaults(func=_cmd_mesh_kick)

    mesh_rotate = mesh_sub.add_parser("rotate", help="refresh epoch secret without changing roster")
    mesh_rotate.add_argument("--reason", default="scheduled")
    mesh_rotate.set_defaults(func=_cmd_mesh_rotate)

    return p


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())