"""Dense retrieval over visit chunks with a FAISS inner-product index.

New chunks are embedded at each memorize(); chunking and the MEMORY block follow bm25_rag.
"""
from __future__ import annotations

import time

import numpy as np

from .base import BaseAgent, MemoryBuildResult, MemoryContext


class DenseRAGAgent(BaseAgent):
    def __init__(self, embedder=None, *, top_k: int, chunk_size: int, chunk_overlap: int,
                 max_context_tokens: int, **params):
        super().__init__(**params)
        if embedder is None:
            raise ValueError("DenseRAGAgent requires an embedder (llm_client.make_embedder())")
        self.embedder = embedder
        self.top_k = top_k
        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._index = None

    @staticmethod
    def _normalize(vecs: list[list[float]]) -> np.ndarray:
        a = np.asarray(vecs, dtype="float32")
        norms = np.linalg.norm(a, axis=1, keepdims=True)
        norms[norms == 0] = 1.0
        return a / norms

    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)
        if new_chunks:
            vecs = self._normalize(self.embedder.embed(new_chunks))
            if self._index is None:
                import faiss
                self._index = faiss.IndexFlatIP(vecs.shape[1])
            self._index.add(vecs)
            self._chunks.extend(new_chunks)
            self._chunk_visits.extend([visit_number] * len(new_chunks))
        return MemoryBuildResult(success=True, method="dense_rag", action="add_to_faiss",
                                 visit_number=visit_number, chunk_count=len(self._chunks),
                                 added_chunks=len(new_chunks), time_cost=time.time() - t0,
                                 extra={"embedding_model": self.embedder.model})

    def _rank(self, query: str) -> list[int]:
        if self._index is None or not self._chunks:
            return []
        qv = self._normalize(self.embedder.embed([query]))
        k = min(self.top_k, len(self._chunks))
        _, idxs = self._index.search(qv, k)
        return [int(i) for i in idxs[0] if 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),
                   "embedding_model": self.embedder.model},
        )