from __future__ import annotations

from dataclasses import dataclass

from rak.bm25 import BM25Index
from rak.embedder import Embedder
from rak.store import VectorStore


@dataclass
class SearchResult:
    doc_id: str
    score: float
    title: str = ""
    source: str = ""
    snippet: str = ""


def rrf_fuse(
    ranked_lists: list[list[dict]],
    limit: int = 10,
    k: int = 60,  # Standard RRF constant; higher k reduces impact of top ranks (60 is the value from the original Cormack et al. paper)
) -> list[SearchResult]:
    scores: dict[str, float] = {}
    titles: dict[str, str] = {}
    snippets: dict[str, str] = {}
    for ranked_list in ranked_lists:
        for rank, item in enumerate(ranked_list):
            doc_id = item["id"]
            scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank + 1)
            if "metadata" in item and "title" in item["metadata"]:
                titles[doc_id] = item["metadata"]["title"]
            elif "title" in item:
                titles[doc_id] = item["title"]
            if doc_id not in snippets and item.get("document"):
                snippets[doc_id] = item["document"]
    sorted_ids = sorted(scores, key=lambda x: scores[x], reverse=True)[:limit]
    return [
        SearchResult(
            doc_id=doc_id, score=scores[doc_id], title=titles.get(doc_id, ""),
            source="fused", snippet=snippets.get(doc_id, ""),
        )
        for doc_id in sorted_ids
    ]


def build_where_filter(
    collection: str | None = None,
    tags: list[str] | None = None,
) -> dict | None:
    filters = []
    if collection:
        filters.append({"collections": {"$contains": collection}})
    if tags:
        if len(tags) == 1:
            filters.append({"tags": {"$contains": tags[0]}})
        else:
            filters.append({"$or": [{"tags": {"$contains": t}} for t in tags]})
    if not filters:
        return None
    if len(filters) == 1:
        return filters[0]
    return {"$and": filters}


def _deduplicate_chunks(results: list[SearchResult]) -> list[SearchResult]:
    """Deduplicate chunk results back to parent papers, keeping the best score and snippet."""
    seen: dict[str, SearchResult] = {}
    for r in results:
        parent_key = r.doc_id.split("_chunk_")[0] if "_chunk_" in r.doc_id else r.doc_id
        if parent_key not in seen or r.score > seen[parent_key].score:
            seen[parent_key] = SearchResult(
                doc_id=parent_key, score=r.score, title=r.title, source=r.source,
                snippet=r.snippet,
            )
    return sorted(seen.values(), key=lambda x: x.score, reverse=True)


class Searcher:
    def __init__(self, embedder: Embedder | None, vector_store: VectorStore | None, bm25_index: BM25Index) -> None:
        self._embedder = embedder
        self._vector_store = vector_store
        self._bm25 = bm25_index

    def bm25_search(
        self, query: str, limit: int = 10,
    ) -> list[SearchResult]:
        """Pure BM25 keyword search. Does not require embedder or vector store."""
        results = self._bm25.search_with_snippet(query, limit=limit)
        raw = [
            SearchResult(
                doc_id=r["id"], score=r["score"],
                title=r.get("title", ""), source="bm25",
                snippet=r.get("snippet", ""),
            )
            for r in results
        ]
        return _deduplicate_chunks(raw)[:limit]

    def vector_search(
        self, query: str, limit: int = 10,
        collection: str | None = None, tags: list[str] | None = None,
    ) -> list[SearchResult]:
        embedding = self._embedder.embed(query)
        where = build_where_filter(collection, tags)
        # Fetch 3x results to ensure enough unique papers after chunk deduplication
        # (multiple chunks per paper are stored as separate vectors)
        results = self._vector_store.search(embedding, limit=limit * 3, where=where)
        raw = [
            SearchResult(
                doc_id=r["id"], score=r["score"],
                title=r.get("metadata", {}).get("title", ""),
                source="vector",
                snippet=r.get("document", ""),
            )
            for r in results
        ]
        return _deduplicate_chunks(raw)[:limit]

    def similar_search(
        self, doc_key: str, limit: int = 10,
        collection: str | None = None, tags: list[str] | None = None,
    ) -> list[SearchResult]:
        """Find papers similar to a given document by its key."""
        if self._vector_store is None:
            return []
        # Try the key directly, then try the first chunk
        embedding = self._vector_store.get_embedding(doc_key)
        if embedding is None:
            embedding = self._vector_store.get_embedding(f"{doc_key}_chunk_0")
        if embedding is None:
            return []
        where = build_where_filter(collection, tags)
        # Fetch extra to account for chunks and self-exclusion
        results = self._vector_store.search(embedding, limit=(limit + 1) * 3, where=where)
        raw = [
            SearchResult(
                doc_id=r["id"], score=r["score"],
                title=r.get("metadata", {}).get("title", ""),
                source="similar",
                snippet=r.get("document", ""),
            )
            for r in results
        ]
        deduped = _deduplicate_chunks(raw)
        # Exclude the source paper
        return [r for r in deduped if r.doc_id != doc_key][:limit]

    def hybrid_search(
        self, query: str, limit: int = 10,
        collection: str | None = None, tags: list[str] | None = None,
    ) -> list[SearchResult]:
        embedding = self._embedder.embed(query)
        where = build_where_filter(collection, tags)
        vector_results = self._vector_store.search(embedding, limit=limit * 3, where=where)
        bm25_results = self._bm25.search(query, limit=limit * 2)
        fused = rrf_fuse([vector_results, bm25_results], limit=limit * 3)
        return _deduplicate_chunks(fused)[:limit]
