"""BM25 over supplied page markdown; no benchmark answers or OCR pipeline."""

import json
import re
from pathlib import Path
from time import perf_counter

from rank_bm25 import BM25Okapi

from visual_doc_rag.data import write_json
from visual_doc_rag.retrieve import rank_scores


def tokenize(text: str) -> list[str]:
    return re.findall(r"\w+", text.lower(), flags=re.UNICODE)


class TextIndex:
    def __init__(self, page_ids: list[int], tokens: list[list[str]]):
        self.page_ids = page_ids
        self.tokens = tokens
        self.bm25 = BM25Okapi(tokens)  # library defaults: k1=1.5, b=0.75, epsilon=0.25

    def search(self, query: str) -> list[int]:
        return rank_scores(self.page_ids, self.bm25.get_scores(tokenize(query)).tolist())

    @classmethod
    def build(cls, corpus, root: Path, source: dict) -> dict:
        started = perf_counter()
        rows = corpus.remove_columns("image")
        payload = {
            "dataset": source,
            "page_ids": list(rows["corpus_id"]),
            "tokens": [tokenize(text or "") for text in rows["markdown"]],
            "tokenizer": "Unicode word regex, lowercase; no stemming/stopword removal",
            "parameters": {"k1": 1.5, "b": 0.75, "epsilon": 0.25},
        }
        # Construct now as well, so build time includes the in-memory BM25 index.
        cls(payload["page_ids"], payload["tokens"])
        write_json(root / "text.json", payload)
        result = {
            "build_seconds": perf_counter() - started,
            "index_bytes": (root / "text.json").stat().st_size,
            "pages": len(rows),
            "peak_vram_bytes": 0,
        }
        write_json(root / "text-build.json", result)
        return result

    @classmethod
    def load(cls, root: Path, source: dict):
        payload = json.loads((root / "text.json").read_text(encoding="utf-8"))
        if payload["dataset"]["revision"] != source["revision"]:
            raise ValueError("Text index dataset revision differs from the downloaded data.")
        return cls(payload["page_ids"], payload["tokens"])