phase1/scripts/corpus_v5.py

"""CPU runtime for an immutable schema-v5 Moonley serving release."""

from __future__ import annotations

import difflib
import json
import math
import os
import re
import sqlite3
import threading
from collections import Counter, OrderedDict, defaultdict
from pathlib import Path
from typing import Any

import numpy as np

from statute_crosswalk import load_default_crosswalk
from statute_library import ExactStatuteLibrary


QUERY_TASK = (
    "Given a legal research query, retrieve relevant passages from judgments "
    "of the Supreme Court of India that answer the query"
)
NAME_STOP = {
    "v", "vs", "of", "and", "the", "ors", "anr", "etc", "state", "union",
    "govt", "government", "in", "re", "ltd", "co", "pvt", "through", "another",
}
NAME_QUERY_NOISE = NAME_STOP | {
    "about", "case", "court", "decision", "did", "give", "held", "holding",
    "for", "is", "judgement", "judgment", "know", "me", "on", "passed", "please",
    "say", "tell", "was", "what", "which",
}
BAD_STATUS = {"overruled", "partly_overruled", "per_incuriam", "doubted"}

ACT_ALIASES = {
    "tpa": "transfer property act",
    "transfer of property act": "transfer property act",
    "transfer property act": "transfer property act",
    "ipc": "indian penal code",
    "crpc": "code criminal procedure",
    "cpc": "code civil procedure",
    "iea": "indian evidence act",
    "ni act": "negotiable instruments act",
    "bns": "bharatiya nyaya sanhita",
    "bnss": "bharatiya nagarik suraksha sanhita",
    "bsa": "bharatiya sakshya adhiniyam",
}


def _clean(value: object) -> str:
    return re.sub(r"\s+", " ", str(value or "")).strip()


def _json(value: object, default: object) -> object:
    try:
        return json.loads(str(value)) if value not in (None, "") else default
    except (TypeError, ValueError, json.JSONDecodeError):
        return default


def _ntok(value: object) -> list[str]:
    out, single = [], ""
    for token in re.findall(r"[a-z]+", str(value or "").lower()):
        if len(token) == 1:
            single += token
        else:
            if single:
                out.append(single); single = ""
            out.append(token)
    if single:
        out.append(single)
    return out


def _norm_identity(value: object) -> str:
    text = str(value or "").lower().replace("versus", " v ").replace("vs.", " v ")
    text = re.sub(r"\bvs?\b", " v ", text)
    return re.sub(r"\s+", " ", re.sub(r"[^a-z0-9]+", " ", text)).strip()


def _norm_citation(value: object) -> str:
    return re.sub(r"\s+", " ", re.sub(r"[^A-Z0-9]+", " ", str(value or "").upper())).strip()


def _norm_act(value: object) -> str:
    """Collapse common abbreviations and harmless title/year variants."""
    text = re.sub(r"\b(?:18|19|20)\d{2}\b", " ", str(value or "").lower())
    text = re.sub(r"[^a-z0-9]+", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    if text in ACT_ALIASES:
        return ACT_ALIASES[text]
    tokens = [token for token in text.split() if token not in {"the", "of", "india"}]
    normalized = " ".join(tokens)
    return ACT_ALIASES.get(normalized, normalized)


def _norm_section(value: object) -> str:
    text = re.sub(r"^\s*(?:sections?|ss?\.?)[\s:-]*", "", str(value or ""), flags=re.I)
    return re.sub(r"\s+", "", text).strip(".,;:")


class CorpusV5:
    """Expose the legacy agent tool contract over the Qwen/schema-v5 bundle.

    Every public judgment method starts at ``eligible_doc_ids``.  The release
    builder has already proved that each member has metadata, stored paragraphs,
    and authoritative search units; the runtime rechecks those counts at boot.
    """

    def __init__(
        self,
        data_dir: str | os.PathLike[str],
        statute_dir: str | os.PathLike[str] | None = None,
        device: str = "cpu",
        *,
        index: Any | None = None,
        query_encoder: Any | None = None,
    ):
        self.data_dir = Path(data_dir)
        self.device = device
        self.manifest = json.loads((self.data_dir / "release_manifest.json").read_text(encoding="utf-8"))
        if self.manifest.get("status") != "complete":
            raise RuntimeError("schema-v5 serving release is not complete")
        self.model_config = self.manifest.get("model") or {}
        self.dimension = int(self.model_config.get("dimension") or 2560)
        self._db_path = self.data_dir / str((self.manifest.get("artifacts") or {}).get("database", {}).get("name") or "corpus.sqlite3")
        if not self._db_path.exists():
            raise RuntimeError(f"serving database missing: {self._db_path}")
        self._local = threading.local()
        self._encoder_lock = threading.Lock()
        self._model_load_lock = threading.Lock()
        self._query_encoder = query_encoder
        self._query_cache: OrderedDict[str, np.ndarray] = OrderedDict()
        self._query_cache_lock = threading.Lock()
        self._query_cache_size = max(8, int(os.environ.get("THEMIS_QUERY_CACHE", "64")))

        if index is None:
            import faiss

            index_path = self.data_dir / str((self.manifest.get("artifacts") or {}).get("faiss_index", {}).get("name") or "index.faiss")
            flags = getattr(faiss, "IO_FLAG_MMAP", 0) | getattr(faiss, "IO_FLAG_READ_ONLY", 0)
            self.index = faiss.read_index(str(index_path), flags)
        else:
            self.index = index
        if int(getattr(self.index, "d", self.dimension)) != self.dimension:
            raise RuntimeError("FAISS dimension does not match release manifest")

        self.meta: dict[str, dict[str, Any]] = {}
        self.goodlaw: dict[str, dict[str, Any]] = {}
        self.decision_year: dict[str, int] = {}
        self.bench_n: dict[str, int] = {}
        self.canonical: set[str] = set()
        self.name_vocab: set[str] = set()
        self.name_postings: dict[str, set[str]] = defaultdict(set)
        self.aliases: dict[str, str] = {}
        self.nc2doc: dict[str, str] = {}
        self.cite_resolver: dict[str, str] = {}
        self._load_metadata()
        self.eligible_doc_ids = set(self.meta)
        self.canonical = set(self.eligible_doc_ids)

        self._doc_rows: dict[str, list[int]] = defaultdict(list)
        self._row_doc: list[str] = []
        self._row_type: list[str] = []
        self._unit_cache: OrderedDict[int, dict[str, Any]] = OrderedDict()
        self._load_unit_map()
        expected_units = int((self.manifest.get("corpus") or {}).get("units") or 0)
        if expected_units and (len(self._row_doc) != expected_units or int(getattr(self.index, "ntotal", expected_units)) != expected_units):
            raise RuntimeError("unit-table, manifest, and FAISS row counts diverge")
        missing = self.eligible_doc_ids - set(self._doc_rows)
        if missing:
            raise RuntimeError(f"{len(missing)} accepted judgments have no serving units")

        self.out_edges: dict[str, list[str]] = defaultdict(list)
        self.in_edges: dict[str, list[str]] = defaultdict(list)
        self.edge_meta: dict[tuple[str, str], dict[str, Any]] = {}
        self.cite_indeg: Counter[str] = Counter()
        self._load_graph()
        self.statute_idx = self._load_statute_index()
        self._provision_act_names: dict[str, list[str]] = defaultdict(list)
        for row in self._connection().execute(
            "SELECT DISTINCT act_name FROM provisions WHERE act_name IS NOT NULL AND act_name != ''"
        ):
            name = str(row["act_name"])
            self._provision_act_names[_norm_act(name)].append(name)
        self.concord = {}
        concordance = Path(statute_dir or "") / "concordance.json" if statute_dir else None
        if concordance and concordance.exists():
            self.concord = json.loads(concordance.read_text(encoding="utf-8"))
        crosswalk_path = os.environ.get("THEMIS_SECTION_CROSSWALK", "").strip() or None
        self.crosswalk = load_default_crosswalk(crosswalk_path)
        provisions_path = Path(statute_dir or "") / "all_statutes.json" if statute_dir else None
        self.statute_library = ExactStatuteLibrary.from_env(fallback_path=provisions_path)
        print(
            f"[corpus-v5] ready — {len(self.eligible_doc_ids)} accepted judgments, "
            f"{len(self._row_doc)} Qwen units, {sum(self.cite_indeg.values())} resolved internal edges",
            flush=True,
        )

    def _connection(self) -> sqlite3.Connection:
        connection = getattr(self._local, "connection", None)
        if connection is None:
            connection = sqlite3.connect(f"file:{self._db_path}?mode=ro", uri=True, check_same_thread=False, timeout=30)
            connection.row_factory = sqlite3.Row
            self._local.connection = connection
        return connection

    def _load_metadata(self) -> None:
        connection = sqlite3.connect(f"file:{self._db_path}?mode=ro", uri=True)
        connection.row_factory = sqlite3.Row
        aliases_by_doc: dict[str, list[str]] = defaultdict(list)
        alias_owners: dict[str, set[str]] = defaultdict(set)
        for row in connection.execute("SELECT alias,normalized_alias,judgment_id FROM aliases"):
            aliases_by_doc[str(row["judgment_id"])].append(str(row["alias"]))
            alias_owners[str(row["normalized_alias"])].add(str(row["judgment_id"]))
        for key, owners in alias_owners.items():
            if len(owners) == 1:
                self.aliases[key] = next(iter(owners))
        citation_owners: dict[str, set[str]] = defaultdict(set)
        for row in connection.execute("SELECT * FROM judgments"):
            d = str(row["judgment_id"])
            equivalents = list(_json(row["equivalent_citations_json"], []))
            bench = list(_json(row["bench_json"], []))
            acts_records = list(_json(row["acts_json"], []))
            provisions = list(_json(row["provisions_json"], []))
            summary = dict(_json(row["summary_json"], {}))
            graph_metrics = dict(_json(row["graph_metrics_json"], {}))
            case_numbers = list(_json(row["case_numbers_json"], []))
            case_number = next((item.get("raw") for item in case_numbers if isinstance(item, dict) and item.get("raw")), None)
            m = {
                "doc_id": d, "judgment_id": d, "case_name": row["case_name"],
                "neutral_citation": row["neutral_citation"], "equivalent_citations": equivalents,
                "date": row["decision_date"], "year": row["year"], "court": row["court"],
                "case_number": case_number, "bench_strength": row["bench_size"] or row["bench_bucket"],
                "bench": bench, "author_judge": None, "disposition": row["disposition"],
                "acts": [item.get("name") for item in acts_records if isinstance(item, dict) and item.get("name")],
                "acts_records": acts_records, "provisions": provisions, "issue": row["issue"], "held": row["held"],
                "summary": summary, "source_url": row["source_url"], "source_provider": row["source_provider"],
                "source_ik_tid": row["source_ik_tid"], "review_status": row["review_status"],
                "aliases": aliases_by_doc.get(d, []), "graph_metrics": graph_metrics,
            }
            self.meta[d] = m
            status = str(row["good_law_status"] or "unknown")
            good_law = dict(_json(row["good_law_json"], {}))
            self.goodlaw[d] = {
                **good_law, "good_law_status": status,
                "display_state": row["display_state"] or "grey",
                "treatment_breakdown": graph_metrics.get("treatment_breakdown") or {},
            }
            try:
                self.decision_year[d] = int(row["year"])
            except (TypeError, ValueError):
                pass
            self.bench_n[d] = int(row["bench_size"] or 0)
            for token in _ntok(row["case_name"]):
                if len(token) >= 4:
                    self.name_vocab.add(token)
                if len(token) > 1:
                    self.name_postings[token].add(d)
            if row["neutral_citation"]:
                self.nc2doc[str(row["neutral_citation"])] = d
            for citation in [row["neutral_citation"], *equivalents]:
                normalized = _norm_citation(citation)
                if normalized:
                    citation_owners[normalized].add(d)
        for key, owners in citation_owners.items():
            if len(owners) == 1:
                self.cite_resolver[key] = next(iter(owners))
        connection.close()

    def _load_unit_map(self) -> None:
        for row in self._connection().execute("SELECT row_id,judgment_id,unit_type FROM units ORDER BY row_id"):
            row_id = int(row["row_id"])
            if row_id != len(self._row_doc):
                raise RuntimeError("serving unit rows are not contiguous")
            judgment_id = str(row["judgment_id"])
            self._row_doc.append(judgment_id)
            self._row_type.append(str(row["unit_type"]))
            self._doc_rows[judgment_id].append(row_id)

    def _load_graph(self) -> None:
        query = "SELECT * FROM graph_edges WHERE target_id IS NOT NULL"
        for row in self._connection().execute(query):
            source, target = str(row["source_id"]), str(row["target_id"])
            if source not in self.eligible_doc_ids or target not in self.eligible_doc_ids:
                continue
            self.out_edges[source].append(target); self.in_edges[target].append(source)
            self.edge_meta[(source, target)] = {
                "treatment": row["relation"] or "referred_to", "scope": row["scope"],
                "confidence": row["confidence"], "evidence": list(_json(row["evidence_json"], [])),
                "method": row["resolution_method"],
            }
            self.cite_indeg[target] += 1

    def _load_statute_index(self) -> list[dict[str, Any]]:
        rows = self._connection().execute(
            "SELECT act_name,number,MIN(raw_mention) title,COUNT(DISTINCT judgment_id) cases "
            "FROM provisions WHERE act_name IS NOT NULL GROUP BY act_name,number ORDER BY cases DESC LIMIT 25000"
        )
        return [
            {"act_short": row["act_name"], "section_number": row["number"], "title": row["title"], "cases": row["cases"]}
            for row in rows
        ]

    def _load_encoder(self) -> Any:
        if self._query_encoder is not None:
            return self._query_encoder
        with self._model_load_lock:
            if self._query_encoder is not None:
                return self._query_encoder
            import torch
            from sentence_transformers import SentenceTransformer

            model_path = os.environ.get("THEMIS_QWEN_MODEL") or self.model_config.get("model_id") or "Qwen/Qwen3-Embedding-4B"
            dtype_name = os.environ.get("THEMIS_QWEN_DTYPE", "bfloat16").lower()
            dtype = torch.bfloat16 if dtype_name == "bfloat16" else torch.float32
            torch.set_num_threads(max(1, int(os.environ.get("THEMIS_TORCH_THREADS", str(os.cpu_count() or 4)))))
            kwargs: dict[str, Any] = {
                "device": "cpu", "trust_remote_code": True,
                "model_kwargs": {"dtype": dtype, "low_cpu_mem_usage": True},
            }
            if Path(str(model_path)).exists():
                kwargs["local_files_only"] = True
            else:
                kwargs["revision"] = self.model_config.get("revision")
            model = SentenceTransformer(str(model_path), **kwargs)
            model.max_seq_length = int(self.model_config.get("max_seq_length") or 2048)
            self._query_encoder = model
        return model

    def warmup(self) -> None:
        self._enc("Supreme Court legal research")

    def _enc(self, query: str) -> np.ndarray:
        normalized = _clean(query)
        with self._query_cache_lock:
            cached = self._query_cache.get(normalized)
            if cached is not None:
                self._query_cache.move_to_end(normalized)
                return cached.copy()
        prompt = f"Instruct: {self.model_config.get('query_task') or QUERY_TASK}\nQuery: {normalized}"
        encoder = self._load_encoder()
        with self._encoder_lock:
            if callable(encoder) and not hasattr(encoder, "encode"):
                vector = encoder(prompt)
            else:
                vector = encoder.encode(prompt, normalize_embeddings=True, convert_to_numpy=True)
        vector = np.asarray(vector, dtype=np.float32).reshape(-1)
        if vector.shape != (self.dimension,):
            raise RuntimeError(f"query encoder returned {vector.shape}; expected {(self.dimension,)}")
        vector /= max(float(np.linalg.norm(vector)), 1e-12)
        with self._query_cache_lock:
            self._query_cache[normalized] = vector.copy()
            self._query_cache.move_to_end(normalized)
            while len(self._query_cache) > self._query_cache_size:
                self._query_cache.popitem(last=False)
        return vector

    def encode_documents(self, texts: list[str]) -> np.ndarray:
        """Encode private knowledge chunks in the Qwen document space."""
        values = [_clean(text) for text in texts if _clean(text)]
        if not values:
            return np.empty((0, self.dimension), dtype=np.float32)
        encoder = self._load_encoder()
        with self._encoder_lock:
            if callable(encoder) and not hasattr(encoder, "encode"):
                matrix = np.vstack([encoder(value) for value in values])
            else:
                matrix = encoder.encode(
                    values,
                    batch_size=max(1, int(os.environ.get("THEMIS_KNOWLEDGE_BATCH", "4"))),
                    normalize_embeddings=True,
                    convert_to_numpy=True,
                    show_progress_bar=False,
                )
        matrix = np.asarray(matrix, dtype=np.float32)
        if matrix.shape != (len(values), self.dimension):
            raise RuntimeError(
                f"document encoder returned {matrix.shape}; expected {(len(values), self.dimension)}"
            )
        return matrix

    def _unit(self, row_id: int) -> dict[str, Any]:
        cached = self._unit_cache.get(int(row_id))
        if cached is not None:
            self._unit_cache.move_to_end(int(row_id)); return cached
        row = self._connection().execute("SELECT * FROM units WHERE row_id=?", (int(row_id),)).fetchone()
        if row is None:
            return {}
        unit = dict(row); unit["paragraph_ids"] = list(_json(unit.pop("paragraph_ids_json", "[]"), []))
        self._unit_cache[int(row_id)] = unit
        if len(self._unit_cache) > 4096:
            self._unit_cache.popitem(last=False)
        return unit

    def _dense_units(self, query: str, n: int = 1500) -> list[tuple[int, float]]:
        limit = max(1, min(int(n), len(self._row_doc)))
        scores, rows = self.index.search(self._enc(query).reshape(1, -1), limit)
        return [(int(row), float(score)) for row, score in zip(rows[0], scores[0]) if int(row) >= 0]

    def _card(self, judgment_id: str, rr: float = 0.0, passage: str | None = None) -> dict[str, Any]:
        d = str(judgment_id); m = self.meta.get(d, {}); gl = self.goodlaw.get(d, {})
        snippet = _clean(passage or m.get("held") or (m.get("summary") or {}).get("one_line") or (m.get("summary") or {}).get("text"))[:420]
        return {
            "doc_id": d, "judgment_id": d, "case_name": m.get("case_name"), "year": m.get("year"),
            "date": m.get("date"), "neutral_citation": m.get("neutral_citation"),
            "equivalent_citations": m.get("equivalent_citations") or [], "court": m.get("court"),
            "bench_strength": m.get("bench_strength"), "disposition": m.get("disposition"),
            "cited_by": self.cite_indeg.get(d, 0), "good_law": gl.get("good_law_status", "unknown"),
            "good_law_status": gl.get("good_law_status", "unknown"), "rr": round(float(rr), 6),
            "snippet": snippet, "passage": snippet,
        }

    def is_retrieval_eligible(self, doc_id: object) -> bool:
        return str(doc_id) in self.eligible_doc_ids

    def coverage(self) -> dict[str, Any]:
        corpus = self.manifest.get("corpus") or {}
        return {
            "accepted_judgments": len(self.eligible_doc_ids), "metadata_only_excluded": 0,
            "units": len(self._row_doc), "paragraphs": int(corpus.get("paragraphs") or 0),
            "scope": corpus.get("source_scope") or "Supreme Court of India judgments stored in this release",
            "release_version": self.manifest.get("release_version"),
        }

    def vector_search(self, query: str, k: int = 8) -> list[dict[str, Any]]:
        out, seen = [], set()
        for row_id, score in self._dense_units(query, max(800, k * 80)):
            d = self._row_doc[row_id]
            if d in seen or not self.is_retrieval_eligible(d):
                continue
            seen.add(d); out.append(self._card(d, score, self._unit(row_id).get("text")))
            if len(out) >= k:
                break
        return out

    def dense_docs(self, query: str, k: int = 60) -> list[str]:
        return [card["doc_id"] for card in self.vector_search(query, k)]

    def keyword_search(self, query: str, k: int = 12, **_: Any) -> list[dict[str, Any]]:
        stop = {"of", "the", "and", "or", "in", "to", "a", "an", "is", "for", "on", "by", "with"}
        tokens = [token for token in re.findall(r"[a-z0-9]+", query.lower()) if token not in stop]
        if not tokens:
            return []
        match = " OR ".join(f'"{token}"' for token in tokens[:24])
        rows = self._connection().execute(
            "SELECT rowid,bm25(unit_fts) score FROM unit_fts WHERE unit_fts MATCH ? ORDER BY score LIMIT ?",
            (match, max(300, k * 40)),
        ).fetchall()
        out, seen = [], set()
        for row in rows:
            row_id = int(row["rowid"]); d = self._row_doc[row_id]
            if d in seen or not self.is_retrieval_eligible(d):
                continue
            seen.add(d); out.append(self._card(d, -float(row["score"]), self._unit(row_id).get("text")))
            if len(out) >= k:
                break
        return out

    def score_docs(self, query: str, doc_ids: list[str], per_doc: int = 8) -> dict[str, float]:
        qv = self._enc(query); scores: dict[str, float] = {}
        for value in doc_ids:
            d = str(value)
            if not self.is_retrieval_eligible(d):
                continue
            rows = self._doc_rows[d][: max(1, per_doc)]
            if not rows:
                continue
            vectors = np.vstack([np.asarray(self.index.reconstruct(int(row)), dtype=np.float32) for row in rows])
            scores[d] = float(np.max(vectors @ qv))
        return scores

    def hybrid_search(self, query: str, k: int = 8, pool: int = 60) -> list[dict[str, Any]]:
        dense = self.vector_search(query, pool); keyword = self.keyword_search(query, pool)
        fused: dict[str, float] = defaultdict(float)
        cards: dict[str, dict[str, Any]] = {}
        for lane in (dense, keyword):
            for rank, card in enumerate(lane, 1):
                d = card["doc_id"]; fused[d] += 1.0 / (60 + rank); cards.setdefault(d, card)
        ranked = sorted(fused, key=lambda d: -fused[d])[: max(k * 5, 40)]
        refined = self.score_docs(query, ranked)
        ranked.sort(key=lambda d: -(refined.get(d, 0.0) + 8 * fused[d]))
        return [self._card(d, refined.get(d, fused[d]), cards[d].get("passage")) for d in ranked[:k]]

    def authority_search(self, query: str, k: int = 8, alpha: float = 0.3) -> list[dict[str, Any]]:
        cards = self.vector_search(query, max(80, k * 10))
        cards.sort(key=lambda card: -(float(card.get("rr") or 0) + alpha * math.log1p(self.cite_indeg.get(card["doc_id"], 0))))
        return cards[:k]

    def search_lanes(self, query: str, frame: dict[str, Any], lane_n: int = 6) -> dict[str, list[dict[str, Any]]]:
        """Build all protected lanes from one Qwen query encoding/index scan.

        CPU serving cannot afford to encode every LLM paraphrase independently.
        The approved lawyer query is the semantic anchor; frame variants shape
        deterministic lane ordering and FTS lookups without another 4B-model pass.
        """
        base = self.hybrid_search(query, max(24, lane_n * 4))
        factual = base[:lane_n]
        doctrine = sorted(
            base,
            key=lambda card: -(
                float(card.get("rr") or 0)
                + 0.18 * math.log1p(self.cite_indeg.get(card["doc_id"], 0))
                + 0.04 * self.bench_n.get(card["doc_id"], 0)
            ),
        )[: max(lane_n, 10)]
        seen = {card["doc_id"] for card in doctrine}
        for authority in frame.get("authorities") or []:
            for card in self.name_lookup(str(authority), 2):
                if card["doc_id"] not in seen:
                    card["named"] = True; card["rr"] = max(1.0, float(card.get("rr") or 0))
                    doctrine.append(card); seen.add(card["doc_id"])
        # A governing provision is a protected metadata route, not a bag of
        # words.  Keep a quota for every inferred section so one route (for
        # example TPA s.41) cannot bury the companion route (TPA s.43).
        statute_runs: list[list[dict[str, Any]]] = []
        for section in (frame.get("sections") or [])[:3]:
            act = str(section.get("act") or "")
            number = str(section.get("section") or "")
            provisions = [{"act": act, "section": number}]
            corresponding = self.statute_crosswalk(act, number).get("corresponding") or []
            provisions.extend(corresponding[:5])
            for provision in provisions:
                mapped_act = str(provision.get("act") or "")
                mapped_number = str(provision.get("section") or "")
                exact = self.provision_cases(mapped_act, mapped_number, max(3, lane_n), query=query)
                if not exact:
                    exact = self.keyword_search(
                        f"{mapped_act} section {mapped_number}", max(3, lane_n)
                    )
                statute_runs.append(exact)
        statute: list[dict[str, Any]] = []
        statute_seen: set[str] = set()
        for rank in range(max((len(run) for run in statute_runs), default=0)):
            for run in statute_runs:
                if rank >= len(run):
                    continue
                card = run[rank]
                if card["doc_id"] not in statute_seen:
                    statute_seen.add(card["doc_id"]); statute.append(card)
                if len(statute) >= lane_n:
                    break
            if len(statute) >= lane_n:
                break
        known = []
        for value in frame.get("known_citations") or []:
            ids, _ = self.identity_hits(str(value))
            for d in ids:
                if all(card["doc_id"] != d for card in known):
                    known.append(self._card(d))
        return {"factual": factual, "doctrine": doctrine, "statute": statute[:lane_n], "known": known[:5]}

    def provision_cases(
        self,
        act: object,
        section: object,
        k: int = 8,
        *,
        query: str | None = None,
    ) -> list[dict[str, Any]]:
        """Return judgments carrying an exact structured act/section match.

        The query embedding remains the Qwen judgment embedding.  Bare-act BGE
        vectors, when enabled, are a separate retrieval space and never enter
        this score calculation.
        """
        act_key, number = _norm_act(act), _norm_section(section)
        names = self._provision_act_names.get(act_key, [])
        if not names or not number:
            return []
        placeholders = ",".join("?" for _ in names)
        rows = self._connection().execute(
            f"SELECT judgment_id,GROUP_CONCAT(DISTINCT salience) saliences "
            f"FROM provisions WHERE act_name IN ({placeholders}) AND number=? GROUP BY judgment_id",
            (*names, number),
        ).fetchall()
        doc_ids = [str(row["judgment_id"]) for row in rows if self.is_retrieval_eligible(row["judgment_id"])]
        if not doc_ids:
            return []
        salience_by_doc = {str(row["judgment_id"]): str(row["saliences"] or "") for row in rows}
        topical = self.score_docs(query, doc_ids) if query else {}

        def rank_signals(doc_id: str) -> tuple[int, int, float]:
            metrics = self.meta.get(doc_id, {}).get("graph_metrics") or {}
            try:
                external = max(0, int(metrics.get("cited_by_count") or 0))
            except (TypeError, ValueError):
                external = 0
            saliences = {value.strip().lower() for value in salience_by_doc.get(doc_id, "").split(",")}
            salience = 2 if saliences & {"core", "ratio", "primary"} else 1 if "supporting" in saliences else 0
            blended = (
                float(topical.get(doc_id, 0.0))
                + 0.35 * math.log1p(external)
                + 0.10 * math.log1p(self.cite_indeg.get(doc_id, 0))
                + 0.025 * self.bench_n.get(doc_id, 0)
            )
            return salience, external, blended

        # Exact-section lanes lead with ratio/core cases and the authority most
        # used by later courts; query fit remains the tie-breaker within that
        # protected legal route.  The final judge still decides relevance.
        signals = {doc_id: rank_signals(doc_id) for doc_id in doc_ids}
        doc_ids.sort(key=lambda doc_id: (-signals[doc_id][0], -signals[doc_id][1], -signals[doc_id][2], str(doc_id)))
        out = []
        for doc_id in doc_ids[: max(1, int(k))]:
            card = self._card(doc_id, topical.get(doc_id, 0.0))
            card["provision_match"] = {"act": act, "section": number, "exact": True}
            card["provision_salience"] = salience_by_doc.get(doc_id, "")
            card["native_cited_by"] = signals[doc_id][1]
            card["protected"] = True
            out.append(card)
        return out

    def held_search(self, query: str, k: int = 12) -> list[str]:
        out, seen = [], set()
        for row_id, _ in self._dense_units(query, max(1600, k * 100)):
            if self._row_type[row_id] not in {"holdings_ratio", "summary_overview", "issues_facts"}:
                continue
            d = self._row_doc[row_id]
            if d not in seen:
                seen.add(d); out.append(d)
            if len(out) >= k:
                break
        return out

    def citectx_search(self, query: str, k: int = 12) -> list[str]:
        out, seen = [], set()
        for row_id, _ in self._dense_units(query, max(2500, k * 140)):
            if "citation" not in self._row_type[row_id]:
                continue
            d = self._row_doc[row_id]
            if d not in seen:
                seen.add(d); out.append(d)
            if len(out) >= k:
                break
        return out

    def identity_hits(self, query: str) -> tuple[list[str], str | None]:
        cite_match = re.search(r"\b\d{4}\s+INSC\s+\d+\b|\[\d{4}\]\s*\d+\s*S\.?C\.?R\.?\s*\d+|\(\d{4}\)\s*\d+\s*SCC\s*\d+|AIR\s+\d{4}\s+SC\s+\d+", query, re.I)
        if cite_match:
            d = self.cite_resolver.get(_norm_citation(cite_match.group(0)))
            if d:
                return [d], "citation"
        normalized = _norm_identity(query)
        if normalized in self.aliases:
            return [self.aliases[normalized]], "case name"
        candidates = [(alias, d) for alias, d in self.aliases.items() if len(alias) >= 8 and alias in normalized and len(normalized) - len(alias) <= 12]
        if candidates:
            candidates.sort(key=lambda item: (-len(item[0]), -self.cite_indeg.get(item[1], 0)))
            return [candidates[0][1]], "case name"
        # A named-case question should survive small spelling errors. Keep this
        # deliberately narrow: a short name-like remainder must match every
        # token in one corpus title, so a broad legal issue still goes through
        # normal hybrid retrieval.
        words = _ntok(query)
        explicit_lookup = bool(re.search(r"\b(?:case|judg(?:e)?ment|holding|held|decision)\b", str(query or ""), re.I))
        name_tokens = [token for token in words if token not in NAME_QUERY_NOISE and len(token) > 1]
        if 1 <= len(name_tokens) <= 5 and (explicit_lookup or len(name_tokens) >= 2):
            candidate_ids: set[str] = set()
            for token in name_tokens:
                vocabulary = [token]
                if token not in self.name_vocab and len(token) >= 5:
                    vocabulary.extend(difflib.get_close_matches(token, self.name_vocab, n=3, cutoff=0.78))
                for value in vocabulary:
                    candidate_ids.update(self.name_postings.get(value, set()))
            scored = []
            for d in candidate_ids:
                case_tokens = [token for token in _ntok(self.meta[d].get("case_name")) if token not in NAME_STOP]
                ratios = [max((difflib.SequenceMatcher(None, token, other, autojunk=False).ratio() for other in case_tokens), default=0.0) for token in name_tokens]
                if ratios and all(score >= 0.78 for score in ratios):
                    score = sum(ratios) / len(ratios)
                    scored.append((score, sum(value == 1.0 for value in ratios), self.cite_indeg.get(d, 0), d))
            scored.sort(reverse=True)
            if scored and scored[0][0] >= 0.88:
                if len(scored) > 1 and scored[1][0] >= 0.88 and scored[0][0] - scored[1][0] <= 0.015:
                    return [item[-1] for item in scored[:3]], "ambiguous case name"
                kind = "case name" if scored[0][0] >= 0.999 else "close case name"
                return [scored[0][-1]], kind
            explicit_named_case = bool(re.search(r"\b(?:case|v(?:s)?\.?|versus)\b", str(query or ""), re.I))
            if explicit_named_case and any(token in self.name_vocab for token in name_tokens):
                return [], "unresolved case name"
        return [], None

    def name_lookup(self, name: str, k: int = 4) -> list[dict[str, Any]]:
        normalized = _norm_identity(name)
        exact = self.aliases.get(normalized)
        if exact:
            return [self._card(exact)]
        raw = [token for token in _ntok(name) if token not in NAME_STOP and len(token) > 1]
        if not raw:
            return []
        expanded = list(raw)
        for token in raw:
            if token not in self.name_vocab and len(token) >= 7:
                expanded.extend(difflib.get_close_matches(token, self.name_vocab, n=2, cutoff=0.84))
        query_tokens = set(expanded)
        candidates: set[str] = set()
        for token in query_tokens:
            candidates.update(self.name_postings.get(token, set()))
        scored = []
        for d in candidates:
            case_tokens = set(_ntok(self.meta[d].get("case_name")))
            overlap = query_tokens & case_tokens
            if overlap:
                scored.append((len(overlap), -abs(len(case_tokens) - len(query_tokens)), self.cite_indeg.get(d, 0), d))
        scored.sort(reverse=True)
        return [self._card(item[-1]) for item in scored[:k]]

    def statute_search(self, query: str, k: int = 3) -> list[dict[str, Any]]:
        tokens = {token for token in re.findall(r"[a-z0-9]+", query.lower()) if len(token) > 1}
        scored = []
        for item in self.statute_idx:
            value = f"{item.get('act_short')} {item.get('section_number')} {item.get('title')}".lower()
            overlap = sum(1 for token in tokens if token in value)
            if overlap:
                scored.append((overlap, int(item.get("cases") or 0), item))
        scored.sort(key=lambda item: (-item[0], -item[1]))
        return [
            {"act": item[2].get("act_short"), "section": item[2].get("section_number"), "title": item[2].get("title"), "i": index}
            for index, item in enumerate(scored[:k])
        ]

    def cases_on_section(self, text: str, k: int = 8) -> list[dict[str, Any]]:
        return self.hybrid_search(text, k)

    def statute_crosswalk(self, code: str, section: object) -> dict[str, Any]:
        return self.crosswalk.lookup(code, section)

    def statute_provision(self, code: str, section: object) -> dict[str, Any] | None:
        """Return exact bare-act text; never substitute a semantic neighbour."""
        return self.statute_library.lookup(code, section)

    def cited_authorities(self, doc_id: str, k: int = 12) -> list[dict[str, Any]]:
        return [self._card(d) for d in dict.fromkeys(self.out_edges.get(str(doc_id), [])) if self.is_retrieval_eligible(d)][:k]

    def progeny(self, doc_id: str, k: int = 12) -> list[dict[str, Any]]:
        values = sorted(set(self.in_edges.get(str(doc_id), [])), key=lambda d: -self.cite_indeg.get(d, 0))
        return [self._card(d) for d in values if self.is_retrieval_eligible(d)][:k]

    def co_cited_cases(self, doc_id: str, k: int = 8) -> list[dict[str, Any]]:
        score: Counter[str] = Counter()
        for target in set(self.out_edges.get(str(doc_id), [])):
            for citer in self.in_edges.get(target, []):
                if citer != str(doc_id) and self.is_retrieval_eligible(citer):
                    score[citer] += 1
        return [self._card(d) for d, _ in score.most_common(k)]

    def good_law_check(self, doc_id: str) -> dict[str, Any]:
        d = str(doc_id); gl = self.goodlaw.get(d, {}); status = gl.get("good_law_status", "unknown")
        overruled_by = None
        if status in BAD_STATUS:
            for source in self.in_edges.get(d, []):
                if self.edge_meta.get((source, d), {}).get("treatment") in {"overruled", "overrules"}:
                    overruled_by = self._card(source); break
        return {"doc_id": d, "good_law": status, "treatment_breakdown": gl.get("treatment_breakdown", {}), "overruled_by": overruled_by}

    def metadata_filter(self, cards: list[dict[str, Any]], min_bench: int | None = None, year_from: int | None = None, year_to: int | None = None) -> list[dict[str, Any]]:
        out = []
        for card in cards:
            d = card["doc_id"]; bench, year = self.bench_n.get(d, 0), self.decision_year.get(d, 0)
            if min_bench and bench < min_bench or year_from and year and year < year_from or year_to and year and year > year_to:
                continue
            out.append(card)
        return out

    def read_case(self, doc_id: str) -> dict[str, Any]:
        d = str(doc_id)
        if not self.is_retrieval_eligible(d):
            return {}
        m = self.meta[d]
        return {"doc_id": d, "case_name": m.get("case_name"), "neutral_citation": m.get("neutral_citation"), "bench_strength": m.get("bench_strength"), "good_law": self.goodlaw[d].get("good_law_status", "unknown"), "issue": _clean(m.get("issue"))[:1600], "held": _clean(m.get("held") or (m.get("summary") or {}).get("text"))[:2600]}

    def front_text(self, doc_id: str, n: int = 1800) -> str:
        m = self.meta.get(str(doc_id), {}); summary = m.get("summary") or {}
        return _clean(m.get("held") or summary.get("text") or summary.get("one_line"))[:n]

    def _rank_doc_rows(self, query: str, doc_id: str, limit: int = 6) -> list[tuple[int, float]]:
        rows = self._doc_rows.get(str(doc_id), [])
        if not rows:
            return []
        qv = self._enc(query)
        vectors = np.vstack([np.asarray(self.index.reconstruct(int(row)), dtype=np.float32) for row in rows])
        scores = vectors @ qv
        order = np.argsort(-scores)[:limit]
        return [(rows[int(i)], float(scores[int(i)])) for i in order]

    def best_chunk_text(self, query: str, doc_id: str, limit: int = 1600) -> str:
        ranked = self._rank_doc_rows(query, str(doc_id), 1)
        return _clean(self._unit(ranked[0][0]).get("text"))[:limit] if ranked else ""

    def full_text_for_read(self, query: str, doc_id: str, cap_chars: int = 90000) -> str:
        d = str(doc_id)
        rows = self._connection().execute("SELECT text FROM paragraphs WHERE judgment_id=? ORDER BY sequence", (d,)).fetchall()
        full = "\n".join(str(row["text"]) for row in rows)
        if len(full) <= cap_chars:
            return full
        relevant = "\n".join(self._unit(row_id).get("text", "") for row_id, _ in self._rank_doc_rows(query, d, 5))
        return (self.front_text(d, 6000) + "\n[...]\n" + relevant + "\n[...]\n" + "\n".join(str(row["text"]) for row in rows[-8:]))[:cap_chars]

    def judgment_paragraphs(self, doc_id: str, offset: int = 0, limit: int = 100) -> dict[str, Any]:
        d = str(doc_id); offset, limit = max(0, int(offset)), max(1, min(int(limit), 2000))
        total = int(self._connection().execute("SELECT COUNT(*) FROM paragraphs WHERE judgment_id=?", (d,)).fetchone()[0])
        rows = self._connection().execute(
            "SELECT * FROM paragraphs WHERE judgment_id=? ORDER BY sequence LIMIT ? OFFSET ?", (d, limit, offset)
        ).fetchall()
        paragraphs = [
            {
                "paragraph_id": row["paragraph_id"], "sequence": row["sequence"],
                "paragraph_number": row["paragraph_number"], "page_number": row["page_number"],
                "label": (f"¶ {row['paragraph_number']}" if row["paragraph_number"] else None) or row["citation_label"] or f"¶ {row['sequence']}",
                "coordinate_status": row["coordinate_status"], "text": row["text"],
                "html_anchor": "paragraph-" + re.sub(r"[^A-Za-z0-9_-]", "-", str(row["paragraph_id"])),
                "source_kind": "stored_paragraph",
            }
            for row in rows
        ]
        return {"judgment_id": d, "paragraphs": paragraphs, "offset": offset, "limit": limit, "total": total, "next_offset": offset + len(paragraphs) if offset + len(paragraphs) < total else None}

    def judgment_view(self, doc_id: str) -> dict[str, Any]:
        d = str(doc_id)
        if not self.is_retrieval_eligible(d):
            return {}
        m = self.meta[d]; page = self.judgment_paragraphs(d, 0, 2000); paragraphs = page["paragraphs"]
        text = "\n\n".join(f"{p.get('paragraph_number') or p['sequence']}. {p['text']}" for p in paragraphs)
        return {
            "doc_id": d, "judgment_id": d, "summary": m.get("summary") or {},
            "case_name": m.get("case_name"), "neutral_citation": m.get("neutral_citation"),
            "equivalent_citations": m.get("equivalent_citations") or [], "court": m.get("court"),
            "date": m.get("date"), "bench_strength": m.get("bench_strength"), "disposition": m.get("disposition"),
            "good_law_status": self.goodlaw[d].get("good_law_status", "unknown"),
            "treatment_breakdown": self.goodlaw[d].get("treatment_breakdown", {}),
            "cited_by": self.cite_indeg.get(d, 0), "issue": _clean(m.get("issue"))[:8000],
            "held": _clean(m.get("held"))[:10000], "text": text, "paragraphs": paragraphs,
            "paragraph_count": page["total"], "paragraphs_truncated": page["next_offset"] is not None,
            "source_url": m.get("source_url"), "source_provider": m.get("source_provider"),
        }

    def case_chat_passages(self, query: str, doc_id: str, k: int = 5, limit: int = 1800) -> list[dict[str, Any]]:
        d = str(doc_id)
        if not self.is_retrieval_eligible(d):
            return []
        paragraph_ids: list[str] = []
        for row_id, _ in self._rank_doc_rows(query, d, max(8, k * 2)):
            for paragraph_id in self._unit(row_id).get("paragraph_ids") or []:
                if paragraph_id not in paragraph_ids:
                    paragraph_ids.append(str(paragraph_id))
                if len(paragraph_ids) >= k:
                    break
            if len(paragraph_ids) >= k:
                break
        out = []
        for paragraph_id in paragraph_ids:
            row = self._connection().execute("SELECT * FROM paragraphs WHERE paragraph_id=? AND judgment_id=?", (paragraph_id, d)).fetchone()
            if row is None:
                continue
            out.append({
                "paragraph_id": paragraph_id, "label": (f"¶ {row['paragraph_number']}" if row["paragraph_number"] else None) or row["citation_label"] or f"¶ {row['sequence']}",
                "text": _clean(row["text"])[:limit], "source_kind": "stored_paragraph",
                "sequence": row["sequence"],
                "html_anchor": "paragraph-" + re.sub(r"[^A-Za-z0-9_-]", "-", paragraph_id),
            })
        return out

    def relevant_passages(self, query: str, doc_id: str, k: int = 6) -> list[dict[str, Any]]:
        """Case-local semantic pinpoints for the query, resolved to stored paragraphs."""
        return [
            {**item, "highlight_kind": "query_relevance"}
            for item in self.case_chat_passages(query, doc_id, k=k, limit=6000)
        ]


__all__ = ["CorpusV5"]