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);

-- 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
);
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);

-- 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
);

-- FTS5 over chunk content for VISUAL-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)
    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


# 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"),
    }