# The MIT License (MIT)
# Copyright (c) 2014 Microsoft Corporation

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Internal class for query execution context implementation in the Azure Cosmos
database service.
"""

from collections import deque
import copy
import logging
from .. import _retry_utility, http_constants, exceptions, _base
from .._constants import _Constants as Constants

_LOGGER = logging.getLogger(__name__)

# pylint: disable=protected-access


class _QueryExecutionContextBase(object):
    """
    This is the abstract base execution context class.
    """

    def __init__(self, client, options):
        """
        :param CosmosClient client:
        :param dict options: The request options for the request.
        """
        self._client = client
        self._options = options
        self._continuation = self._get_initial_continuation()
        self._has_started = False
        self._has_finished = False
        self._buffer = deque()
        self._resource_link = None

    def _get_initial_continuation(self):
        if "continuation" in self._options:
            return self._options["continuation"]
        return None

    def _has_more_pages(self):
        return not self._has_finished

    def _ensure(self):
        if not self._has_more_pages():
            return

        if not self._buffer:
            results = self._fetch_next_block()
            self._buffer.extend(results)

        if not self._buffer:
            self._has_finished = True

    def fetch_next_block(self):
        """Returns a block of results with respecting retry policy.

        This method only exists for backward compatibility reasons. (Because
        QueryIterable has exposed fetch_next_block api).

        :return: List of results.
        :rtype: list
        """
        self._ensure()
        res = list(self._buffer)
        self._buffer.clear()
        return res

    def _fetch_next_block(self):
        raise NotImplementedError

    def __iter__(self):
        """Returns itself as an iterator
        :returns: Query as an iterator.
        :rtype: Iterator
        """
        return self

    def __next__(self):
        """Return the next query result.

        :return: The next query result.
        :rtype: dict
        :raises StopIteration: If no more result is left.
        """
        self._ensure()

        if not self._buffer:
            raise StopIteration

        return self._buffer.popleft()

    def _fetch_items_helper_no_retries(self, fetch_function):
        """Fetches more items and doesn't retry on failure

        :param Callable fetch_function: The function that fetches the items.
        :return: List of fetched items.
        :rtype: list
        """
        fetched_items = []
        new_options = copy.deepcopy(self._options)
        while self._continuation or not self._has_started:
            new_options["continuation"] = self._continuation

            response_headers = {}
            (fetched_items, response_headers) = fetch_function(new_options)
            if not self._has_started:
                self._has_started = True

            continuation_key = http_constants.HttpHeaders.Continuation
            self._continuation = response_headers.get(continuation_key)

            if fetched_items:
                break
        return fetched_items

    def _fetch_items_helper_with_retries(self, fetch_function):
        # TODO: Properly propagate kwargs from retry utility to fetch function
        # the callback keep the **kwargs parameter to maintain compatibility with the retry utility's execution pattern.
        # Execute passes retry context parameters (timeout, operation start time, logger, etc.)
        # The callback need to accept these parameters even if unused
        # Removing **kwargs results in a TypeError when Execute tries to pass these parameters
        def execute_fetch():
            def callback(**kwargs):  # pylint: disable=unused-argument
                return self._fetch_items_helper_no_retries(fetch_function)

            return _retry_utility.Execute(
                self._client, self._client._global_endpoint_manager, callback, **self._options
            )

        # Check if this is an internal partition key range fetch - skip 410 retry logic to avoid recursion
        # When we call refresh_routing_map_provider(), it triggers _ReadPartitionKeyRanges which would
        # come through this same code path. If that also gets a 410 and tries to refresh, we get infinite recursion.
        is_pk_range_fetch = self._options.get("_internal_pk_range_fetch", False)
        if is_pk_range_fetch:
            # For partition key range queries, just execute without 410 partition split retry
            # The underlying retry utility will still handle other transient errors
            _LOGGER.debug("Partition split retry: Skipping 410 retry for internal PK range fetch")
            return execute_fetch()

        max_retries = 3
        attempt = 0

        while attempt <= max_retries:
            try:
                return execute_fetch()
            except exceptions.CosmosHttpResponseError as e:
                if exceptions._partition_range_is_gone(e):
                    attempt += 1
                    if attempt > max_retries:
                        _LOGGER.error(
                            "Partition split retry: Exhausted all %d retries. "
                            "state: _has_started=%s, _continuation=%s",
                            max_retries, self._has_started, self._continuation
                        )
                        raise  # Exhausted retries, propagate error

                    _LOGGER.warning(
                        "Partition split retry: 410 error (sub_status=%s). Attempt %d of %d. "
                        "Refreshing routing map and resetting state.",
                        getattr(e, 'sub_status', 'N/A'),
                        attempt,
                        max_retries
                    )

                    # Refresh routing map to get new partition key ranges.
                    # When resource_link is available, do a targeted refresh for just this collection
                    # instead of the global refresh of all collections' cached routing maps.
                    if self._resource_link:
                        collection_id = _base.GetResourceIdOrFullNameFromLink(self._resource_link)
                        previous_map = self._client._routing_map_provider._collection_routing_map_by_item.get(
                            collection_id)
                        _LOGGER.debug(
                            "Partition split retry: Targeted refresh for collection %s (has_previous_map=%s)",
                            self._resource_link,
                            previous_map is not None,
                        )
                        refresh_feed_options = {}
                        if Constants.ContainerRID in self._options:
                            refresh_feed_options[Constants.ContainerRID] = self._options[Constants.ContainerRID]
                        if "excludedLocations" in self._options:
                            refresh_feed_options["excludedLocations"] = self._options["excludedLocations"]
                        self._client.refresh_routing_map_provider(
                            self._resource_link,
                            previous_map,
                            refresh_feed_options if refresh_feed_options else None,
                        )
                    else:
                        # No resource_link available — defensive fallback to global nuke.
                        # This branch should not be reached in practice since all callers now pass resource_link.
                        _LOGGER.debug("Partition split retry: No resource_link available, using global refresh")
                        self._client.refresh_routing_map_provider()

                    # Reset execution context state to allow retry from the beginning
                    self._has_started = False
                    self._continuation = None
                    # Retry immediately (no backoff needed for partition splits)
                    continue
                raise  # Not a partition split error, propagate immediately

        # This should never be reached, but added for safety
        _LOGGER.warning(
            "Partition split retry: Unexpectedly exited retry loop without returning results. "
            "This indicates a potential logic error."
        )
        return []
    next = __next__  # Python 2 compatibility.


class _DefaultQueryExecutionContext(_QueryExecutionContextBase):
    """
    This is the default execution context.
    """

    def __init__(self, client, options, fetch_function, resource_link=None):
        """
        :param CosmosClient client:
        :param dict options: The request options for the request.
        :param method fetch_function:
            Will be invoked for retrieving each page
        :param str resource_link:
            Optional collection link used for targeted routing map refresh on 410 partition split.
            When provided, the 410 retry loop refreshes only this collection's cached routing map
            instead of destroying all collections' cached maps.

            Example of `fetch_function`:

            >>> def result_fn(result):
            >>>     return result['Databases']

        """
        super(_DefaultQueryExecutionContext, self).__init__(client, options)
        self._fetch_function = fetch_function
        self._resource_link = resource_link

    def _fetch_next_block(self):  # pylint: disable=inconsistent-return-statements
        while super(_DefaultQueryExecutionContext, self)._has_more_pages() and not self._buffer:
            return self._fetch_items_helper_with_retries(self._fetch_function)
