"""
    Slixmpp: The Slick XMPP Library
    Copyright (C) 2010 Nathanael C. Fritz, Erik Reuterborg Larsson
    This file is part of Slixmpp.

    See the file LICENSE for copying permission.
"""

import logging

import slixmpp
from slixmpp import Iq
from slixmpp.plugins import BasePlugin, register_plugin
from slixmpp.xmlstream import register_stanza_plugin
from slixmpp.plugins.xep_0059 import stanza, Set
from slixmpp.exceptions import XMPPError


log = logging.getLogger(__name__)


class ResultIterator:

    """
    An iterator for Result Set Management
    """

    def __init__(self, query, interface, results='substanzas', amount=10,
                       start=None, reverse=False, recv_interface=None,
                       pre_cb=None, post_cb=None):
        """
        Arguments:
           query     -- The template query
           interface -- The substanza of the query to send, for example disco_items
           recv_interface -- The substanza of the query to receive, for example disco_items
           results   -- The query stanza's interface which provides a
                        countable list of query results.
           amount    -- The max amounts of items to request per iteration
           start     -- From which item id to start
           reverse   -- If True, page backwards through the results
           pre_cb    -- Callback to run before sending the stanza
           post_cb   -- Callback to run after receiving the reply

        Example:
           q = Iq()
           q['to'] = 'pubsub.example.com'
           q['disco_items']['node'] = 'blog'
           for i in ResultIterator(q, 'disco_items', '10'):
               print i['disco_items']['items']

        """
        self.query = query
        self.amount = amount
        self.start = start
        self.interface = interface
        if recv_interface:
            self.recv_interface = recv_interface
        else:
            self.recv_interface = interface
        self.pre_cb = pre_cb
        self.post_cb = post_cb
        self.results = results
        self.reverse = reverse
        self._stop = False

    def __aiter__(self):
        return self

    async def __anext__(self):
        return await self.next()

    async def next(self):
        """
        Return the next page of results from a query.

        Note: If using backwards paging, then the next page of
              results will be the items before the current page
              of items.
        """
        if self._stop:
            raise StopAsyncIteration
        if self.query[self.interface]['rsm']['before'] is None:
            self.query[self.interface]['rsm']['before'] = self.reverse
        self.query['id'] = self.query.stream.new_id()
        self.query[self.interface]['rsm']['max'] = str(self.amount)

        if self.start and self.reverse:
            self.query[self.interface]['rsm']['before'] = self.start
        elif self.start:
            self.query[self.interface]['rsm']['after'] = self.start

        try:
            if self.pre_cb:
                self.pre_cb(self.query)
            r = await self.query.send()

            if not r[self.recv_interface]['rsm']['first'] and \
               not r[self.recv_interface]['rsm']['last']:
                raise StopAsyncIteration

            if r[self.recv_interface]['rsm']['count'] and \
               r[self.recv_interface]['rsm']['first_index']:
                count = int(r[self.recv_interface]['rsm']['count'])
                first = int(r[self.recv_interface]['rsm']['first_index'])
                num_items = len(r[self.recv_interface][self.results])
                if first + num_items == count:
                    self._stop = True

            if self.reverse:
                self.start = r[self.recv_interface]['rsm']['first']
            else:
                self.start = r[self.recv_interface]['rsm']['last']

            if self.post_cb:
                self.post_cb(r)
            return r
        except XMPPError:
            raise StopAsyncIteration


[docs]class XEP_0059(BasePlugin):

    """
    XEP-0050: Result Set Management
    """

    name = 'xep_0059'
    description = 'XEP-0059: Result Set Management'
    dependencies = {'xep_0030'}
    stanza = stanza

    def plugin_init(self):
        """
        Start the XEP-0059 plugin.
        """
        register_stanza_plugin(self.xmpp['xep_0030'].stanza.DiscoItems,
                               self.stanza.Set)

    def plugin_end(self):
        self.xmpp['xep_0030'].del_feature(feature=Set.namespace)

    def session_bind(self, jid):
        self.xmpp['xep_0030'].add_feature(Set.namespace)

[docs]    def iterate(self, stanza, interface, results='substanzas', amount=10, reverse=False,
                recv_interface=None, pre_cb=None, post_cb=None):
        """
        Create a new result set iterator for a given stanza query.

        Arguments:
            stanza    -- A stanza object to serve as a template for
                         queries made each iteration. For example, a
                         basic disco#items query.
            interface -- The name of the substanza to which the
                         result set management stanza should be
                         appended in the query stanza. For example,
                         for disco#items queries the interface
                         'disco_items' should be used.
            recv_interface -- The name of the substanza from which the
                              result set management stanza should be
                              read in the result stanza. If unspecified,
                              it will be set to the same value as the
                              ``interface`` parameter.
            pre_cb    -- Callback to run before sending each stanza e.g.
                         setting the MAM queryid and starting a stanza
                         collector.
            post_cb   -- Callback to run after receiving each stanza e.g.
                         stopping a MAM stanza collector in order to
                         gather results.
            results   -- The name of the interface containing the
                         query results (typically just 'substanzas').
        """
        return ResultIterator(stanza, interface, results, amount, reverse=reverse,
                              recv_interface=recv_interface, pre_cb=pre_cb,
                              post_cb=post_cb)