"""Elasticsearch query builder for constructing complex search queries."""

from typing import Type, TypeVar, Generic, TYPE_CHECKING

from common.database.search_fields import Sort

if TYPE_CHECKING:
    from common.database.search_model import BaseSearchModel


T = TypeVar("T", bound="BaseSearchModel")


class QueryBuilder(Generic[T]):
    """Query builder for constructing Elasticsearch queries with the ES 8.16+ retriever format."""

    def __init__(self, model_class: Type[T]):
        self.model_class = model_class
        self.query = {"bool": {"must": [], "should": [], "must_not": [], "filter": []}}
        self.from_value = 0
        self.sort_fields = []
        self.retriever = None
        self.size_value = None
        self.allowed_ids = None

    def allow_ids(self, ids: list[str]) -> "QueryBuilder[T]":
        if len(ids) == 0:
            return self
        if self.allowed_ids is None:
            self.allowed_ids = set(ids)
        else:
            self.allowed_ids.intersection_update(set(ids))
        return self

    def must(self, **kwargs) -> "QueryBuilder[T]":
        """Field must match the given value(s)"""
        for field, value in kwargs.items():
            self.query["bool"]["must"].append({"match": {field: value}})
        return self

    def should(self, **kwargs) -> "QueryBuilder[T]":
        """Field should match any of the given values"""
        for field, value in kwargs.items():
            self.query["bool"]["should"].append({"match": {field: value}})
        return self

    def must_not(self, **kwargs) -> "QueryBuilder[T]":
        """Field must not match the given value(s)"""
        for field, value in kwargs.items():
            self.query["bool"]["must_not"].append({"match": {field: value}})
        return self

    def filter(self, **kwargs):
        """Field must match the given value(s) without scoring"""
        for field, value in kwargs.items():
            self.query["bool"]["filter"].append({"term": {field: value}})
        return self

    def filter_range(
        self, field: str, gt=None, gte=None, lt=None, lte=None
    ) -> "QueryBuilder[T]":
        """Add range filters"""
        range_params = {}
        if gt is not None:
            range_params["gt"] = gt
        if gte is not None:
            range_params["gte"] = gte
        if lt is not None:
            range_params["lt"] = lt
        if lte is not None:
            range_params["lte"] = lte

        if range_params:
            self.query["bool"]["filter"].append({"range": {field: range_params}})
        return self

    def sort(self, field: str, order: Sort = Sort.ASC) -> "QueryBuilder[T]":
        """Add sort criteria"""
        self.sort_fields.append({field: order})
        return self

    def offset(self, offset: int) -> "QueryBuilder[T]":
        """Set starting offset"""
        self.from_value = offset
        return self

    def size(self, size: int) -> "QueryBuilder[T]":
        """Set number of results to return"""
        self.size_value = size
        return self

    def knn(
        self,
        field: str,
        query_vector: list[float],
        k: int = 10,
        num_candidates: int = 100,
        **filter_kwargs,
    ) -> "QueryBuilder[T]":
        """Add a KNN retriever for vector search (ES 8.16+ format).

        Args:
            field: Name of the dense vector field
            query_vector: The query vector
            k: Number of results to return
            num_candidates: Number of candidates to consider
            **filter_kwargs: Additional keyword filters (converted to term filters)

        Returns:
            self: For method chaining
        """
        knn_config = {
            "field": field,
            "query_vector": query_vector,
            "k": k,
            "num_candidates": num_candidates,
        }

        # Build filters if provided
        filters = []

        if filter_kwargs:
            for key, value in filter_kwargs.items():
                filters.append({"term": {key: value}})

        if filters:
            if len(filters) == 1:
                knn_config["filter"] = filters[0]
            else:
                knn_config["filter"] = {"bool": {"must": filters}}

        self.retriever = {"knn": knn_config}
        return self

    def build(self) -> dict:
        """Build the final query dict using the ES 8.16+ retriever format.

        This method handles:
        - Single retrievers (including single RRF retrievers)
        - Transforms a single query into a retriever-based format
        """
        query_dict = {}

        # Use the new retriever-based format
        if self.retriever:
            # Single retriever (could be a standard, knn, or rrf retriever)
            query_dict["retriever"] = self.retriever
        elif any(self.query["bool"].values()):
            # Fallback to standard query if no retrievers but bool query exists
            query_dict["retriever"] = {
                "standard": {"query": {"bool": self.query["bool"]}}
            }
        else:
            # No retrievers or query specified - raise an error
            raise ValueError(
                "Query must have at least one retriever or query condition"
            )

        # Handle the case where allowed_ids is not None
        if self.allowed_ids is not None:
            ids_values = list(self.allowed_ids)
            if len(ids_values) == 0:
                # No allowed docs → force a no-op search
                query_dict["retriever"] = {"standard": {"query": {"match_none": {}}}}
            else:
                ids_filter = {"ids": {"values": ids_values}}
                retriever = query_dict.get("retriever", {})
                if "knn" in retriever:
                    knn_cfg = retriever["knn"]
                    existing_filter = knn_cfg.get("filter")
                    if existing_filter is None:
                        knn_cfg["filter"] = ids_filter
                    else:
                        # Combine with bool must
                        knn_cfg["filter"] = {"bool": {"must": [existing_filter, ids_filter]}}
                elif "standard" in retriever:
                    std = retriever["standard"]
                    q = std.get("query", {})
                    bool_q = q.get("bool")
                    if bool_q is None:
                        std["query"] = {"bool": {"filter": [ids_filter]}}
                    else:
                        # Ensure filter is a list
                        if "filter" not in bool_q:
                            bool_q["filter"] = []
                        elif isinstance(bool_q["filter"], dict):
                            bool_q["filter"] = [bool_q["filter"]]
                        bool_q["filter"].append(ids_filter)
                query_dict["retriever"] = retriever

        if self.size_value:
            query_dict["size"] = self.size_value

        if self.sort_fields:
            query_dict["sort"] = self.sort_fields

        if self.from_value:
            query_dict["from"] = self.from_value

        return query_dict
