app/engine/search.py

"""
search.py — the retrieval orchestrator.

Pipeline per query:
  parse (SLM/rules) -> BM25+ lexical top-K  ∥  bge/FAISS dense top-K
    -> fusion (RRF | weighted min-max) -> intent-aware soft boosts
    -> cross-encoder rerank of the shortlist -> grounded "why" per result.

All artifacts are precomputed by scripts/build_index.py; this module just loads
them and serves.
"""
from __future__ import annotations

import math
import pickle
import re
import time

import numpy as np

from app.config import Settings
from app.engine.corpus import Corpus
from app.engine.fusion import Fused, fuse
from app.engine.lexical import LexicalIndex
from app.engine.preprocess import preprocessor
from app.engine.rerank import Reranker
from app.engine.semantic import SemanticIndex
from app.engine.slm import GENRE_CANON, SLM
from app.schemas import (
    EngineInfo,
    GenreCount,
    HealthResponse,
    Movie,
    ParsedQuery,
    ScoreBreakdown,
    SearchRequest,
    SearchResponse,
    SearchResult,
    StatsResponse,
    Suggestion,
    Timing,
)

_SAMPLE_QUERIES = [
    "mind-bending sci-fi like Inception but with more heart",
    "cozy 90s comedies for a rainy sunday",
    "slow-burn neo-noir crime thrillers",
    "movies that feel like a fever dream",
    "gritty revenge thrillers with a female lead",
    "feel-good found-family space adventures",
]


def _norm(s: str) -> str:
    return re.sub(r"[^a-z0-9]", "", str(s).lower())


# Meta words the parser consumes into `min_rating` — they must not leak into the
# BM25 query as content terms (e.g. "rated" lemmatizes to "rat" and would match
# a movie about a rat). Filtered only when a rating intent was actually parsed.
_RATING_WORDS = {"highly", "rated", "rate", "acclaimed", "critically",
                 "masterpiece", "toprated"}


def _minmax_list(vals: list[float]) -> list[float]:
    if not vals:
        return []
    lo, hi = min(vals), max(vals)
    if hi - lo < 1e-12:
        return [1.0 if hi > 0 else 0.0 for _ in vals]
    return [(v - lo) / (hi - lo) for v in vals]


class SearchEngine:
    def __init__(self, settings: Settings):
        self.s = settings
        self.corpus = Corpus(settings.corpus_path)

        # ---- lexical: load precomputed tokenized corpus, build BM25+ ----
        tokenized = self._load_tokens()
        self.lexical = LexicalIndex(
            tokenized, settings.bm25_k1, settings.bm25_b, settings.bm25_delta
        )

        # ---- semantic: load embeddings, build FAISS flat index ----
        embeddings = np.load(settings.embeddings_path)
        self.semantic = SemanticIndex(settings.embedding_model, embeddings)

        # ---- reranker + SLM ----
        self.reranker = Reranker(settings.reranker_model) if settings.enable_reranker else None
        self.slm = SLM(settings.ollama_host, settings.ollama_model,
                       settings.enable_slm, settings.slm_timeout)

        self.title_index = {_norm(r["title"]): i for i, r in enumerate(self.corpus.records)}
        self._ready = False

    # --------------------------------------------------------------------- #
    def _load_tokens(self) -> list[list[str]]:
        tok_path = self.s.embeddings_path.parent / "lexical_tokens.pkl"
        if tok_path.exists():
            with open(tok_path, "rb") as f:
                return pickle.load(f)
        # fallback: tokenize on the fly (slower cold start)
        return [preprocessor.tokenize(t) for t in self.corpus.lexical_texts()]

    def warmup(self) -> None:
        self.semantic.warmup()
        if self.reranker:
            self.reranker.warmup()
        self._ready = True

    def _resolve_similar(self, title: str | None) -> int | None:
        if not title:
            return None
        key = _norm(title)
        if key in self.title_index:
            return self.title_index[key]
        # prefix / containment fallback
        for k, idx in self.title_index.items():
            if key and (k.startswith(key) or key.startswith(k)) and abs(len(k) - len(key)) <= 4:
                return idx
        return None

    def _boost(self, parsed: ParsedQuery, rec: dict) -> float:
        """Intent-aware additive score adjustment in ~[-0.3, +0.25]."""
        adj = 0.0
        genres = set(rec.get("genres") or [])
        low_g = {g.lower() for g in genres}
        kws = {k.lower() for k in (rec.get("keywords") or [])}

        gmatch = len(set(parsed.genres) & genres)
        adj += min(0.12, 0.07 * gmatch)

        if parsed.era_from or parsed.era_to:
            y = rec.get("year")
            if y:
                lo = parsed.era_from or 0
                hi = parsed.era_to or 3000
                adj += 0.10 if lo <= y <= hi else -0.13

        if parsed.min_rating:
            r = rec.get("vote_average")
            if r is not None:
                adj += 0.06 if r >= parsed.min_rating else -0.40

        for neg in parsed.negations:
            n = neg.lower()
            # negated *genre* is handled by a hard filter upstream; here we also
            # penalize soft matches in keywords / plain adjectives.
            if n in kws or (n in low_g and GENRE_CANON.get(n) not in genres):
                adj -= 0.25
        return max(-0.6, min(0.25, adj))

    # --------------------------------------------------------------------- #
    def _rank_candidates(
        self, query_for_semantic: str, lex_tokens: list[str], parsed: ParsedQuery,
        req: SearchRequest, timing: Timing, exclude: set[int], seed_idx: int | None,
    ) -> list[Fused]:
        pool = self.s.candidate_pool

        t = time.perf_counter()
        lex = self.lexical.search(lex_tokens, pool)
        timing.lexical = round((time.perf_counter() - t) * 1000, 2)

        t = time.perf_counter()
        if seed_idx is not None:
            qv = self.semantic.encode_query(query_for_semantic)[0]
            blended = 0.55 * qv + 0.45 * self.semantic.vector_of(seed_idx)
            blended = blended / (np.linalg.norm(blended) + 1e-9)
            sem = self.semantic.search_by_vector(blended, pool)
        else:
            sem = self.semantic.search(query_for_semantic, pool)
        timing.semantic = round((time.perf_counter() - t) * 1000, 2)

        if exclude:
            lex = [(i, s) for i, s in lex if i not in exclude]
            sem = [(i, s) for i, s in sem if i not in exclude]

        t = time.perf_counter()
        fused = fuse(lex, sem, method=req.fusion, rrf_k=self.s.rrf_k,
                     w_lex=req.lexical_weight, w_sem=req.semantic_weight)
        timing.fusion = round((time.perf_counter() - t) * 1000, 2)
        return fused

    @staticmethod
    def _display_scores(finals: list[float]) -> list[float]:
        """Turn raw ranking scores into a satisfying 'match confidence' for the
        UI: the page top is anchored to its absolute quality (0.80..0.99), the
        rest spread down to a floor via a soft sqrt curve. The truthful
        lexical / semantic / fusion / rerank sub-scores stay untouched in
        ScoreBreakdown."""
        if not finals:
            return []
        hi, lo = max(finals), min(finals)
        top_disp = 0.80 + 0.19 * max(0.0, min(1.0, hi))
        floor = 0.62
        if hi - lo < 1e-6:
            return [round(top_disp, 4)] * len(finals)
        return [round(floor + (top_disp - floor) * math.sqrt((f - lo) / (hi - lo)), 4)
                for f in finals]

    def _assemble(
        self, query: str, parsed: ParsedQuery, fused: list[Fused], lex_tokens: list[str],
        req: SearchRequest, timing: Timing,
    ) -> list[SearchResult]:
        # hard-exclude explicitly negated genres ("... not horror") when doing
        # so still leaves enough results; otherwise fall back to soft penalties.
        neg_genres = {GENRE_CANON[w.lower()] for w in parsed.negations
                      if w.lower() in GENRE_CANON}
        if neg_genres:
            kept = [f for f in fused
                    if not (set(self.corpus.record(f.idx)["genres"]) & neg_genres)]
            if len(kept) >= req.limit:
                fused = kept

        # normalize fused scores over the shortlist for display
        shortlist = fused[: max(self.s.rerank_pool, req.limit)]
        fused_norm = _minmax_list([f.fused for f in shortlist])
        for f, fn in zip(shortlist, fused_norm):
            f.fused = fn  # reuse field to hold normalized fusion score for display

        # ---- rerank the shortlist ----
        rr_scores: dict[int, float] | None = None
        rerank_used = False
        t = time.perf_counter()
        if req.rerank and self.reranker:
            docs = [(f.idx, self.corpus.record(f.idx)["semantic_text"][:600]) for f in shortlist]
            rr_scores = self.reranker.rerank(query, docs)
            rerank_used = rr_scores is not None
        timing.rerank = round((time.perf_counter() - t) * 1000, 2)

        # ---- final ranking score = base (rerank|fusion) + intent boosts ----
        exact_idx = self.title_index.get(_norm(query))
        scored: list[tuple[float, Fused, float | None]] = []
        for f in shortlist:
            rec = self.corpus.record(f.idx)
            base = rr_scores[f.idx] if rerank_used else f.fused
            adj = self._boost(parsed, rec)
            if exact_idx is not None and f.idx == exact_idx:
                adj += 0.6  # an exact-title query should surface that exact film
            final = max(0.0, min(1.0, base + adj))
            scored.append((final, f, rr_scores[f.idx] if rerank_used else None))
        scored.sort(key=lambda x: x[0], reverse=True)

        top = scored[: req.limit]
        disp = self._display_scores([x[0] for x in top])

        # ---- build results + grounded reasons ----
        results: list[SearchResult] = []
        reason_t = 0.0
        for rank_i, (final, f, rr) in enumerate(top):
            rec = self.corpus.record(f.idx)
            movie = self.corpus.to_movie(f.idx)
            match_terms = self.lexical.match_terms(f.idx, lex_tokens)
            reason = None
            if req.explain:
                rt = time.perf_counter()
                reason = self.slm.explain(query, parsed, rec, match_terms)
                reason_t += (time.perf_counter() - rt) * 1000
            results.append(SearchResult(
                **movie.model_dump(),
                score=disp[rank_i],
                scores=ScoreBreakdown(
                    lexical=round(f.lex_norm, 4),
                    semantic=round(f.sem_norm, 4),
                    fusion=round(f.fused, 4),
                    rerank=round(rr, 4) if rr is not None else None,
                    lexical_rank=f.lex_rank,
                    semantic_rank=f.sem_rank,
                ),
                match_terms=match_terms,
                reason=reason,
            ))
        timing.reason = round(reason_t, 2)
        self._last_rerank_used = rerank_used
        return results

    # --------------------------------------------------------------------- #
    def search(self, req: SearchRequest) -> SearchResponse:
        t0 = time.perf_counter()
        timing = Timing()

        t = time.perf_counter()
        parsed = self.slm.parse(req.query)
        timing.parse = round((time.perf_counter() - t) * 1000, 2)

        lex_tokens = preprocessor.tokenize(req.query)
        # negated terms must not count as positive lexical matches
        neg_lemmas: set[str] = set()
        for w in parsed.negations:
            neg_lemmas.update(preprocessor.tokenize(w))
        if neg_lemmas:
            lex_tokens = [t for t in lex_tokens if t not in neg_lemmas]
        if parsed.min_rating:
            rating_lemmas = {lem for w in _RATING_WORDS for lem in preprocessor.tokenize(w)}
            lex_tokens = [t for t in lex_tokens if t not in rating_lemmas]
        seed_idx = self._resolve_similar(parsed.similar_to)
        exclude = {seed_idx} if seed_idx is not None else set()
        semantic_query = parsed.clean_query or req.query

        fused = self._rank_candidates(
            semantic_query, lex_tokens, parsed, req, timing, exclude, seed_idx
        )
        results = self._assemble(req.query, parsed, fused, lex_tokens, req, timing)
        timing.total = round((time.perf_counter() - t0) * 1000, 2)

        return SearchResponse(
            query=req.query, parsed=parsed, count=len(results), results=results,
            timing_ms=timing, engine=self._engine_info(req, getattr(self, "_last_rerank_used", False)),
        )

    def similar(self, movie_id: str, limit: int) -> SearchResponse:
        idx = self.corpus.index_of(movie_id)
        if idx is None:
            raise KeyError(movie_id)
        rec = self.corpus.record(idx)
        req = SearchRequest(query=rec["title"], limit=limit, fusion="rrf", rerank=True)
        timing = Timing()
        parsed = ParsedQuery(
            intent="similar", clean_query=rec["title"],
            similar_to=rec["title"], genres=rec.get("genres", [])[:2], source="rules",
        )
        lex_tokens = preprocessor.tokenize(rec["lexical_text"])[:40]
        fused = self._rank_candidates(
            rec["semantic_text"][:400], lex_tokens, parsed, req, timing, {idx}, idx
        )
        results = self._assemble(rec["title"], parsed, fused, lex_tokens, req, timing)
        return SearchResponse(
            query=rec["title"], parsed=parsed, count=len(results), results=results,
            timing_ms=timing, engine=self._engine_info(req, getattr(self, "_last_rerank_used", False)),
        )

    # --------------------------------------------------------------------- #
    def suggest(self, q: str, limit: int) -> list[Suggestion]:
        key = _norm(q)
        if not key:
            return []
        starts, contains = [], []
        for i, rec in enumerate(self.corpus.records):
            tk = _norm(rec["title"])
            if tk.startswith(key):
                starts.append(i)
            elif key in tk:
                contains.append(i)
            if len(starts) >= limit and len(contains) >= limit:
                break
        # popularity-rank within each bucket
        starts.sort(key=lambda i: -self.corpus.records[i]["vote_count"])
        contains.sort(key=lambda i: -self.corpus.records[i]["vote_count"])
        picks = (starts + contains)[:limit]
        return [Suggestion(id=self.corpus.records[i]["id"], title=self.corpus.records[i]["title"],
                           year=self.corpus.records[i]["year"],
                           poster_url=self.corpus.records[i]["poster_url"]) for i in picks]

    def get_movie(self, movie_id: str) -> Movie:
        idx = self.corpus.index_of(movie_id)
        if idx is None:
            raise KeyError(movie_id)
        return self.corpus.to_movie(idx)

    def stats(self) -> StatsResponse:
        from collections import Counter
        gc: Counter = Counter()
        dc: Counter = Counter()
        for r in self.corpus.records:
            for g in r["genres"]:
                gc[g] += 1
            if r["year"]:
                dc[f"{r['year'] // 10 * 10}s"] += 1
        top = sorted(
            [r for r in self.corpus.records if (r["vote_count"] or 0) > 3000],
            key=lambda r: -(r["vote_average"] or 0),
        )[:10]
        return StatsResponse(
            corpus_size=self.corpus.size,
            genres=[GenreCount(name=n, count=c) for n, c in gc.most_common(12)],
            decades=[GenreCount(name=n, count=c) for n, c in sorted(dc.items())],
            sample_queries=_SAMPLE_QUERIES,
            top_rated=[Suggestion(id=r["id"], title=r["title"], year=r["year"],
                                  poster_url=r["poster_url"]) for r in top],
        )

    def health(self) -> HealthResponse:
        return HealthResponse(
            status="ok", ready=self._ready, corpus_size=self.corpus.size,
            embedding_model=self.s.embedding_model,
            reranker=bool(self.reranker and self.reranker.available),
            reranker_model=self.s.reranker_model if self.reranker else None,
            slm=self.slm.label, faiss=True,
        )

    def _engine_info(self, req: SearchRequest, rerank_used: bool) -> EngineInfo:
        return EngineInfo(
            fusion=req.fusion, rerank=rerank_used,
            embedding_model=self.s.embedding_model,
            reranker_model=self.s.reranker_model if (self.reranker and rerank_used) else None,
            slm=self.slm.label, corpus_size=self.corpus.size,
        )