#!/usr/bin/env python
# Copyright (C) 2015, 2019 IBM. All rights reserved.
#
# Licensed 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.
"""
API module for composing and executing Cloudant queries.
"""

import json
import contextlib

from ._2to3 import iteritems_
from .result import QueryResult
from .error import CloudantArgumentError
from ._common_util import QUERY_ARG_TYPES
from ._common_util import response_to_json_dict

class Query(dict):
    """
    Encapsulates a query as a dictionary based object, providing a sliceable
    and iterable query result collection that can be used to process query
    output data through the ``result`` attribute.

    For example:

    .. code-block:: python

        # Slicing to skip/limit:
        query.result[100:200]
        query.result[:200]
        query.result[100:]
        query.result[:]

        # Iteration is supported via the result attribute:
        for doc in query.result:
            print doc

    The query ``result`` collection provides basic functionality,
    which can be customized with other arguments using the
    :func:`~cloudant.query.Query.custom_result` context.

    For example:

    .. code-block:: python

        # Setting the read quorum as part of a custom result
        with query.custom_result(r=3) as rslt:
            rslt[100:200] # slice the result

            # Iteration
            for doc in rslt:
                print doc

        # Iteration over a query result sorted by the "name" field:
        with query.custom_result(sort=[{'name': 'asc'}]) as rslt:
            for doc in rslt:
                print doc

    :param CloudantDatabase database: A Cloudant database instance used by the
        Query.
    :param str bookmark: A string that enables you to specify which page of
        results you require. Only valid for queries using indexes of type
        *text*.
    :param list fields: A list of fields to be returned by the query.
    :param int limit: Maximum number of results returned.
    :param int r: Read quorum needed for the result.  Each document is read from
        at least 'r' number of replicas before it is returned in the results.
    :param dict selector: Dictionary object describing criteria used to select
        documents.
    :param int skip: Skip the first 'n' results, where 'n' is the value
        specified.
    :param list sort: A list of fields to sort by.  Optionally the list can
        contain elements that are single member dictionary structures that
        specify sort direction.  For example ``sort=['name', {'age': 'desc'}]``
        means to sort the query results by the "name" field in ascending order
        and the "age" field in descending order.
    :param str use_index: Identifies a specific index for the query to run
        against, rather than using the Cloudant Query algorithm which finds
        what it believes to be the best index.
    :param str partition_key: Optional. Specify a query partition key. Defaults
        to ``None`` resulting in global queries.
    """

    def __init__(self, database, **kwargs):
        super(Query, self).__init__()
        self._database = database
        self._partition_key = kwargs.pop('partition_key', None)
        self._r_session = self._database.r_session
        self._encoder = self._database.client.encoder
        if kwargs.get('fields', True) is None:
            del kwargs['fields']  # delete `None` fields kwarg
        if kwargs:
            super(Query, self).update(kwargs)
        self.result = QueryResult(self)

    @property
    def url(self):
        """
        Constructs and returns the Query URL.

        :returns: Query URL
        """
        if self._partition_key:
            base_url = self._database.database_partition_url(
                self._partition_key)
        else:
            base_url = self._database.database_url

        return base_url + '/_find'

    def __call__(self, **kwargs):
        """
        Makes the Query object callable and retrieves the raw JSON content
        from the remote database based on the current Query definition,
        and any additional kwargs provided as query parameters.

        For example:

        .. code-block:: python

            # Construct a Query
            query = Query(database, selector={'_id': {'$gt': 0}})
            # Use query as a callable limiting results to 100,
            # skipping the first 100.
            for doc in query(limit=100, skip=100)['docs']:
                # Process query data (in JSON format).

        Note:  Rather than using the Query callable directly, if you wish to
        retrieve query results in raw JSON format use the provided database API
        of :func:`~cloudant.database.CouchDatabase.get_query_result`
        and set ``raw_result=True`` instead.

        :param str bookmark: A string that enables you to specify which page of
            results you require. Only valid for queries using indexes of type
            *text*.
        :param list fields: A list of fields to be returned by the query.
        :param int limit: Maximum number of results returned.
        :param int r: Read quorum needed for the result.  Each document is read
            from at least 'r' number of replicas before it is returned in the
            results.
        :param dict selector: Dictionary object describing criteria used to
            select documents.
        :param int skip: Skip the first 'n' results, where 'n' is the value
            specified.
        :param list sort: A list of fields to sort by.  Optionally the list can
            contain elements that are single member dictionary structures that
            specify sort direction.  For example
            ``sort=['name', {'age': 'desc'}]`` means to sort the query results
            by the "name" field in ascending order and the "age" field in
            descending order.
        :param str use_index: Identifies a specific index for the query to run
            against, rather than using the Cloudant Query algorithm which finds
            what it believes to be the best index.

        :returns: Query result data in JSON format
        """
        data = dict(self)
        data.update(kwargs)

        # Validate query arguments and values
        for key, val in iteritems_(data):
            if key not in list(QUERY_ARG_TYPES.keys()):
                raise CloudantArgumentError(129, key)
            if not isinstance(val, QUERY_ARG_TYPES[key]):
                raise CloudantArgumentError(130, key, QUERY_ARG_TYPES[key])
        if data.get('selector', None) is None or data.get('selector') == {}:
            raise CloudantArgumentError(131)

        # Execute query find
        headers = {'Content-Type': 'application/json'}
        resp = self._r_session.post(
            self.url,
            headers=headers,
            data=json.dumps(data, cls=self._encoder)
        )
        resp.raise_for_status()
        return response_to_json_dict(resp)

    @contextlib.contextmanager
    def custom_result(self, **options):
        """
        Customizes the :class:`~cloudant.result.QueryResult` behavior and
        provides a convenient context manager for the QueryResult.  QueryResult
        customizations can be made by providing extra options to the query
        result call using this context manager.  The use of ``skip`` and
        ``limit`` as options are not valid when using a QueryResult since the
        ``skip`` and ``limit`` functionality is handled in the QueryResult.

        For example:

        .. code-block:: python

            with query.custom_result(sort=[{'name': 'asc'}]) as rslt:
                data = rslt[100:200]

        :param str bookmark: A string that enables you to specify which page of
            results you require. Only valid for queries using indexes of type
            *text*.
        :param list fields: A list of fields to be returned by the query.
        :param int page_size: Sets the page size for result iteration.  Default
            is 100.
        :param int r: Read quorum needed for the result.  Each document is read
            from at least 'r' number of replicas before it is returned in the
            results.
        :param dict selector: Dictionary object describing criteria used to
            select documents.
        :param list sort: A list of fields to sort by.  Optionally the list can
            contain elements that are single member dictionary structures that
            specify sort direction.  For example
            ``sort=['name', {'age': 'desc'}]`` means to sort the query results
            by the "name" field in ascending order and the "age" field in
            descending order.
        :param str use_index: Identifies a specific index for the query to run
            against, rather than using the Cloudant Query algorithm which finds
            what it believes to be the best index.

        :returns: Query result data wrapped in a QueryResult instance
        """
        rslt = QueryResult(self, **options)
        yield rslt
        del rslt
