--- a/src/plugins/plugin_xep_0313.py	Fri Jan 09 10:50:11 2015 +0100
+++ b/src/plugins/plugin_xep_0313.py	Fri Jan 09 10:51:12 2015 +0100
@@ -23,17 +23,18 @@
 from sat.core.log import getLogger
 log = getLogger(__name__)
 
-from wokkel import disco, iwokkel, compat, data_form
-from wokkel.rsm import RSMRequest
-from wokkel.generic import parseXml
 try:
     from twisted.words.protocols.xmlstream import XMPPHandler
 except ImportError:
     from wokkel.subprotocols import XMPPHandler
 from twisted.words.xish import domish
+from twisted.words.protocols.jabber import jid
+
 from zope.interface import implements
 
-from dateutil.tz import tzutc
+from wokkel import disco, data_form, mam
+from wokkel.rsm import RSMRequest
+from wokkel.generic import parseXml
 
 
 NS_MAM = 'urn:xmpp:mam:0'
@@ -59,92 +60,120 @@
     def __init__(self, host):
         log.info(_("Message Archive Management plugin initialization"))
         self.host = host
-        host.bridge.addMethod("MAMqueryFields", ".plugin", in_sign='s', out_sign='s',
-                              method=self.queryFields,
+        self.clients = {}  # bind profile name to SatMAMClient
+        host.bridge.addMethod("MAMqueryFields", ".plugin", in_sign='ss', out_sign='s',
+                              method=self._queryFields,
                               async=True,
                               doc={})
-        host.bridge.addMethod("MAMqueryArchive", ".plugin", in_sign='ssss', out_sign='s',
+        host.bridge.addMethod("MAMqueryArchive", ".plugin", in_sign='ssa{ss}ss', out_sign='s',
                               method=self._queryArchive,
                               async=True,
                               doc={})
+        host.bridge.addMethod("MAMgetPrefs", ".plugin", in_sign='ss', out_sign='s',
+                              method=self._getPrefs,
+                              async=True,
+                              doc={})
+        host.bridge.addMethod("MAMsetPrefs", ".plugin", in_sign='ssasass', out_sign='s',
+                              method=self._setPrefs,
+                              async=True,
+                              doc={})
         host.trigger.add("MessageReceived", self.messageReceivedTrigger)
 
     def getHandler(self, profile):
-        return XEP_0313_handler(self, profile)
+        self.clients[profile] = SatMAMClient(self, profile)
+        return self.clients[profile]
 
-    def _queryArchive(self, form=None, rsm=None, node=None, profile_key=C.PROF_KEY_NONE):
-        form_elt = parseXml(form) if form else None
-        rsm_inst = RSMRequest(**rsm) if rsm else None
-        return self.queryArchive(form_elt, rsm_inst, node, profile_key)
+    def profileDisconnected(self, profile):
+        try:
+            del self.clients[profile]
+        except KeyError:
+            pass
 
-    def queryArchive(self, form=None, rsm=None, node=None, profile_key=C.PROF_KEY_NONE):
-        """Query a user, MUC or pubsub archive.
+    def _queryFields(self, service_s=None, profile_key=C.PROF_KEY_NONE):
+        service = jid.JID(service_s) if service_s else None
+        return self.queryFields(service, profile_key)
 
-        @param form (domish.Element): data form to filter the request
-        @param rsm (RSMRequest): RSM request instance
-        @param node (unicode): pubsub node to query, or None if inappropriate
+    def queryFields(self, service=None, profile_key=C.PROF_KEY_NONE):
+        """Ask the server about additional supported fields.
+
+        @param service: entity offering the MAM service (None for user archives)
         @param profile_key (unicode): %(doc_profile_key)s
-        @return: a Deferred when the message has been sent
+        @return: the server response as a Deferred domish.Element
         """
-        client = self.host.getClient(profile_key)
-        iq = compat.IQ(client.xmlstream, 'set')
-        query_elt = iq.addElement((NS_MAM, 'query'))
-        if form:
-            query_elt.addChild(form)
-        if rsm:
-            rsm.render(query_elt)
-        if node:
-            query_elt['node'] = node
-        d = iq.send()
-
+        # http://xmpp.org/extensions/xep-0313.html#query-form
         def eb(failure):
             # typically StanzaError with condition u'service-unavailable'
             log.error(failure.getErrorMessage())
             return ''
 
+        profile = self.host.memory.getProfileName(profile_key)
+        d = self.clients[profile].queryFields(service)
         return d.addCallbacks(lambda elt: elt.toXml(), eb)
 
-    def queryFields(self, profile_key=C.PROF_KEY_NONE):
-        """Ask the server about additional supported fields.
+    def _queryArchive(self, service_s=None, form_xml=None, rsm_dict=None, node=None, profile_key=C.PROF_KEY_NONE):
+        service = jid.JID(service_s) if service_s else None
+        if form_xml:
+            form = data_form.Form.fromElement(parseXml(form_xml))
+            if form.formNamespace != NS_MAM:
+                log.error(_("Expected a MAM Data Form, got instead: %s") % form.formNamespace)
+                form = None
+        else:
+            form = None
+        rsm = RSMRequest(**rsm_dict) if rsm_dict else None
+        return self.queryArchive(service, form, rsm, node, profile_key)
 
+    def queryArchive(self, service=None, form=None, rsm=None, node=None, profile_key=C.PROF_KEY_NONE):
+        """Query a user, MUC or pubsub archive.
+
+        @param service: entity offering the MAM service (None for user archives)
+        @param form (Form): data form to filter the request
+        @param rsm (RSMRequest): RSM request instance
+        @param node (unicode): pubsub node to query, or None if inappropriate
         @param profile_key (unicode): %(doc_profile_key)s
-        @return: the server response as a Deferred domish.Element
+        @return: a Deferred when the message has been sent
         """
-        # http://xmpp.org/extensions/xep-0313.html#query-form
-        client = self.host.getClient(profile_key)
-        iq = compat.IQ(client.xmlstream, 'get')
-        iq.addElement((NS_MAM, 'query'))
-        d = iq.send()
-
         def eb(failure):
             # typically StanzaError with condition u'service-unavailable'
             log.error(failure.getErrorMessage())
             return ''
 
+        profile = self.host.memory.getProfileName(profile_key)
+        d = self.clients[profile].queryArchive(service, form, rsm, node)
         return d.addCallbacks(lambda elt: elt.toXml(), eb)
+        # TODO: add the handler for receiving the final message
 
-    def queryPrefs(self, profile_key=C.PROF_KEY_NONE):
+    def _getPrefs(self, service_s=None, profile_key=C.PROF_KEY_NONE):
+        service = jid.JID(service_s) if service_s else None
+        return self.getPrefs(service, profile_key)
+
+    def getPrefs(self, service=None, profile_key=C.PROF_KEY_NONE):
         """Retrieve the current user preferences.
 
+        @param service: entity offering the MAM service (None for user archives)
         @param profile_key (unicode): %(doc_profile_key)s
         @return: the server response as a Deferred domish.Element
         """
         # http://xmpp.org/extensions/xep-0313.html#prefs
-        client = self.host.getClient(profile_key)
-        iq = compat.IQ(client.xmlstream, 'get')
-        iq.addElement((NS_MAM, 'prefs'))
-        d = iq.send()
-
         def eb(failure):
             # typically StanzaError with condition u'service-unavailable'
             log.error(failure.getErrorMessage())
             return ''
 
+        profile = self.host.memory.getProfileName(profile_key)
+        d = self.clients[profile].queryPrefs(service)
         return d.addCallbacks(lambda elt: elt.toXml(), eb)
 
-    def setPrefs(self, default='roster', always=None, never=None, profile_key=C.PROF_KEY_NONE):
+    def _setPrefs(self, service_s=None, default='roster', always=None, never=None, profile_key=C.PROF_KEY_NONE):
+        service = jid.JID(service_s) if service_s else None
+        always_jid = [jid.JID(entity) for entity in always]
+        never_jid = [jid.JID(entity) for entity in never]
+        #TODO: why not build here a MAMPrefs object instead of passing the args separately?
+        return self.setPrefs(service, default, always_jid, never_jid, profile_key)
+
+    def setPrefs(self, service=None, default='roster', always=None, never=None, profile_key=C.PROF_KEY_NONE):
         """Set news user preferences.
 
+        @param service: entity offering the MAM service (None for user archives)
         @param default (unicode): a value in ('always', 'never', 'roster')
         @param always (list): a list of JID instances
         @param never (list): a list of JID instances
@@ -152,61 +181,15 @@
         @return: the server response as a Deferred domish.Element
         """
         # http://xmpp.org/extensions/xep-0313.html#prefs
-        assert(default in ('always', 'never', 'roster'))
-        client = self.host.getClient(profile_key)
-        iq = compat.IQ(client.xmlstream, 'set')
-        prefs = iq.addElement((NS_MAM, 'prefs'))
-        prefs['default'] = default
-
-        for var, attr in ((always, 'always'), (never, 'never')):
-            if var is not None:
-                elt = prefs.addElement((None, attr))
-                for entity in var:
-                    elt.addElement((None, 'jid')).addContent(entity.full())
-        d = iq.send()
-
         def eb(failure):
             # typically StanzaError with condition u'service-unavailable'
             log.error(failure.getErrorMessage())
             return ''
 
+        profile = self.host.memory.getProfileName(profile_key)
+        d = self.clients[profile].setPrefs(service, default, always, never)
         return d.addCallbacks(lambda elt: elt.toXml(), eb)
 
-    @classmethod
-    def datetime2utc(cls, datetime):
-        """Convert a datetime to a XEP-0082 compliant UTC datetime.
-
-        @param datetime (datetime): offset-aware timestamp to convert.
-        @return: unicode
-        """
-        # receipt from wokkel.delay.Delay.toElement
-        stampFormat = '%Y-%m-%dT%H:%M:%SZ'
-        return datetime.astimezone(tzutc()).strftime(stampFormat)
-
-    @classmethod
-    def buildForm(cls, start=None, end=None, with_jid=None, extra=None):
-        """Prepare a data form for MAM query.
-
-        @param start (datetime): offset-aware timestamp to filter out older messages.
-        @param end (datetime): offset-aware timestamp to filter out later messages.
-        @param with_jid (JID): JID against which to match messages.
-        @param extra (list): list of extra fields that are not defined by the
-            specification. Each element must be a 3-tuple containing the field
-            type, name and value.
-        @return: a XEP-0004 data form as domish.Element
-        """
-        form = data_form.Form('submit', formNamespace=NS_MAM)
-        if start:
-            form.addField(data_form.Field('text-single', 'start', XEP_0313.datetime2utc(start)))
-        if end:
-            form.addField(data_form.Field('text-single', 'end', XEP_0313.datetime2utc(end)))
-        if with_jid:
-            form.addField(data_form.Field('jid-single', 'with', with_jid.full()))
-        if extra is not None:
-            for field in extra:
-                form.addField(data_form.Field(*field))
-        return form.toElement()
-
     def messageReceivedTrigger(self, message, post_treat, profile):
         """Check if the message is a MAM result. If so, extract the original
         message, stop processing the current message and process the original
@@ -234,13 +217,14 @@
         return False
 
 
-class XEP_0313_handler(XMPPHandler):
-    implements(iwokkel.IDisco)
+class SatMAMClient(mam.MAMClient):
+    implements(disco.IDisco)
 
     def __init__(self, plugin_parent, profile):
         self.plugin_parent = plugin_parent
         self.host = plugin_parent.host
         self.profile = profile
+        mam.MAMClient.__init__(self)
 
     def getDiscoInfo(self, requestor, target, nodeIdentifier=''):
         return [disco.DiscoFeature(NS_MAM)]