"""BM25 sparse retrieval over visit chunks.

Text is tokenized with nltk word_tokenize (split on whitespace when nltk is not available). Each visit is
chunked on its own into `chunk_size`-token chunks, and continuation chunks repeat the visit header. The top
chunks form the MEMORY block.
"""
from __future__ import annotations

import time

from .base import BaseAgent, MemoryBuildResult, MemoryContext


class BM25RAGAgent(BaseAgent):
    def __init__(self, *, top_k: int, k1: float, b: float, chunk_size: int, chunk_overlap: int,
                 max_context_tokens: int, **params):
        super().__init__(**params)
        self.top_k = top_k
        self.k1 = k1
        self.b = b
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.max_context_tokens = max_context_tokens
        self._chunks: list[str] = []
        self._chunk_visits: list[int] = []
        self._tokenized: list[list[str]] = []
        self._bm25 = None

    @staticmethod
    def _words(text: str) -> list[str]:
        try:
            from nltk.tokenize import word_tokenize
            return word_tokenize(text.lower())
        except (ImportError, LookupError):
            return text.lower().split()

    def memorize(self, text: str, visit_number: int, date: str) -> MemoryBuildResult:
        t0 = time.time()
        header = text.splitlines()[0] if text else ""
        new_chunks = self.split_tokens(text, self.chunk_size, self.chunk_overlap, header=header)
        self._chunks.extend(new_chunks)
        self._chunk_visits.extend([visit_number] * len(new_chunks))
        self._tokenized.extend(self._words(c) for c in new_chunks)
        from rank_bm25 import BM25Okapi
        self._bm25 = BM25Okapi(self._tokenized, k1=self.k1, b=self.b)
        return MemoryBuildResult(success=True, method="bm25_rag", action="build_bm25_index",
                                 visit_number=visit_number, chunk_count=len(self._chunks),
                                 added_chunks=len(new_chunks), time_cost=time.time() - t0)

    def _rank(self, query: str) -> list[int]:
        if self._bm25 is None or not self._chunks:
            return []
        scores = self._bm25.get_scores(self._words(query))
        top = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:self.top_k]
        return [i for i in top if scores[i] > 0]

    def build_memory_context(self, question: str,
                             as_of: str = "", display_question: str | None = None) -> MemoryContext:
        idxs = self._rank(question)
        kept, used = [], 0
        for i in idxs:
            n = self.count_tokens(self._chunks[i])
            if used + n > self.max_context_tokens:
                break
            kept.append(i)
            used += n
        mem = "\n\n".join(f"[Memory {j + 1}]\n{self._chunks[i]}" for j, i in enumerate(kept))
        visits = [self._chunk_visits[i] for i in kept]
        return MemoryContext(
            mem=mem, n_tok=self.count_tokens(mem), retrieved_count=len(kept),
            retrieved_visits=sorted(set(visits)),
            retrieved_memories=[{"visit": self._chunk_visits[i],
                                 "memory": self._chunks[i][:300]} for i in kept],
            context_truncated=len(kept) < len(idxs),
            extra={"original_retrieved_count": len(idxs), "k1": self.k1, "b": self.b},
        )