"""Deterministic ranking and reciprocal-rank fusion."""

import json
from pathlib import Path


def rank_scores(page_ids: list[int], scores: list[float]) -> list[int]:
    if len(page_ids) != len(scores) or len(set(page_ids)) != len(page_ids):
        raise ValueError("Scores must align with unique page IDs.")
    return [
        page for page, _ in sorted(zip(page_ids, scores, strict=True), key=lambda x: (-x[1], x[0]))
    ]


def rrf(rankings: list[list[int]], k: int = 60) -> list[int]:
    """Fuse full rankings, with one-based ranks and page-ID tie breaking."""
    if k < 1:
        raise ValueError("RRF k must be positive.")
    scores: dict[int, float] = {}
    for ranking in rankings:
        if len(set(ranking)) != len(ranking):
            raise ValueError("Each ranking must contain unique page IDs.")
        for rank, page in enumerate(ranking, start=1):
            scores[page] = scores.get(page, 0.0) + 1.0 / (k + rank)
    return rank_scores(list(scores), list(scores.values()))


class PageRetriever:
    """Keep local data, indexes and the CPU encoder warm across queries."""

    def __init__(self, subset: str = "hr", data_dir: Path = Path("data"), query_encoder=None):
        from visual_doc_rag.data import load_tables
        from visual_doc_rag.pages import load_pages

        demo_manifest = data_dir / subset / "demo.json"
        self.sample_rankings = {}
        self.query_encoder = query_encoder
        if demo_manifest.exists():
            from datasets import Dataset

            payload = json.loads(demo_manifest.read_text(encoding="utf-8"))
            if payload["format_version"] != 1:
                raise ValueError("Unsupported demo bundle format.")
            self.source = payload["dataset"]
            self.sample_rankings = payload.get("sample_rankings", {})
            corpus = json.loads((demo_manifest.parent / "corpus.json").read_text(encoding="utf-8"))
            self.tables = {"corpus": Dataset.from_list(corpus)}
        else:
            self.source, self.tables = load_tables(subset, data_dir)
        self.page_root, manifest = load_pages(subset, data_dir, self.source["revision"])
        self.lookup = {page["corpus_id"]: page for page in manifest["pages"]}
        self.root = data_dir / subset / "indexes"
        self.data_dir = data_dir
        self.text = self.visual = self.encoder = None

    def search(self, query: str, method: str = "hybrid", top_k: int = 3) -> list[dict]:
        if method not in {"text", "visual", "hybrid"} or not query.strip() or top_k < 1:
            raise ValueError("Provide a nonempty query, a valid method and positive top-k.")
        sample = self.sample_rankings.get(" ".join(query.split()))
        if method == "hybrid" and sample and top_k <= len(sample):
            return self._results(sample[:top_k])
        rankings = []
        if method in {"text", "hybrid"}:
            from visual_doc_rag.text_index import TextIndex

            if self.text is None:
                self.text = TextIndex.load(self.root, self.source)
            rankings.append(self.text.search(query))
        if method in {"visual", "hybrid"}:
            import torch

            from visual_doc_rag.visual_index import Encoder, VisualIndex

            torch.set_num_threads(4)
            if self.visual is None:
                self.visual = VisualIndex(self.root, self.source)
            if self.query_encoder is not None:
                embedding = self.query_encoder(query)
            else:
                if self.encoder is None:
                    self.encoder = Encoder(
                        self.data_dir, self.visual.manifest["signature"]["model"], "cpu"
                    )
                embedding = self.encoder.encode(query=query)
            rankings.append(self.visual.search(embedding))
        ranking = rrf(rankings) if method == "hybrid" else rankings[0]
        return self._results(ranking[:top_k])

    def _results(self, ranking):
        return [
            {
                "rank": rank,
                **self.lookup[page],
                "image_path": str(self.page_root / self.lookup[page]["file"]),
            }
            for rank, page in enumerate(ranking, start=1)
        ]