"""Hybrid retrieval: vector search + BM25, fused with RRF, optionally re-ranked."""
import logging

from . import store
from .bm25 import bm25
from .config import settings

log = logging.getLogger("rag.retrieval")


def rrf(result_lists: list[list[dict]], k: int = 60) -> list[dict]:
    """Reciprocal Rank Fusion: merge several ranked lists into one."""
    scores: dict[str, float] = {}
    items: dict[str, dict] = {}
    for lst in result_lists:
        for rank, hit in enumerate(lst):
            scores[hit["id"]] = scores.get(hit["id"], 0.0) + 1.0 / (k + rank + 1)
            items.setdefault(hit["id"], hit)
    fused = sorted(items.values(), key=lambda h: -scores[h["id"]])
    return [{**h, "score": scores[h["id"]]} for h in fused]


def _interleave(groups: list[list[dict]], limit: int) -> list[dict]:
    """Round-robin across documents so one big PDF cannot crowd out the others."""
    out: list[dict] = []
    i = 0
    while len(out) < limit and any(i < len(g) for g in groups):
        for g in groups:
            if i < len(g) and len(out) < limit:
                out.append(g[i])
        i += 1
    return out


def hybrid_search(query: str, doc_ids: list[str], pages: list[int] | None = None, top_k: int | None = None) -> list[dict]:
    top_k = top_k or settings.top_k
    vec = store.embed([query])[0]

    if 1 < len(doc_ids) <= settings.max_doc_groups:
        groups = [[d] for d in doc_ids]  # search each document on its own
        per_group = settings.per_doc_k
    else:
        groups = [list(doc_ids)]
        per_group = settings.candidate_k

    ranked_groups: list[list[dict]] = []
    for g in groups:
        dense = store.search(vec, store.make_filter(g, pages), per_group)
        sparse = bm25.search(query, g, pages, per_group)
        dense_score = {h["id"]: h["score"] for h in dense}
        fused = rrf([dense, sparse])[:per_group]
        # keep the raw cosine similarity: it is comparable across documents, RRF scores are not
        ranked_groups.append([{**h, "dense": dense_score.get(h["id"], 0.0)} for h in fused])

    candidates = [h for grp in ranked_groups for h in grp]
    if not candidates:
        return []

    reranker = store.get_reranker()
    if reranker is not None and len(candidates) > 1:
        scores = reranker.predict([(query, c["text"]) for c in candidates], batch_size=16, show_progress_bar=False)
        ranked = sorted(zip(candidates, scores), key=lambda x: -float(x[1]))
        return [{**c, "score": float(s)} for c, s in ranked[:top_k]]

    # no re-ranker: put the documents with the best semantic match first, then round-robin
    ranked_groups = [g for g in ranked_groups if g]
    ranked_groups.sort(key=lambda g: -max(h["dense"] for h in g))
    return _interleave(ranked_groups, top_k)


def fetch_pages(doc_ids: list[str], pages: list[int]) -> list[dict]:
    """Every chunk on the given pages (exact lookup, no similarity search)."""
    chunks = store.scroll_all(store.make_filter(doc_ids, pages))
    chunks.sort(key=lambda c: (c["filename"], c["page"], c.get("chunk_index", 0)))
    return chunks