arborist-viz/arborist_viz/views/roots.py

"""Root explorer + leaves (ticket §6.2, §13). Reads through ``arborist.read``."""

from __future__ import annotations

from dataclasses import asdict

from pyramid.response import Response
from pyramid.view import view_config


def _json(body, status=200) -> Response:
    return Response(json_body=body, status=status)


def _no_shard() -> Response:
    return _json(
        {"error": "no_shard", "note": "configure arborist.shards or set ARBORIST_VIZ_SHARDS"},
        status=503,
    )


@view_config(route_name="api_root", request_method="GET")
def get_root(request):
    if request.shards is None:
        return _no_shard()
    root_hash = request.matchdict["root_hash"]
    r = request.shards.root(root_hash)
    if r is not None:
        return _json({
            "root_hash": r.document_root,
            "document_uri": r.document_uri,
            "title": r.title,
            "source_type": r.source_type,
            "kind": r.kind,
            "leaf_count": r.leaf_count,
            "hash_algorithm": "sha256",
            "domain_separator": "leaf=0x00 / node=0x03 (arborist convention)",
            "canonicalization_version": r.canonicalization_version,
            "schema_version": r.schema_version,
            "chunking_version": r.chunking_version,
            "ingest_ts": r.ingest_ts,
            "hit_count": r.hit_count,
            "shard": r.shard_path,
            "source_roots": [r.document_root],
        })
    # Not a documents row — fall through to the synthetic context root
    # the QA pipeline produces when an answer was assembled from chunks
    # of several documents.
    ctx = request.shards.context(root_hash)
    if ctx is None:
        return _json({"error": "not_found", "root_hash": root_hash}, status=404)
    return _json({
        "root_hash": ctx.context_root,
        "document_uri": ctx.document_uri,
        "title": (ctx.question_text[:140] + "…") if len(ctx.question_text) > 140 else ctx.question_text,
        "source_type": "context",
        "kind": "context_root",
        "leaf_count": len(ctx.sources),
        "hash_algorithm": "sha256",
        "domain_separator": "leaf=0x00 / node=0x03 (arborist convention)",
        "canonicalization_version": ctx.canonicalization_version,
        "schema_version": ctx.schema_version,
        "chunking_version": ctx.chunking_version,
        "ingest_ts": ctx.created_at,
        "hit_count": ctx.hit_count,
        "shard": ctx.shard_path,
        "source_roots": [ctx.context_root],
        "sources": ctx.sources,
        "cache_key": ctx.cache_key,
        "question_text": ctx.question_text,
    })


@view_config(route_name="api_root_leaves", request_method="GET")
def get_root_leaves(request):
    if request.shards is None:
        return _no_shard()
    root_hash = request.matchdict["root_hash"]
    leaves = request.shards.leaves(
        root_hash, reveal_private=request.reveal_private_leaves
    )
    if leaves:
        return _json({
            "root_hash": root_hash,
            "leaf_count": len(leaves),
            "private_revealed": request.reveal_private_leaves,
            "leaves": [asdict(L) for L in leaves],
            "source_roots": [root_hash],
        })
    # Synthesize leaves from a context root's sources manifest. Each
    # source is one leaf of the context tree; the leaf hash is the
    # source's document_root. We load the ACTUAL chunk text from the
    # source document (the chunk_idx that retrieval surfaced) so the
    # content panel shows the real prose the answer drew on, not just a
    # title placeholder. Falls back to title+URI when the source shard
    # isn't mounted (e.g. web.db absent) or content is private/cold.
    # Each leaf also carries its usage signals (used / used_pointer_ids /
    # source_role) so the lattice + leaves strip can highlight the chunks
    # that actually fed the answer.
    ctx = request.shards.context(root_hash)
    if ctx is None:
        return _json({"error": "not_found", "root_hash": root_hash}, status=404)
    sorted_sources = sorted(ctx.sources, key=lambda s: (s.get("document_root") or "").lower())
    syn = []
    for idx, s in enumerate(sorted_sources):
        sroot = (s.get("document_root") or "").lower()
        title = s.get("title") or s.get("document_uri") or sroot[:16]
        uri = s.get("document_uri") or ""
        chunk_idx = s.get("chunk_idx")
        content = None
        prose = None
        base_version = None
        if request.reveal_private_leaves and sroot:
            try:
                src_leaves = request.shards.leaves(sroot, reveal_private=True)
                cl = next((L for L in src_leaves if L.idx == chunk_idx), None)
                if cl is None and src_leaves:
                    cl = src_leaves[0]
                if cl is not None:
                    content = cl.content
                    prose = cl.prose
                    base_version = cl.base_version
            except Exception:
                pass
        if content is None:
            content = "{0}\n{1}".format(title, uri)
        syn.append({
            "document_root": ctx.context_root,
            "idx": idx,
            "leaf_hash": sroot,
            "tier": "context",
            "content": content,
            "prose": prose,
            "base_version": base_version,
            "source_root": sroot,
            "source_chunk_idx": chunk_idx,
            "document_uri": uri,
            "title": title,
            "used": bool(s.get("used")),
            "used_pointer_ids": s.get("used_pointer_ids") or [],
            "source_role": s.get("source_role"),
        })
    return _json({
        "root_hash": root_hash,
        "leaf_count": len(syn),
        "private_revealed": request.reveal_private_leaves,
        "leaves": syn,
        "source_roots": [root_hash],
        "kind": "context_root",
    })


@view_config(route_name="api_roots", request_method="GET")
def list_roots(request):
    if request.shards is None:
        return _no_shard()
    try:
        limit = int(request.GET.get("limit", "50"))
    except ValueError:
        limit = 50
    roots = request.shards.roots(limit=limit)
    counts = request.shards.counts()
    return _json({
        "roots": [asdict(r) for r in roots],
        "count": len(roots),
        "total_documents": counts.documents,
        "shards": request.shards.paths,
        "source_roots": [r.document_root for r in roots],
    })


@view_config(route_name="api_lattice", request_method="GET")
def get_lattice(request):
    if request.shards is None:
        return _no_shard()
    root_hash = request.matchdict["root_hash"]
    layers = request.shards.tree_layers(root_hash)
    if layers is None:
        # Fall through to the context tree synthesized from a QA's
        # sources manifest (same canonical builder, sorted source roots).
        layers = request.shards.context_layers(root_hash)
    if layers is None:
        return _json({"error": "not_found", "root_hash": root_hash}, status=404)
    return _json({
        "root_hash": root_hash,
        "layers": layers,
        "layer_count": len(layers),
        "leaf_count": len(layers[0]) if layers else 0,
        "source_roots": [root_hash],
    })


@view_config(route_name="api_resolve", request_method="GET")
def resolve_hash(request):
    """Drop-a-hash dispatcher. Returns whichever arborist object owns the hash.

    Resolution order: corpus shards first (document_root / leaf_hash /
    audit_event / merkle_interior / qa cache_key / context_root /
    run_dag_root), then sessions shard (cache_key referenced by N
    session nodes; node_hash / subtree_hash on a session node).
    """
    h = request.matchdict["hex_hash"].strip().lower()
    if not h or len(h) > 128 or not all(c in "0123456789abcdef" for c in h):
        return _json({"error": "not_hex", "hash": h}, status=400)

    # First: corpus shards (when configured).
    if request.shards is not None:
        res = request.shards.resolve(h)
        if res.kind != "unknown":
            return _json({
                "kind": res.kind, "hash": res.hash,
                "shard": res.shard_path, "url": _nav_url(res),
                **res.extra,
            })

    # Second: sessions shard. cache_key OR node_hash OR subtree_hash.
    sessions = request.sessions
    if sessions is not None and sessions.is_open:
        hits = sessions.find_by_cache_key(h)
        if hits:
            first = hits[0]
            return _json({
                "kind": "session_cache_key",
                "hash": h,
                "shard": sessions.path,
                "n_session_nodes": len(hits),
                "sids": sorted({n.sid for n in hits}),
                "bates_list": [n.bates for n in hits],
                "audit_modes": sorted({n.audit_mode for n in hits}),
                "question": first.question,
                "url": "/proof/{0}".format(h),
            })
        # node_hash / subtree_hash search: scan one row.
        # (No index on these; rare path — direct query.)
        try:
            row = sessions._conn.execute(  # type: ignore[attr-defined]
                "SELECT bates, sid, audit_mode, question, node_hash, "
                "       subtree_hash FROM nodes "
                "WHERE node_hash=? OR subtree_hash=? LIMIT 1",
                (h, h),
            ).fetchone()
        except Exception:
            row = None
        if row:
            kind = (
                "session_node_hash" if row["node_hash"] == h
                else "session_subtree_hash"
            )
            return _json({
                "kind": kind,
                "hash": h,
                "shard": sessions.path,
                "bates": row["bates"],
                "sid": row["sid"],
                "audit_mode": row["audit_mode"],
                "question": row["question"],
                "url": "/proof/{0}".format(h),
            })

    return _json({"kind": "unknown", "hash": h, "url": "/"})


def _nav_url(res) -> str:
    """Map a ResolveResult to a VIZ page.

    VIZ owns two first-class visualizers — ``/proof/{root}`` (Merkle
    lattice + leaf inclusion) and ``/audit/{event}`` (audit chain). Every
    node kind arborist surfaces resolves to one of those via its linked
    first-class hash. As arborist's recursive falsification adds new node
    kinds, this map grows by routing them to whichever existing
    visualizer is most relevant — never by minting new kind-named URLs.
    """
    if res.kind == "document_root":
        return "/proof/{0}".format(res.hash)
    if res.kind == "leaf_hash":
        return "/proof/{0}?leaf={1}".format(res.extra.get("document_root"), res.extra.get("leaf_index"))
    if res.kind == "audit_event":
        return "/audit/{0}".format(res.hash)
    if res.kind == "merkle_interior":
        return "/proof/{0}".format(res.extra.get("document_root"))
    # QA-record hashes — surface whichever root the answer cites. Carry
    # the cache_key through as a query param so the proof page can also
    # display the question + answer alongside the merkle data.
    if res.kind == "qa_cache_key":
        sr = res.extra.get("source_root")
        return "/proof/{0}?cache_key={1}".format(sr, res.hash) if sr else "/"
    if res.kind == "context_root":
        ck = res.extra.get("cache_key")
        return "/proof/{0}?cache_key={1}".format(res.hash, ck) if ck else "/proof/{0}".format(res.hash)
    if res.kind == "qa_run_dag_root":
        # No source_root in extras yet (run_dag_root only carries its
        # cache_key). Land on home until the seam enriches the extras.
        return "/"
    return "/"