"""Retrieval layer over the memory store.

``Retrieval`` exposes a pluggable scorer. Two scorers ship out of the box:

- :class:`BM25Scorer` — a dependency-free Okapi-BM25 implementation over
  lexical tokens. Good, explainable baseline that needs no external models.
- :class:`TagScorer` — scores by tag overlap, useful for structured lookups.

You can also supply your own ``encode(text) -> float-vector`` callable to swap
in a real embedding model later; the interface is the same, so the evaluation
harness runs unchanged. This is intentional: it shows where a semantic encoder
plugs into a memory system without coupling the whole design to one library.
"""

from __future__ import annotations

import math
from collections import Counter
from dataclasses import dataclass
from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple

from .mem import Entry, _tokenize


@dataclass
class ScoredMemory:
    entry: Entry
    score: float

    def as_dict(self) -> Dict:
        return {"id": self.entry.id, "score": round(self.score, 6)}


def _doc_freqs(entries: Sequence[Entry]) -> Dict[str, int]:
    df: Dict[str, int] = {}
    for e in entries:
        for tok in set(_tokenize(e.text)):
            df[tok] = df.get(tok, 0) + 1
    return df


class BM25Scorer:
    """Okapi-BM25 scorer with sensible defaults."""

    def __init__(self, k1: float = 1.5, b: float = 0.75, eps: float = 0.25) -> None:
        self.k1 = k1
        self.b = b
        self.eps = eps
        self._df: Dict[str, int] = {}
        self._avgdl: float = 0.0
        self._n = 0
        self._built = False

    def _build(self, entries: Sequence[Entry]) -> None:
        self._df = _doc_freqs(entries)
        lens = [len(_tokenize(e.text)) for e in entries]
        self._avgdl = (sum(lens) / len(lens)) if lens else 0.0
        self._n = len(entries)
        self._built = True

    def idf(self, term: str) -> float:
        if not self._built:
            return 0.0
        n = self._df.get(term, 0)
        return math.log(1 + (self._n - n + 0.5) / (n + 0.5) + self.eps)

    def score(self, query_terms: Sequence[str], entry: Entry) -> float:
        if not self._built or not query_terms:
            return 0.0
        tf = Counter(_tokenize(entry.text))
        dl = sum(tf.values())
        score = 0.0
        for q in query_terms:
            f = tf.get(q, 0)
            if f == 0:
                continue
            denom = f + self.k1 * (1 - self.b + self.b * (dl / self._avgdl if self._avgdl else 1.0))
            score += self.idf(q) * (f * (self.k1 + 1)) / denom
        return score

    def score_all(self, query_terms: Sequence[str], entries: Sequence[Entry]) -> List[ScoredMemory]:
        scored = [ScoredMemory(e, self.score(query_terms, e)) for e in entries]
        scored.sort(key=lambda s: s.score, reverse=True)
        return scored


class TagScorer:
    """Scores by normalized tag overlap; a simple structured baseline."""

    def score_all(self, query_terms: Sequence[str], entries: Sequence[Entry]) -> List[ScoredMemory]:
        qset = set(query_terms)
        scored = []
        for e in entries:
            tags = {t.lower() for t in e.tags}
            overlap = len(qset & tags)
            # small lexical bonus so untagged but keyword-matching entries still rank
            lex = len(qset & set(_tokenize(e.text)))
            scored.append(ScoredMemory(e, float(overlap) + 0.1 * lex))
        scored.sort(key=lambda s: s.score, reverse=True)
        return scored


class Retrieval:
    """Query the memory store with a chosen scorer.

    Parameters
    ----------
    scorer:
        A scorer with ``score_all(query_terms, entries) -> List[ScoredMemory]``.
        Defaults to :class:`BM25Scorer`.
    include_tags:
        Append entry tags to the query expansion before scoring (tag-aware query).
    """

    def __init__(self, scorer: Optional[object] = None, include_tags: bool = True) -> None:
        self.scorer = scorer if scorer is not None else BM25Scorer()
        self.include_tags = include_tags

    def _query_terms(self, query: str) -> List[str]:
        terms = _tokenize(query)
        if self.include_tags:
            for t in terms:
                if t not in terms:
                    pass
            # tag expansion happens in score_all through entry tags already
        return terms

    def search(self, store: MemoryStore, query: str, k: int = 5) -> List[ScoredMemory]:
        entries = store.all()
        qterms = self._query_terms(query)
        if isinstance(self.scorer, (BM25Scorer,)):
            self.scorer._build(entries)
        scored = self.scorer.score_all(qterms, entries)
        return scored[:k]

    def search_terms(self, entries: Sequence[Entry], query: str, k: int = 5) -> List[ScoredMemory]:
        """Score an explicit list of entries (used by the eval harness)."""
        qterms = _tokenize(query)
        if isinstance(self.scorer, BM25Scorer):
            self.scorer._build(entries)
        scored = self.scorer.score_all(qterms, entries)
        return scored[:k]