"""Base search model for Elasticsearch-backed models."""

import uuid
from typing import ClassVar, Optional, List

from pydantic import BaseModel, Field

from common.database.es_manager import GlobalESManager as ESManager
from common.database.search_fields import SearchField
from common.database.query_builder import QueryBuilder
from common.database.search_type import SearchType


class BaseSearchModel(BaseModel):
    """Base model for Elasticsearch-backed search models.

    This class provides:
    - Automatic index management
    - Rich query building via QueryBuilder
    - Vector, lexical, and hybrid search methods
    - CRUD operations
    """

    __es_index_name__: ClassVar[str | None] = None
    __es_index__: ClassVar[ESManager.ESIndex | None] = None

    id: str = SearchField(
        Field(default_factory=lambda: str(uuid.uuid4().hex)),
        es_ignore=True,  # Don't store in _source since we'll use ES _id
    )

    @classmethod
    def query(cls) -> QueryBuilder["BaseSearchModel"]:
        """Start building a query for this model"""
        return QueryBuilder(cls)

    @classmethod
    async def get_index(cls) -> ESManager.ESIndex:
        if cls.__es_index__ is None:
            # Build the index
            cls.__es_index__ = await ESManager.get_or_create_index(cls)
        return cls.__es_index__

    @classmethod
    def get_index_name(cls) -> str:
        if cls.__es_index_name__ is None:
            cls.__es_index_name__ = cls.__name__.lower()
        return cls.__es_index_name__

    @classmethod
    async def get(cls, id: str) -> Optional["BaseSearchModel"]:
        """
        Retrieve a document by its ID.

        Args:
            id: Document ID

        Returns:
            BaseSearchModel | None: The document if found, None otherwise
        """
        results = await ESManager.search(
            cls, {"query": {"ids": {"values": [id]}}}, size=1
        )
        return results[0].model if results else None

    @classmethod
    async def delete(cls, id: str) -> None:
        await ESManager.delete(cls, id)

    @classmethod
    async def create(cls, **kwargs) -> "BaseSearchModel":
        obj = cls(**kwargs)
        await obj.save_interface()
        return obj

    async def update(self) -> "BaseSearchModel":
        await self.save_interface()
        return self

    async def save_interface(self, **kwargs) -> None:
        """
        Save/update the document in Elasticsearch.

        Args:
            **kwargs: Additional parameters to pass to ES index API
        """
        await ESManager.index(self, **kwargs)

    @classmethod
    async def _perform_single_field_search(
        cls,
        query: str,
        field: str,
        search_type: SearchType,
        num_results: int,
        vector: list[float] | None = None,
        vector_field: str | None = None,
        vector_weight: float = 0.5,
        **search_kwargs,
    ) -> List:
        """Perform a search on a single field.

        Args:
            query: Query for field
            field: Field to search
            search_type: Search type (SearchType.VECTOR, SearchType.LEXICAL, or SearchType.HYBRID)
            num_results: Number of results to return
            vector: Query vector to search with
            vector_field: Field to search for vector similarity
            vector_weight: Weight for vector search, applies only to hybrid search
            **search_kwargs: Additional search parameters

        Returns:
            List of search results
        """

        # Determine the search type, prepare the search parameters, and call the appropriate search method
        if search_type == SearchType.VECTOR:
            if search_kwargs.get("pre_filters") is None:
                search_kwargs["pre_filters"] = {}
            # FIXME: This shouldn't be here; subclasses have no obligation to have a ..._valid field.
            embedding_valid_field = f"{vector_field}_valid"
            search_kwargs["pre_filters"][embedding_valid_field] = True

            return await cls.vector_search(
                field=vector_field,
                vector=vector,
                k=num_results,
                **search_kwargs,
            )

        elif search_type == SearchType.LEXICAL:
            return await cls.lexical_search(
                field=field, text=query, k=num_results, **search_kwargs
            )

        else:  # SearchType.HYBRID
            if search_kwargs.get("knn_pre_filters") is None:
                search_kwargs["knn_pre_filters"] = {}
            embedding_valid_field = f"{vector_field}_valid"
            search_kwargs["knn_pre_filters"][embedding_valid_field] = True

            return await cls.hybrid_search(
                text_field=field,
                vector_field=vector_field,
                text=query,
                vector=vector,
                k=num_results,
                vector_weight=vector_weight,
                **search_kwargs,
            )

    @classmethod
    async def lexical_search(
        cls,
        field: str,
        text: str,
        *,
        k: int = 10,
        score_threshold: float | None = None,
        allowed_ids: List[str] | None = None,
        **kwargs,
    ) -> list[ESManager.SearchResult]:
        """
        Perform a lexical (text-based) search.

        Args:
            field: Name of the text field to search
            text: The text query to search for
            k: Number of results to return
            score_threshold: Remove all results with scores below this threshold
            allowed_ids: If provided, only return results with these IDs.
            **kwargs: Field-value pairs to filter results (e.g., category="books")

        Returns:
            list[ESManager.SearchResult]: List of search results with scores and metadata
        """
        query = cls.query().should(**{field: text}).size(k)

        if kwargs:
            query.filter(**kwargs)

        if allowed_ids is not None:
            query.allow_ids(allowed_ids)

        return await ESManager.search(
            cls,
            query.build(),
            score_threshold=score_threshold,
        )

    @classmethod
    async def vector_search(
        cls,
        field: str,
        vector: list[float],
        *,
        k: int = 10,
        num_candidates: int = 100,
        score_threshold: float | None = None,
        pre_filters: dict | None = None,
        allowed_ids: List[str] | None = None,
        **kwargs,
    ) -> list[ESManager.SearchResult]:
        """
        Perform a vector similarity search with optional filters.

        Args:
            field: Name of the dense vector field to search
            vector: Query vector to search with
            k: Number of results to return (default: 10)
            num_candidates: Number of candidates to consider (default: 100)
            score_threshold: Remove all results with scores below this threshold
            pre_filters: Additional filters to apply before the knn search
            allowed_ids: If provided, only return results with these IDs.
            **kwargs: Field-value pairs to filter results (e.g., category="books")

        Returns:
            list[ESManager.SearchResult]: List of search results with scores and metadata
        """
        query = cls.query()

        # Build filters for KNN
        filters = {}
        if pre_filters:
            filters.update(pre_filters)
        if kwargs:
            filters.update(kwargs)

        # Add KNN retriever
        query.knn(
            field=field,
            query_vector=vector,
            k=k,
            num_candidates=num_candidates,
            **filters,
        ).size(k)

        if allowed_ids is not None:
            query.allow_ids(allowed_ids)

        return await ESManager.search(
            cls,
            query.build(),
            score_threshold=score_threshold,
        )

    @classmethod
    async def hybrid_search(
        cls,
        text_field: str | None = None,
        text: str | None = None,
        vector_field: str | None = None,
        vector: list[float] | None = None,
        *,
        k: int = 10,
        vector_weight: float = 0.5,
        num_knn_candidates: int = 100,
        knn_pre_filters: dict | None = None,
        rrf_window_size: int = 100,
        rrf_rank_constant: int = 60,
        score_threshold: float | None = None,
        allowed_ids: List[str] | None = None,
        **kwargs,
    ) -> list[ESManager.SearchResult]:
        """
        Perform a hybrid search using a single text/vector combination.

        Args:
            text_field: Field(s) for text search. Single string.
            text: Query text(s). Single string.
            vector_field: Field(s) for vector search. Single string.
            vector: Query vector(s). Single vector.
            k: Number of results to return
            vector_weight: Weight(s) for vector vs text search (0.0=text only, 1.0=vector only).
                Single float. Default: 0.5
            num_knn_candidates: KNN candidates. Single int.
            knn_pre_filters: KNN filters. Single dict.
            rrf_window_size: RRF window for inner fusion. Single int.
            rrf_rank_constant: RRF constant for inner fusion. Single int.
            score_threshold: Remove results below this score threshold.
            allowed_ids: If provided, only return results with these IDs.
            **kwargs: Additional filters to apply to all searches.
        """

        if not (text_field or text or vector_field or vector):
            raise ValueError(
                "text_field, text, vector_field, and vector must be provided"
            )

        # Build text query
        text_query = cls.query().should(**{text_field: text}).size(k)

        if kwargs:
            text_query.filter(**kwargs)

        if allowed_ids is not None:
            text_query.allow_ids(allowed_ids)

        # Build KNN query
        knn_filters = {}
        if knn_pre_filters:
            knn_filters.update(knn_pre_filters)
        if kwargs:
            knn_filters.update(kwargs)

        vector_query = (
            cls.query()
            .knn(
                field=vector_field,
                query_vector=vector,
                k=rrf_window_size,
                num_candidates=max(num_knn_candidates, rrf_window_size),
                **knn_filters,
            )
            .size(k)
        )

        if allowed_ids is not None:
            vector_query.allow_ids(allowed_ids)

        result_sets = await ESManager.multi_search(
            cls,
            [vector_query.build(), text_query.build()],
            score_threshold=score_threshold,
        )

        return await cls._fuse_results_using_rrf(
            result_sets,
            result_weights=[vector_weight, 1 - vector_weight],
            rrf_window_size=rrf_window_size,
            rrf_rank_constant=rrf_rank_constant,
        )

    @classmethod
    async def _fuse_results_using_rrf(
        cls,
        result_sets: list[list[ESManager.SearchResult]],
        result_weights: list[float] | None = None,
        rrf_window_size: int = 100,
        rrf_rank_constant: int = 60,
        enforce_geometric_mean: bool = False,
    ) -> list[ESManager.SearchResult]:
        """Fuse results from multiple result sets using Reciprocal Rank Fusion.

        Args:
            result_sets: List of result sets to fuse
            result_weights: Optional weights for each result set (will be normalized to sum to 1)
            rrf_window_size: Window size for RRF fusion
            rrf_rank_constant: Constant k in RRF formula (1/(k+rank))
            enforce_geometric_mean: If True, use geometric mean; if False, use arithmetic mean

        Returns:
            Fused results with RRF scores
        """
        import math

        if not result_sets:
            return []

        num_lists = len(result_sets)

        # Normalize weights if provided, otherwise use equal weights
        if result_weights:
            # Normalize weights to sum to 1
            weight_sum = sum(result_weights)
            if weight_sum > 0:
                normalized_weights = [w / weight_sum for w in result_weights]
            else:
                normalized_weights = [1.0 / num_lists] * num_lists
        else:
            normalized_weights = [1.0 / num_lists] * num_lists

        # Track models and scores
        doc_models: dict[str, ESManager.SearchResult] = {}

        if enforce_geometric_mean:
            # For weighted geometric mean: product(score_i^weight_i)
            # In log space: sum(weight_i * log(score_i))
            doc_weighted_log_sums: dict[str, float] = {}
            doc_seen_in_lists: dict[str, set[int]] = {}

            for list_idx, result_set in enumerate(result_sets):
                weight = normalized_weights[list_idx]
                for rank, result in enumerate(result_set[:rrf_window_size], start=1):
                    doc_id = (
                        result.model.id
                        if hasattr(result.model, "id")
                        else str(result.model)
                    )
                    # RRF score for this document in this list
                    rrf_score = 1.0 / (rrf_rank_constant + rank)
                    doc_models.setdefault(doc_id, result)

                    # Add weighted log contribution
                    doc_weighted_log_sums[doc_id] = doc_weighted_log_sums.get(
                        doc_id, 0.0
                    ) + weight * math.log(rrf_score)

                    # Track which lists this document appeared in
                    if doc_id not in doc_seen_in_lists:
                        doc_seen_in_lists[doc_id] = set()
                    doc_seen_in_lists[doc_id].add(list_idx)

            # Handle documents missing from some lists
            epsilon = 1.0 / (rrf_rank_constant + rrf_window_size + 1)
            for doc_id in doc_weighted_log_sums.keys():
                # For each list this document is missing from, add weighted penalty
                for list_idx in range(num_lists):
                    if list_idx not in doc_seen_in_lists[doc_id]:
                        weight = normalized_weights[list_idx]
                        doc_weighted_log_sums[doc_id] += weight * math.log(epsilon)

            # Convert back from log space - no division by num_lists needed since weights are already normalized
            scored_items = [
                (doc_id, math.exp(weighted_log_sum))
                for doc_id, weighted_log_sum in doc_weighted_log_sums.items()
            ]

            sorted_docs = sorted(scored_items, key=lambda x: x[1], reverse=True)

        else:
            # Arithmetic mean with weights
            doc_scores: dict[str, float] = {}

            for list_idx, result_set in enumerate(result_sets):
                weight = normalized_weights[list_idx]
                for rank, result in enumerate(result_set[:rrf_window_size], start=1):
                    doc_id = (
                        result.model.id
                        if hasattr(result.model, "id")
                        else str(result.model)
                    )
                    # Apply weight to the RRF score
                    rrf_score = weight * (1.0 / (rrf_rank_constant + rank))
                    if doc_id in doc_scores:
                        doc_scores[doc_id] += rrf_score
                    else:
                        doc_scores[doc_id] = rrf_score
                        doc_models[doc_id] = result

            sorted_docs = sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)

        # Create result objects with fused scores
        results = []
        response_uuid = uuid.uuid4().hex
        total = len(sorted_docs)

        for doc_id, score in sorted_docs:
            # Clone the result object with new score
            original_result = doc_models[doc_id]
            results.append(
                ESManager.SearchResult(
                    model=original_result.model,
                    score=score,
                    total=total,
                    response_uuid=response_uuid,
                )
            )

        return results
