"""Hybrid search: keyword (BM25-like) + optional vector search with RRF fusion."""

from __future__ import annotations

import math
from collections import defaultdict
from typing import Any

from thryve.core.memory.storage import SQLiteStorage
from thryve.core.memory.types import MemoryChunk
from thryve.utils import get_logger

logger = get_logger("memory.searcher")


class HybridSearcher:
    """Hybrid search combining keyword matching and (optionally) vector similarity.

    Weights are configurable per the ``MemoryConfig`` (default 70/30).
    If no embedding provider is available, falls back to keyword-only search.
    """

    VECTOR_WEIGHT = 0.7
    BM25_WEIGHT = 0.3
    RRF_K = 60  # Reciprocal Rank Fusion constant

    def __init__(
        self,
        storage: SQLiteStorage,
        embedding_fn: Any | None = None,
        vector_weight: float = 0.7,
        bm25_weight: float = 0.3,
    ) -> None:
        self._storage = storage
        self._embedding_fn = embedding_fn
        self.VECTOR_WEIGHT = vector_weight
        self.BM25_WEIGHT = bm25_weight

    # ------------------------------------------------------------------
    # Public
    # ------------------------------------------------------------------

    def search(
        self,
        query: str,
        top_k: int = 5,
        source: str | None = None,
    ) -> list[MemoryChunk]:
        """Search memory using keyword + vector (when available)."""
        # 1. Keyword search
        kw_results = self._keyword_search(query, top_k=top_k * 3, source=source)

        # 2. Vector search (if available)
        if self._embedding_fn is not None:
            vec_results = self._vector_search(query, top_k=top_k * 3, source=source)
            fused = self._rrf_fuse(kw_results, vec_results)
            return fused[:top_k]

        return kw_results[:top_k]

    # ------------------------------------------------------------------
    # Keyword search
    # ------------------------------------------------------------------

    def _keyword_search(
        self, query: str, top_k: int = 15, source: str | None = None,
    ) -> list[MemoryChunk]:
        """Simple keyword LIKE search via SQLite."""
        return self._storage.search_by_keyword(query, top_k=top_k, source=source)

    # ------------------------------------------------------------------
    # Vector search (placeholder for now)
    # ------------------------------------------------------------------

    def _vector_search(
        self, query: str, top_k: int = 15, source: str | None = None,
    ) -> list[MemoryChunk]:
        """Brute-force cosine similarity search over stored embeddings.

        For production, consider using a dedicated vector store.
        """
        if self._embedding_fn is None:
            return []

        try:
            query_emb = self._embedding_fn([query])[0]
        except Exception:
            logger.warning("Failed to generate query embedding; skipping vector search")
            return []

        all_chunks = self._storage.get_all(source=source, limit=1000)
        scored: list[tuple[float, MemoryChunk]] = []
        for chunk in all_chunks:
            if chunk.embedding:
                sim = self._cosine_similarity(query_emb, chunk.embedding)
                scored.append((sim, chunk))

        scored.sort(key=lambda x: x[0], reverse=True)
        return [c for _, c in scored[:top_k]]

    # ------------------------------------------------------------------
    # RRF fusion
    # ------------------------------------------------------------------

    def _rrf_fuse(
        self, *result_lists: list[MemoryChunk],
    ) -> list[MemoryChunk]:
        """Reciprocal Rank Fusion across multiple ranked result lists."""
        scores: dict[str, float] = defaultdict(float)
        chunk_map: dict[str, MemoryChunk] = {}

        for results in result_lists:
            for rank, chunk in enumerate(results):
                scores[chunk.id] += 1.0 / (self.RRF_K + rank)
                chunk_map[chunk.id] = chunk

        ranked_ids = sorted(scores, key=lambda cid: scores[cid], reverse=True)
        return [chunk_map[cid] for cid in ranked_ids]

    # ------------------------------------------------------------------
    # Utility
    # ------------------------------------------------------------------

    @staticmethod
    def _cosine_similarity(a: list[float], b: list[float]) -> float:
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))
        if norm_a == 0 or norm_b == 0:
            return 0.0
        return dot / (norm_a * norm_b)
