arborist/aborist/store.py

"""SQLite-backed v9.8 store.

Schema implements the Merkle-AGI v9.8 admissibility ledger:
- 8-dim providence_cache key (source_root, question_hash, model_profile_hash,
  conversation_hash, governance_policy_hash, schema_version,
  canonicalization_version, chunking_version)
- falsification_state ∈ {live, failed, stale, quarantined}
- audit_events append-only chain (event_hash chains via prev_event_hash)
- documents.kind ∈ {surface, core} for layered compression
- chunks.tier ∈ {hot, warm, cold} for reversible eviction
- derivations table binds core docs back to source surface roots

The providence_cache layer is schema-only in Phase 0 — no Q&A inference yet.
"""

from __future__ import annotations

import json
import sqlite3
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator


DEFAULT_DB_PATH = Path.home() / ".aborist" / "aborist.db"


SCHEMA_SQL = """
PRAGMA journal_mode = WAL;
PRAGMA foreign_keys = ON;

CREATE TABLE IF NOT EXISTS schema_meta (
    key   TEXT PRIMARY KEY,
    value TEXT NOT NULL
);

-- Free-form per-DB metadata. Used by the resume mechanic to track each
-- source's high-water mark so a stopped ingest can rsync forward without
-- re-parsing rows that are already in this DB.
CREATE TABLE IF NOT EXISTS meta (
    key   TEXT PRIMARY KEY,
    value TEXT NOT NULL,
    updated_at INTEGER
);

-- Documents: surface (raw ingest) or core (distilled, Merkle-signed back).
CREATE TABLE IF NOT EXISTS documents (
    document_root            TEXT PRIMARY KEY,    -- hex sha256 of merkle root
    document_uri             TEXT NOT NULL,
    source_type              TEXT NOT NULL,
    kind                     TEXT NOT NULL DEFAULT 'surface'
                                CHECK (kind IN ('surface','core')),
    compression_depth        INTEGER NOT NULL DEFAULT 0,
    title                    TEXT,
    chunking_version         TEXT NOT NULL,
    canonicalization_version TEXT NOT NULL,
    schema_version           TEXT NOT NULL,
    ingest_ts                INTEGER NOT NULL,
    hit_count                INTEGER NOT NULL DEFAULT 0,
    last_hit_at              INTEGER
);
CREATE INDEX IF NOT EXISTS idx_documents_uri  ON documents(document_uri);
CREATE INDEX IF NOT EXISTS idx_documents_kind ON documents(kind);

-- Per-document HTTP metadata for cheap recrawl-checks. ETag and
-- Last-Modified come from the response headers at ingest time and let a
-- recrawl pass send conditional HEAD requests (If-None-Match /
-- If-Modified-Since); a 304 lets us skip the body fetch entirely.
-- last_status / last_checked_at record the most recent recheck so
-- operators can audit "when was this URL last verified."
CREATE TABLE IF NOT EXISTS document_http_meta (
    document_root    TEXT PRIMARY KEY,
    etag             TEXT,
    last_modified    TEXT,
    last_fetched_at  INTEGER NOT NULL,
    last_status      INTEGER,
    last_checked_at  INTEGER,
    FOREIGN KEY (document_root) REFERENCES documents(document_root) ON DELETE CASCADE
);

-- Chunks with tier-based reversible eviction.
-- content nullable: cold tier evicts content but retains leaf_hash + URI for
-- rehydration. Identity verified on rehydrate by recomputing leaf_hash.
--
-- chunk_id INTEGER PRIMARY KEY AUTOINCREMENT serves dual duty: it's both the
-- primary key and the rowid that the contentless FTS5 virtual table joins
-- against. The (document_root, idx) UNIQUE constraint preserves the prior
-- "one chunk per (doc, position)" invariant for callers that look up by it.
CREATE TABLE IF NOT EXISTS chunks (
    chunk_id      INTEGER PRIMARY KEY AUTOINCREMENT,
    document_root TEXT NOT NULL,
    idx           INTEGER NOT NULL,
    leaf_hash     TEXT NOT NULL,
    content       TEXT,
    tier          TEXT NOT NULL DEFAULT 'hot'
                     CHECK (tier IN ('hot','warm','cold')),
    UNIQUE (document_root, idx),
    FOREIGN KEY (document_root) REFERENCES documents(document_root) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_chunks_leaf ON chunks(leaf_hash);

-- Interior Merkle nodes (layer >= 1). Layer 0 lives in chunks.leaf_hash.
CREATE TABLE IF NOT EXISTS merkle_nodes (
    document_root TEXT NOT NULL,
    layer         INTEGER NOT NULL,
    idx           INTEGER NOT NULL,
    hash          TEXT NOT NULL,
    PRIMARY KEY (document_root, layer, idx),
    FOREIGN KEY (document_root) REFERENCES documents(document_root) ON DELETE CASCADE
);

-- Cross-links between documents (the forest).
-- Unresolved forward links (dst not yet ingested) carry dst_root='' and the
-- ingest pass backfills dst_root when the target appears.
--
-- WITHOUT ROWID: the PK covers every column, so a default rowid-based table
-- would near-duplicate the row data in the PK index. WITHOUT ROWID makes
-- the table itself a B-tree keyed on the PK and saves ~50% of edge storage
-- on real Wikipedia ingests (measured: 38 MB -> 21 MB / 1000 docs).
-- Behaviorally identical; only the on-disk layout changes.
CREATE TABLE IF NOT EXISTS edges (
    src_root  TEXT NOT NULL,
    dst_root  TEXT NOT NULL DEFAULT '',      -- '' = unresolved, backfilled later
    dst_uri   TEXT NOT NULL DEFAULT '',      -- always present so we can resolve later
    edge_type TEXT NOT NULL,                 -- wikilink, citation, derived_from, ...
    anchor    TEXT NOT NULL DEFAULT '',      -- chunk index or fragment, '' if N/A
    PRIMARY KEY (src_root, edge_type, dst_root, dst_uri, anchor)
) WITHOUT ROWID;
CREATE INDEX IF NOT EXISTS idx_edges_dst_root ON edges(dst_root) WHERE dst_root <> '';
-- idx_edges_dst_uri intentionally omitted: only the gravity_top_inbound
-- analytical query in cli.py filters on dst_uri alone, and a full scan +
-- sort over edges is acceptable for that one-shot reporting path.

-- Distillation: core_root <- src_root with Merkle-signed proof binding.
CREATE TABLE IF NOT EXISTS derivations (
    core_root    TEXT NOT NULL,
    src_root     TEXT NOT NULL,
    proof_blob   TEXT NOT NULL,              -- JSON merkle proof
    process_id   TEXT NOT NULL,              -- distillation process identifier
    distilled_at INTEGER NOT NULL,
    PRIMARY KEY (core_root, src_root, process_id),
    FOREIGN KEY (core_root) REFERENCES documents(document_root) ON DELETE CASCADE,
    FOREIGN KEY (src_root)  REFERENCES documents(document_root) ON DELETE CASCADE
);

-- v9.8 providence cache: 8-dim admissibility key + falsification state.
-- Schema-only in Phase 0 (no Q&A runs yet); ready for Phase 1.
CREATE TABLE IF NOT EXISTS providence_cache (
    cache_key                TEXT PRIMARY KEY,
    source_root              TEXT NOT NULL,
    document_uri             TEXT NOT NULL,
    question_hash            TEXT NOT NULL,
    question_text            TEXT NOT NULL,
    answer_text              TEXT NOT NULL,
    merkle_proof             TEXT NOT NULL,   -- JSON
    model_profile_hash       TEXT NOT NULL,   -- model_id + revision + quantization
    conversation_hash        TEXT NOT NULL,
    governance_policy_hash   TEXT NOT NULL,
    schema_version           TEXT NOT NULL,
    canonicalization_version TEXT NOT NULL,
    chunking_version         TEXT NOT NULL,
    falsification_state      TEXT NOT NULL DEFAULT 'live'
                                CHECK (falsification_state IN ('live','failed','stale','quarantined')),
    chain                    TEXT NOT NULL DEFAULT 'private'
                                CHECK (chain IN ('private','public')),
    audit_event_hash         TEXT,            -- latest audit event for this record
    created_at               INTEGER NOT NULL,
    last_hit_at              INTEGER,
    hit_count                INTEGER NOT NULL DEFAULT 0,
    -- v9.8 audit_mode trichotomy (RAG-adapted vocabulary; substrate calls
    -- UNGROUNDED "VISUAL"): STRICT (every quote in answer verified against
    -- context), HYBRID (some claims verified, some emergent), UNGROUNDED
    -- (no verbatim grounding — purely emergent from training).
    -- Default UNGROUNDED: an unclassified record is the weakest claim.
    audit_mode               TEXT NOT NULL DEFAULT 'UNGROUNDED'
                                CHECK (audit_mode IN ('STRICT','HYBRID','UNGROUNDED')),
    n_quotes                 INTEGER NOT NULL DEFAULT 0,
    n_verified               INTEGER NOT NULL DEFAULT 0,
    -- JSON array of quoted spans the model produced but we couldn't find
    -- verbatim in context. Mining these surfaces "what the model emerged
    -- beyond the corpus" — candidate ingest targets.
    unverified_quotes        TEXT,
    -- Which verifier strategy classified this record. 'quote' = model
    -- followed the format and wrapped claims in double quotes. 'span' =
    -- bullet/sentence-level substring match in context. 'entity' = no
    -- spans matched but multi-word proper nouns did. 'none' = no
    -- evidence at all (truly emergent).
    verifier_method          TEXT NOT NULL DEFAULT 'none'
                                CHECK (verifier_method IN ('quote','span','entity','none'))
);
CREATE INDEX IF NOT EXISTS idx_providence_root  ON providence_cache(source_root);
CREATE INDEX IF NOT EXISTS idx_providence_state ON providence_cache(falsification_state);
-- idx_providence_audit lives in _migrate_audit_mode() so legacy shards
-- (where audit_mode column gets added by ALTER TABLE) don't trip this
-- script before the migration runs.

-- Append-only audit chain. event_hash = sha256(prev_event_hash || canonical(body)).
CREATE TABLE IF NOT EXISTS audit_events (
    seq             INTEGER PRIMARY KEY AUTOINCREMENT,
    event_hash      TEXT NOT NULL UNIQUE,
    prev_event_hash TEXT,                    -- NULL for genesis
    event_type      TEXT NOT NULL,           -- ingest|falsify|evict_warm|evict_cold|derive|rehydrate|...
    subject_root    TEXT,                    -- document_root or cache_key
    body            TEXT NOT NULL,           -- canonical JSON
    ts              INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_audit_subject ON audit_events(subject_root);

-- Falsification log: which records were marked failed/stale/quarantined and why.
CREATE TABLE IF NOT EXISTS falsifications (
    cache_key        TEXT NOT NULL,
    state            TEXT NOT NULL,
    reason           TEXT,
    by_actor         TEXT,
    at               INTEGER NOT NULL,
    audit_event_hash TEXT NOT NULL,
    PRIMARY KEY (cache_key, at)
);

-- Snapshots: corpus-level Merkle root pinning a forest state at a point in
-- time. snapshot_root = MerkleTree.build([sorted document_roots]). Audit-
-- chain-linked so peers can verify a claimed snapshot was actually witnessed
-- by this instance. parent_snapshot lets snapshots chain (A -> B -> C) for
-- diff/replay. doc_count is informational; the root is the canonical id.
CREATE TABLE IF NOT EXISTS snapshots (
    snapshot_root      TEXT PRIMARY KEY,
    taken_at           INTEGER NOT NULL,
    audit_event_hash   TEXT NOT NULL,
    doc_count          INTEGER NOT NULL,
    parent_snapshot    TEXT,
    reason             TEXT
);
CREATE INDEX IF NOT EXISTS idx_snapshots_taken_at ON snapshots(taken_at);

-- Mesh layer tables. Off by default — populated only when the user runs
-- `aborist mesh init`. Never accessed by ingest / query / distill paths;
-- mesh state is opt-in plumbing for federated peers (see aborist.mesh).
CREATE TABLE IF NOT EXISTS mesh_identity (
    id           INTEGER PRIMARY KEY CHECK (id = 1),  -- singleton
    member_id    TEXT NOT NULL UNIQUE,
    sign_priv    BLOB NOT NULL,                       -- ed25519 32B raw
    sign_pub     BLOB NOT NULL,                       -- ed25519 32B raw
    dh_priv      BLOB NOT NULL,                       -- x25519 32B raw
    dh_pub       BLOB NOT NULL,                       -- x25519 32B raw
    group_name   TEXT NOT NULL,
    created_at   INTEGER NOT NULL
);

-- Per-epoch roster. epoch 0 = group genesis (founder only). Each membership
-- mutation (join, kick, scheduled rotate) bumps the epoch_id by 1 and writes
-- a fresh row-set capturing the new roster.
CREATE TABLE IF NOT EXISTS mesh_roster (
    epoch_id     INTEGER NOT NULL,
    member_id    TEXT NOT NULL,
    sign_pub     BLOB NOT NULL,
    dh_pub       BLOB NOT NULL,
    role         TEXT NOT NULL DEFAULT 'member'
                    CHECK (role IN ('admin','member')),
    PRIMARY KEY (epoch_id, member_id)
);
CREATE INDEX IF NOT EXISTS idx_mesh_roster_member ON mesh_roster(member_id);

-- Epoch lifecycle log. secret_envelope is JSON of the form
--    {"member_id": {"nonce_b64": "...", "ct_b64": "..."}, ...}
-- where each entry is the symmetric epoch secret AEAD-wrapped to that
-- member's X25519 pubkey via ECDH. Eviction happens by NOT including the
-- evicted member's entry in the next epoch's envelope.
CREATE TABLE IF NOT EXISTS mesh_epochs (
    epoch_id           INTEGER PRIMARY KEY,
    started_at         INTEGER NOT NULL,
    started_event_hash TEXT NOT NULL,
    secret_envelope    TEXT NOT NULL,
    reason             TEXT
);

-- Per-peer audit-chain tracking. Each row records the most recent
-- event_hash a given peer has broadcast to us; we enforce that every
-- subsequent gossip envelope carries `prev_event_hash == last_event_hash`
-- of that peer. A mismatch is a fork — the gossip is rejected (409).
-- last_seq is the local count of accepted envelopes from that peer
-- (informational; the canonical chain identity is last_event_hash).
CREATE TABLE IF NOT EXISTS mesh_peer_chains (
    peer_member_id    TEXT PRIMARY KEY,
    last_event_hash   TEXT NOT NULL,
    last_seq          INTEGER NOT NULL,
    last_seen_at      INTEGER NOT NULL
);

-- FTS5 over chunk content for UNGROUNDED-mode keyword search.
--
-- Contentless mode (`content=''`): FTS5 stores ONLY the inverted index, no
-- copy of the indexed text. This eliminates the ~28 MB / 1000 docs that the
-- prior schema spent on chunks_fts_content (the stored copy was redundant
-- with chunks.content). The trade: snippet() / highlight() return empty
-- in contentless mode, so the FTS5 backend builds snippets in Python by
-- joining `chunks_fts.rowid = chunks.chunk_id`, decompressing chunks.content,
-- and locating query tokens.
--
-- Inserts use `INSERT INTO chunks_fts (rowid, content) VALUES (chunk_id, plain)`
-- — the rowid must equal the chunks.chunk_id of the underlying row so the
-- search-time join lines up.
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
    content,
    content='',
    contentless_delete=1,
    tokenize = 'porter unicode61'
);
"""


def connect(db_path: Path | str = DEFAULT_DB_PATH) -> sqlite3.Connection:
    """Open a writable connection, creating the parent dir + schema if needed.

    Performance pragmas applied per-connection. Under WAL (set in the schema):
    - synchronous=NORMAL skips the per-commit fsync; durable up to the last
      checkpoint (SQLite auto-checkpoints at WAL ~1000 frames).
    - cache_size=-65536 = 64 MB page cache (reduces re-reads).
    - temp_store=MEMORY keeps temp tables in RAM (no /tmp churn).
    - mmap_size=256 MB lets reads come from page-cache without read() syscalls.
    """
    p = Path(db_path)
    p.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(p, isolation_level=None)  # autocommit; we'll BEGIN manually
    conn.row_factory = sqlite3.Row
    conn.executescript(SCHEMA_SQL)
    _migrate_audit_mode(conn)
    _migrate_mesh_peer_chains(conn)
    _migrate_document_http_meta(conn)
    conn.execute("PRAGMA synchronous = NORMAL")
    conn.execute("PRAGMA cache_size = -65536")
    conn.execute("PRAGMA temp_store = MEMORY")
    conn.execute("PRAGMA mmap_size = 268435456")
    return conn


def _migrate_audit_mode(conn: sqlite3.Connection) -> None:
    """Forward-migrate pre-v9.8-audit-mode providence_cache shards.

    Adds audit_mode + n_quotes + n_verified + unverified_quotes + verifier_method
    columns to DBs that pre-date the faithfulness-classification rollout. SQLite
    ALTER TABLE ADD COLUMN is O(1) (metadata-only) so this is cheap on every
    open. Idempotent — checks PRAGMA before each ADD.

    Also handles the VISUAL → UNGROUNDED rename for the audit_mode value
    space. SQLite cannot ALTER a column's CHECK in place, so legacy tables
    with the old `CHECK (audit_mode IN ('STRICT','HYBRID','VISUAL'))` get
    rebuilt via the standard temp-table dance, with values translated.
    """
    cols = {row["name"] for row in conn.execute("PRAGMA table_info(providence_cache)")}
    if "audit_mode" not in cols:
        conn.execute(
            "ALTER TABLE providence_cache ADD COLUMN audit_mode TEXT "
            "NOT NULL DEFAULT 'UNGROUNDED' "
            "CHECK (audit_mode IN ('STRICT','HYBRID','UNGROUNDED'))"
        )
    if "n_quotes" not in cols:
        conn.execute(
            "ALTER TABLE providence_cache ADD COLUMN n_quotes INTEGER NOT NULL DEFAULT 0"
        )
    if "n_verified" not in cols:
        conn.execute(
            "ALTER TABLE providence_cache ADD COLUMN n_verified INTEGER NOT NULL DEFAULT 0"
        )
    if "unverified_quotes" not in cols:
        conn.execute(
            "ALTER TABLE providence_cache ADD COLUMN unverified_quotes TEXT"
        )
    if "verifier_method" not in cols:
        conn.execute(
            "ALTER TABLE providence_cache ADD COLUMN verifier_method TEXT "
            "NOT NULL DEFAULT 'none' "
            "CHECK (verifier_method IN ('quote','span','entity','none'))"
        )

    # VISUAL → UNGROUNDED rename. Detect legacy CHECK by inspecting DDL.
    ddl_row = conn.execute(
        "SELECT sql FROM sqlite_master WHERE type='table' AND name='providence_cache'"
    ).fetchone()
    if ddl_row and "'VISUAL'" in (ddl_row[0] or ""):
        _rebuild_providence_cache_ungrounded(conn)

    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_providence_audit "
        "ON providence_cache(audit_mode)"
    )


def _migrate_document_http_meta(conn: sqlite3.Connection) -> None:
    """Forward-migrate pre-recrawl-check shards.

    Adds ``document_http_meta`` to DBs that pre-date the recrawl-check
    feature. Mirrors ``_migrate_mesh_peer_chains``: idempotent
    table-existence probe, then CREATE if missing.
    """
    row = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name='document_http_meta'"
    ).fetchone()
    if row is None:
        conn.execute(
            "CREATE TABLE document_http_meta ("
            " document_root    TEXT PRIMARY KEY,"
            " etag             TEXT,"
            " last_modified    TEXT,"
            " last_fetched_at  INTEGER NOT NULL,"
            " last_status      INTEGER,"
            " last_checked_at  INTEGER,"
            " FOREIGN KEY (document_root) REFERENCES documents(document_root)"
            "  ON DELETE CASCADE"
            ")"
        )


def _migrate_mesh_peer_chains(conn: sqlite3.Connection) -> None:
    """Forward-migrate pre-mesh-fork-detection shards.

    Adds the `mesh_peer_chains` table to DBs that pre-date per-peer
    audit-chain tracking on the mesh wire. CREATE TABLE IF NOT EXISTS
    in SCHEMA_SQL covers brand-new shards; this migration is a belt-
    and-suspenders idempotency check for callers that bypass the full
    SCHEMA_SQL pass (cross-shard query views, etc.). Idempotent —
    PRAGMA-checks before issuing CREATE.
    """
    row = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name='mesh_peer_chains'"
    ).fetchone()
    if row is None:
        conn.execute(
            "CREATE TABLE mesh_peer_chains ("
            " peer_member_id    TEXT PRIMARY KEY,"
            " last_event_hash   TEXT NOT NULL,"
            " last_seq          INTEGER NOT NULL,"
            " last_seen_at      INTEGER NOT NULL"
            ")"
        )


def _rebuild_providence_cache_ungrounded(conn: sqlite3.Connection) -> None:
    """One-time table rebuild: rename audit_mode value VISUAL → UNGROUNDED.

    SQLite cannot modify a column's CHECK constraint in place. Standard
    pattern: create new table with new CHECK, copy data while translating
    values, drop old, rename new. Wrapped in IMMEDIATE transaction so a
    failure rolls back cleanly without leaving the DB half-migrated.
    """
    new_create = """
    CREATE TABLE providence_cache_new (
        cache_key                TEXT PRIMARY KEY,
        source_root              TEXT NOT NULL,
        document_uri             TEXT NOT NULL,
        question_hash            TEXT NOT NULL,
        question_text            TEXT NOT NULL,
        answer_text              TEXT NOT NULL,
        merkle_proof             TEXT NOT NULL,
        model_profile_hash       TEXT NOT NULL,
        conversation_hash        TEXT NOT NULL,
        governance_policy_hash   TEXT NOT NULL,
        schema_version           TEXT NOT NULL,
        canonicalization_version TEXT NOT NULL,
        chunking_version         TEXT NOT NULL,
        falsification_state      TEXT NOT NULL DEFAULT 'live'
                                    CHECK (falsification_state IN ('live','failed','stale','quarantined')),
        chain                    TEXT NOT NULL DEFAULT 'private'
                                    CHECK (chain IN ('private','public')),
        audit_event_hash         TEXT,
        created_at               INTEGER NOT NULL,
        last_hit_at              INTEGER,
        hit_count                INTEGER NOT NULL DEFAULT 0,
        audit_mode               TEXT NOT NULL DEFAULT 'UNGROUNDED'
                                    CHECK (audit_mode IN ('STRICT','HYBRID','UNGROUNDED')),
        n_quotes                 INTEGER NOT NULL DEFAULT 0,
        n_verified               INTEGER NOT NULL DEFAULT 0,
        unverified_quotes        TEXT,
        verifier_method          TEXT NOT NULL DEFAULT 'none'
                                    CHECK (verifier_method IN ('quote','span','entity','none'))
    )
    """
    conn.execute("BEGIN IMMEDIATE")
    try:
        conn.execute(new_create)
        conn.execute(
            "INSERT INTO providence_cache_new "
            "(cache_key, source_root, document_uri, question_hash, question_text, "
            " answer_text, merkle_proof, model_profile_hash, conversation_hash, "
            " governance_policy_hash, schema_version, canonicalization_version, "
            " chunking_version, falsification_state, chain, audit_event_hash, "
            " created_at, last_hit_at, hit_count, audit_mode, n_quotes, "
            " n_verified, unverified_quotes, verifier_method) "
            "SELECT "
            " cache_key, source_root, document_uri, question_hash, question_text, "
            " answer_text, merkle_proof, model_profile_hash, conversation_hash, "
            " governance_policy_hash, schema_version, canonicalization_version, "
            " chunking_version, falsification_state, chain, audit_event_hash, "
            " created_at, last_hit_at, hit_count, "
            " CASE WHEN audit_mode = 'VISUAL' THEN 'UNGROUNDED' ELSE audit_mode END, "
            " n_quotes, n_verified, unverified_quotes, verifier_method "
            "FROM providence_cache"
        )
        conn.execute("DROP TABLE providence_cache")
        conn.execute("ALTER TABLE providence_cache_new RENAME TO providence_cache")
        conn.execute("COMMIT")
    except Exception:
        conn.execute("ROLLBACK")
        raise


# Tables that exist in every shard with the same schema. Used to build
# cross-shard UNION views in connect_query().
_SHARDABLE_TABLES = (
    "documents",
    "chunks",
    "merkle_nodes",
    "edges",
    "derivations",
    "providence_cache",
    "audit_events",
    "falsifications",
)

# Per-table column lists for cross-shard UNION views. The `chunks` table
# is pinned explicitly because the column order matters for cross-shard
# search: chunks_fts is contentless and joins back to `chunks.chunk_id`.
# Mixing legacy (composite-PK, no chunk_id column) shards with current
# (chunk_id-keyed) shards in the same --shards-dir is unsupported — run
# the migration on legacy shards first or keep them in a separate dir.
_SHARED_COLUMNS = {
    "chunks": "chunk_id, document_root, idx, leaf_hash, content, tier",
}


def discover_shards(shards_dir: Path | str) -> list[Path]:
    """Enumerate shard DB files in `shards_dir`. Returns sorted list of paths."""
    p = Path(shards_dir)
    if not p.is_dir():
        return []
    return sorted(p.glob("*.db"))


def connect_query(
    db_path: Path | str | None = None,
    shards_dir: Path | str | None = None,
) -> sqlite3.Connection:
    """Open a read-only-style connection that surfaces ALL shards as one DB.

    If `shards_dir` is set, every `*.db` in it is ATTACHed and UNION ALL views
    are created over the standard tables so existing queries (`SELECT * FROM
    documents`) work unchanged across shards. Reads only — writes still go
    through `connect()` against a specific shard.

    If `shards_dir` is None, returns a normal `connect(db_path)` for back-compat.
    """
    if shards_dir is None:
        return connect(db_path or DEFAULT_DB_PATH)

    shard_paths = discover_shards(shards_dir)
    conn = sqlite3.connect(":memory:", isolation_level=None)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA temp_store = MEMORY")

    if not shard_paths:
        # Nothing attached; create empty placeholder tables so callers don't crash.
        conn.executescript(SCHEMA_SQL)
        return conn

    aliases: list[str] = []
    for i, sp in enumerate(shard_paths):
        alias = f"sh{i:03d}"
        conn.execute(f"ATTACH DATABASE ? AS {alias}", (str(sp.resolve()),))
        aliases.append(alias)

    # UNION ALL views over the shardable tables. Columns are listed
    # explicitly (not `SELECT *`) so a shard cluster that mixes the prior
    # composite-PK chunks layout with the newer chunk_id-keyed layout still
    # unions cleanly — the explicit list is the intersection of columns
    # present in both schema generations.
    for table in _SHARDABLE_TABLES:
        cols = _SHARED_COLUMNS.get(table, "*")
        unions = " UNION ALL ".join(
            f"SELECT {cols} FROM {a}.{table}" for a in aliases
        )
        conn.execute(f"CREATE TEMP VIEW {table} AS {unions}")

    # Stash the shard list for tools that want it.
    conn.execute(
        "CREATE TEMP TABLE _shards (shard_id TEXT, path TEXT, alias TEXT)"
    )
    conn.executemany(
        "INSERT INTO _shards (shard_id, path, alias) VALUES (?, ?, ?)",
        [(p.stem, str(p.resolve()), a) for p, a in zip(shard_paths, aliases)],
    )
    return conn


@contextmanager
def transaction(conn: sqlite3.Connection) -> Iterator[sqlite3.Connection]:
    """BEGIN IMMEDIATE / COMMIT / ROLLBACK around a block."""
    conn.execute("BEGIN IMMEDIATE")
    try:
        yield conn
    except Exception:
        conn.execute("ROLLBACK")
        raise
    else:
        conn.execute("COMMIT")


def _canonical_json(obj) -> str:
    """Stable JSON for audit hashing: sorted keys, no whitespace."""
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def get_meta(conn: sqlite3.Connection, key: str) -> str | None:
    """Read a value from the per-DB meta table; None if missing."""
    row = conn.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone()
    return row["value"] if row else None


def set_meta(conn: sqlite3.Connection, key: str, value: str) -> None:
    """Upsert a (key, value) into meta. Caller wraps in a transaction."""
    conn.execute(
        "INSERT INTO meta (key, value, updated_at) VALUES (?, ?, ?) "
        "ON CONFLICT(key) DO UPDATE SET value = excluded.value, "
        "updated_at = excluded.updated_at",
        (key, value, int(time.time())),
    )


def latest_event_hash(conn: sqlite3.Connection) -> str | None:
    """Return the last event_hash in the audit chain, or None for genesis."""
    row = conn.execute(
        "SELECT event_hash FROM audit_events ORDER BY seq DESC LIMIT 1"
    ).fetchone()
    return row["event_hash"] if row else None


def chain_audit_events(
    prev_event_hash: str | None,
    events: list[dict],
) -> tuple[list[tuple], str | None]:
    """Compute the event_hash chain for a batch in pure Python.

    Each event dict needs: `event_type`, `body` (dict), `subject_root` (str|None), `ts` (int).
    Returns (rows_for_executemany, last_event_hash). Insert with:

        executemany("INSERT INTO audit_events
                     (event_hash, prev_event_hash, event_type, subject_root,
                      body, ts) VALUES (?, ?, ?, ?, ?, ?)", rows)

    All chain SHA-256s are computed locally — zero DB round-trips per event.
    """
    import hashlib

    rows: list[tuple] = []
    prev = prev_event_hash
    for ev in events:
        body_json = _canonical_json(ev["body"])
        h = hashlib.sha256()
        if prev is not None:
            h.update(bytes.fromhex(prev))
        h.update(body_json.encode("utf-8"))
        event_hash = h.hexdigest()
        rows.append(
            (
                event_hash,
                prev,
                ev["event_type"],
                ev.get("subject_root"),
                body_json,
                ev["ts"],
            )
        )
        prev = event_hash
    return rows, prev


def append_audit(
    conn: sqlite3.Connection,
    event_type: str,
    body: dict,
    subject_root: str | None = None,
    ts: int | None = None,
) -> str:
    """Append one event to the audit chain. Returns the new event_hash (hex).

    Convenience wrapper for one-off events. Bulk inserts should use
    chain_audit_events() + executemany() for ~10x throughput on large batches.
    """
    import hashlib

    if ts is None:
        ts = int(time.time())
    prev = latest_event_hash(conn)
    body_json = _canonical_json(body)
    h = hashlib.sha256()
    if prev is not None:
        h.update(bytes.fromhex(prev))
    h.update(body_json.encode("utf-8"))
    event_hash = h.hexdigest()
    conn.execute(
        "INSERT INTO audit_events (event_hash, prev_event_hash, event_type, subject_root, body, ts) "
        "VALUES (?, ?, ?, ?, ?, ?)",
        (event_hash, prev, event_type, subject_root, body_json, ts),
    )
    return event_hash


def stats(conn: sqlite3.Connection) -> dict:
    """Quick landscape report."""
    def one(sql: str, *args) -> int:
        return conn.execute(sql, args).fetchone()[0]

    return {
        "documents_total": one("SELECT COUNT(*) FROM documents"),
        "documents_surface": one("SELECT COUNT(*) FROM documents WHERE kind='surface'"),
        "documents_core": one("SELECT COUNT(*) FROM documents WHERE kind='core'"),
        "chunks_total": one("SELECT COUNT(*) FROM chunks"),
        "chunks_hot": one("SELECT COUNT(*) FROM chunks WHERE tier='hot'"),
        "chunks_warm": one("SELECT COUNT(*) FROM chunks WHERE tier='warm'"),
        "chunks_cold": one("SELECT COUNT(*) FROM chunks WHERE tier='cold'"),
        "edges_total": one("SELECT COUNT(*) FROM edges"),
        "providence_total": one("SELECT COUNT(*) FROM providence_cache"),
        "audit_events_total": one("SELECT COUNT(*) FROM audit_events"),
    }