#  Licensed to Elasticsearch B.V. under one or more contributor
#  license agreements. See the NOTICE file distributed with
#  this work for additional information regarding copyright
#  ownership. Elasticsearch B.V. licenses this file to you under
#  the Apache License, Version 2.0 (the "License"); you may
#  not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
# 	http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing,
#  software distributed under the License is distributed on an
#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
#  KIND, either express or implied.  See the License for the
#  specific language governing permissions and limitations
#  under the License.

import contextlib
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Iterator,
    List,
    Optional,
    cast,
)

from typing_extensions import Self

from elasticsearch.exceptions import ApiError
from elasticsearch.helpers import scan

from ..connections import get_connection
from ..response import Response
from ..search_base import MultiSearchBase, SearchBase
from ..utils import _R, AttrDict, UsingType


class Search(SearchBase[_R]):
    _using: UsingType

    def __iter__(self) -> Iterator[_R]:
        """
        Iterate over the hits.
        """

        class ResultsIterator(Iterator[_R]):
            def __init__(self, search: Search[_R]):
                self.search = search
                self.iterator: Optional[Iterator[_R]] = None

            def __next__(self) -> _R:
                if self.iterator is None:
                    self.iterator = iter(self.search.execute())
                try:
                    return next(self.iterator)
                except StopIteration:
                    raise StopIteration()

        return ResultsIterator(self)

    def count(self) -> int:
        """
        Return the number of hits matching the query and filters. Note that
        only the actual number is returned.
        """
        if hasattr(self, "_response") and self._response.hits.total.relation == "eq":  # type: ignore[attr-defined]
            return cast(int, self._response.hits.total.value)  # type: ignore[attr-defined]

        es = get_connection(self._using)

        d = self.to_dict(count=True)
        # TODO: failed shards detection
        resp = es.count(
            index=self._index,
            query=cast(Optional[Dict[str, Any]], d.get("query", None)),
            **self._params,
        )

        return cast(int, resp["count"])

    def execute(self, ignore_cache: bool = False) -> Response[_R]:
        """
        Execute the search and return an instance of ``Response`` wrapping all
        the data.

        :arg ignore_cache: if set to ``True``, consecutive calls will hit
            ES, while cached result will be ignored. Defaults to `False`
        """
        if ignore_cache or not hasattr(self, "_response"):
            es = get_connection(self._using)

            self._response = self._response_class(
                self,
                (
                    es.search(index=self._index, body=self.to_dict(), **self._params)
                ).body,
            )
        return self._response

    def scan(self) -> Iterator[_R]:
        """
        Turn the search into a scan search and return a generator that will
        iterate over all the documents matching the query.

        Use the ``params`` method to specify any additional arguments you wish to
        pass to the underlying ``scan`` helper from ``elasticsearch-py`` -
        https://elasticsearch-py.readthedocs.io/en/latest/helpers.html#scan

        The ``iterate()`` method should be preferred, as it provides similar
        functionality using an Elasticsearch point in time.
        """
        es = get_connection(self._using)

        for hit in scan(es, query=self.to_dict(), index=self._index, **self._params):
            yield self._get_result(cast(AttrDict[Any], hit))

    def delete(self) -> AttrDict[Any]:
        """
        ``delete()`` executes the query by delegating to ``delete_by_query()``.

        Use the ``params`` method to specify any additional arguments you wish to
        pass to the underlying ``delete_by_query`` helper from ``elasticsearch-py`` -
        https://elasticsearch-py.readthedocs.io/en/latest/api/elasticsearch.html#elasticsearch.Elasticsearch.delete_by_query
        """

        es = get_connection(self._using)
        assert self._index is not None

        return AttrDict(
            cast(
                Dict[str, Any],
                es.delete_by_query(
                    index=self._index, body=self.to_dict(), **self._params
                ),
            )
        )

    @contextlib.contextmanager
    def point_in_time(self, keep_alive: str = "1m") -> Iterator[Self]:
        """
        Open a point in time (pit) that can be used across several searches.

        This method implements a context manager that returns a search object
        configured to operate within the created pit.

        :arg keep_alive: the time to live for the point in time, renewed with each search request
        """
        es = get_connection(self._using)

        pit = es.open_point_in_time(index=self._index or "*", keep_alive=keep_alive)
        search = self.index().extra(pit={"id": pit["id"], "keep_alive": keep_alive})
        if not search._sort:
            search = search.sort("_shard_doc")
        yield search
        es.close_point_in_time(id=pit["id"])

    def iterate(self, keep_alive: str = "1m") -> Iterator[_R]:
        """
        Return a generator that iterates over all the documents matching the query.

        This method uses a point in time to provide consistent results even when
        the index is changing. It should be preferred over ``scan()``.

        :arg keep_alive: the time to live for the point in time, renewed with each new search request
        """
        with self.point_in_time(keep_alive=keep_alive) as s:
            while True:
                r = s.execute()
                for hit in r:
                    yield hit
                if len(r.hits) == 0:
                    break
                s = s.search_after()


class MultiSearch(MultiSearchBase[_R]):
    """
    Combine multiple :class:`~elasticsearch.dsl.Search` objects into a single
    request.
    """

    _using: UsingType

    if TYPE_CHECKING:

        def add(self, search: Search[_R]) -> Self: ...  # type: ignore[override]

    def execute(
        self, ignore_cache: bool = False, raise_on_error: bool = True
    ) -> List[Response[_R]]:
        """
        Execute the multi search request and return a list of search results.
        """
        if ignore_cache or not hasattr(self, "_response"):
            es = get_connection(self._using)

            responses = es.msearch(
                index=self._index, body=self.to_dict(), **self._params
            )

            out: List[Response[_R]] = []
            for s, r in zip(self._searches, responses["responses"]):
                if r.get("error", False):
                    if raise_on_error:
                        raise ApiError("N/A", meta=responses.meta, body=r)
                    r = None
                else:
                    r = Response(s, r)
                out.append(r)

            self._response = out

        return self._response


class EmptySearch(Search[_R]):
    def count(self) -> int:
        return 0

    def execute(self, ignore_cache: bool = False) -> Response[_R]:
        return self._response_class(self, {"hits": {"total": 0, "hits": []}})

    def scan(self) -> Iterator[_R]:
        return
        yield  # a bit strange, but this forces an empty generator function

    def delete(self) -> AttrDict[Any]:
        return AttrDict[Any]({})