"""Small retrieval pipeline: upsert, tombstone, ACL, hybrid rank.

Vectors are feature hashes of tokens. That is a stand-in so the pipeline runs
with no model download. Swap `embed` for a real embedding function without
changing ids, ACL filters, or rank fusion.
"""

from __future__ import annotations

import hashlib
import math
import re

DIM = 64
RRF_K = 60
TOKEN = re.compile(r"[a-z0-9]+")


def tokenize(text: str) -> list[str]:
    return TOKEN.findall(text.lower())


def embed(text: str) -> list[float]:
    vec = [0.0] * DIM
    for tok in tokenize(text):
        digest = hashlib.sha256(tok.encode()).digest()
        idx = int.from_bytes(digest[:2], "big") % DIM
        sign = 1.0 if digest[2] % 2 == 0 else -1.0
        vec[idx] += sign
    norm = math.sqrt(sum(v * v for v in vec)) or 1.0
    return [v / norm for v in vec]


def _cos(a, b) -> float:
    return sum(x * y for x, y in zip(a, b))


class Index:
    def __init__(self):
        self.chunks: dict[str, dict] = {}

    def upsert(self, *, chunk_id: str, doc_id: str, text: str, acl: list[str], version: int):
        if not acl:
            raise ValueError("acl must name at least one group")
        self.chunks[chunk_id] = {
            "chunk_id": chunk_id,
            "doc_id": doc_id,
            "text": text,
            "acl": list(acl),
            "version": version,
            "vector": embed(text),
            "tombstoned": False,
        }

    def tombstone_doc(self, doc_id: str, *, below_version: int | None = None):
        for row in self.chunks.values():
            if row["doc_id"] != doc_id:
                continue
            if below_version is None or row["version"] < below_version:
                row["tombstoned"] = True

    def query(self, text: str, *, groups: list[str], k: int = 5) -> list[dict]:
        allowed = set(groups)
        live = [
            row
            for row in self.chunks.values()
            if not row["tombstoned"] and allowed.intersection(row["acl"])
        ]
        if not live:
            return []
        qv = embed(text)
        q_tokens = set(tokenize(text))
        by_vec = sorted(live, key=lambda r: _cos(qv, r["vector"]), reverse=True)
        def kw(row):
            toks = set(tokenize(row["text"]))
            if not q_tokens:
                return 0
            return len(q_tokens & toks) / len(q_tokens)

        by_kw = sorted(live, key=kw, reverse=True)
        scores: dict[str, float] = {}
        for rank, row in enumerate(by_vec, start=1):
            scores[row["chunk_id"]] = scores.get(row["chunk_id"], 0.0) + 1.0 / (RRF_K + rank)
        for rank, row in enumerate(by_kw, start=1):
            scores[row["chunk_id"]] = scores.get(row["chunk_id"], 0.0) + 1.0 / (RRF_K + rank)
        # Equal RRF scores are not a coin flip. chunk_id is the written tie-break.
        # Insertion order would prefer the vector list and would change if that
        # sort changed. Every live chunk is scored on both lists; there is no
        # top-k cutoff on either leg in this file.
        ordered = sorted(scores, key=lambda cid: (-scores[cid], cid))[:k]
        out = []
        for cid in ordered:
            row = self.chunks[cid]
            out.append(
                {
                    "chunk_id": cid,
                    "doc_id": row["doc_id"],
                    "text": row["text"],
                    "version": row["version"],
                    "score": scores[cid],
                }
            )
        return out