"""
Reference components that might be generally useful.
"""

from zope.interface import implements
from zope.interface import implementedBy
from zope.component import getUtility

from zope.component.interfaces import IFactory
from archetypes.uid.interfaces import IUID
from archetypes.uid.interfaces import IUIDQuery
from archetypes.reference.interfaces import IReferenceQuery
from archetypes.reference.interfaces import IReferenceStorage
from archetypes.reference.interfaces import IReferenceSource
from archetypes.reference.interfaces import IReferenceTarget

from BTrees.OOBTree import OOBTree
from BTrees.OOBTree import OOSet
from BTrees.OOBTree import union
from BTrees.OOBTree import intersection


class ReferenceSource(object):
    """A base implementation of an IReferenceSource adapter that
    relies on a IReferenceQuery local utility to look up targets and
    on an IReferenceStorage adapter to add references.

    Basic setup:

        >>> import archetypes.testing
        >>> archetypes.testing.setupPortal()
        >>> archetypes.testing.wireUp()

    We create an object that we adapt to our ReferenceSource:

        >>> obj = archetypes.testing.FakeObject('sourceobj')
        >>> source = ReferenceSource(obj)

    We don't have any targets or relationships at this point:

        >>> source.getTargets()
        []
        >>> source.getRelationships()
        set([])

    Let's add a target:

        >>> target = archetypes.testing.FakeObject('targetobj')
        >>> ref = source.addReference(target, 'pooh')
        >>> ref
        <archetypes.testing.FakeReference object at ...>
        >>> ref.source is obj
        True
        >>> ref.target is target
        True

    When we 'getTargets' now we retrieve the target that we just
    added.  Also, 'getRelationships' returns our relationship now:

        >>> source.getTargets()
        [<archetypes.testing.FakeObject object at ...>]
        >>> source.getTargets()[0] is target
        True
        >>> source.getRelationships()
        set(['pooh'])

    The target should still be there if we readapt our source object:

        >>> source = ReferenceSource(obj)
        >>> source.getTargets()
        [<archetypes.testing.FakeObject object at ...>]
        >>> source.getTargets()[0] is target
        True
        
    Teardown:

        >>> archetypes.testing.teardownPortal()
    """
    implements(IReferenceSource)

    def __init__(self, context):
        self.context = context
        self.query = getUtility(IReferenceQuery, context=context)

    def getTargets(self, relationship=None):
        refs = self.query(source=self.context, relationship=relationship)
        return [r.target for r in refs]

    def addReference(self, target, relationship):
        storage = IReferenceStorage(self.context)
        ref = storage.insert(self.context, target, relationship)
        return ref

    def deleteReferences(self, target=None, relationship=None):
        storage = IReferenceStorage(self.context)
        refs = self.query(source=self.context, target=target,
                          relationship=relationship)

        for reference in refs:
            storage.remove(reference)

    def getRelationships(self):
        relationships = set()

        for ref in self.query(source=self.context):
            relationships.add(ref.relationship)

        return relationships


class ReferenceTarget(object):
    """A base implementation of an IReferenceTarget adapter that
    relies entirely on a IReferenceQuery local utility to look up
    source object.

    Basic setup:

        >>> import archetypes.testing
        >>> archetypes.testing.setupPortal()
        >>> archetypes.testing.wireUp()

    Let's retrieve the utility that our ReferenceTarget uses to look
    up sources:
        
        >>> query = getUtility(IReferenceQuery)

    We create an object that we adapt to  our ReferenceTarget:

        >>> obj = archetypes.testing.FakeObject('targetobj')
        >>> target = ReferenceTarget(obj)

    Nothing is pointing to target yet:

        >>> target.getSources()
        []

    Let's create a reference from a source to our object 'obj' and
    index it with the IReferenceQuery utility:

        >>> source = archetypes.testing.FakeObject('sourceobj')
        >>> ref = archetypes.testing.FakeReference(source, obj, 'pooh')
        >>> query.indexObject(ref)

    Now source is pointing to us:

        >>> target.getSources()
        [<archetypes.testing.FakeObject object at ...>]
        >>> target.getSources('pooh')
        [<archetypes.testing.FakeObject object at ...>]
        >>> target.getSources('other')
        []

    The source should still be there if we readapt our target object:

        >>> target = ReferenceTarget(obj)
        >>> target.getSources()
        [<archetypes.testing.FakeObject object at ...>]
    
    Teardown:

        >>> archetypes.testing.teardownPortal()
    """
    implements(IReferenceTarget)

    def __init__(self, context):
        self.context = context
        self.query = getUtility(IReferenceQuery, context=context)

    def getSources(self, relationship=None):
        refs = self.query(target=self.context, relationship=relationship)
        return [ref.source for ref in refs]


class ReferenceCatalogQuery(object):
    """An implementation of a IReferenceQuery utility that indexes and
    unindexes reference objects on creation / deletion using OOBtrees.

    Basic setup:

        >>> import archetypes.testing
        >>> archetypes.testing.setupPortal()
        >>> archetypes.testing.wireUp()

    Let's create our ReferenceCatalogQuery, which is of course empty
    just afterwards:

        >>> query = ReferenceCatalogQuery()
        >>> query()
        []

    Now we create source and target and make a reference:

        >>> source = archetypes.testing.FakeObject('sourceobj')
        >>> target = archetypes.testing.FakeObject('targetobj')
        >>> ref = archetypes.testing.FakeReference(
        ...     source, target, 'pooh')

    Normally, a reference would be added by an IReferenceStorage and
    then an event would call our ReferenceCatalogQuery's indexObject.
    We need to call our indexObject manually, however:

        >>> query.indexObject(ref)

    Now we can do all kinds of funny queries to retrieve our
    reference:

        >>> refs = query()
        >>> refs
        [<archetypes.testing.FakeReference object at ...>]
        >>> refs[0] is ref
        True
        >>> query(source=source) == [ref]
        True
        >>> query(source=source, target=target) == [ref]
        True
        >>> query(source=source, target=target, relationship='pooh') == [ref]
        True
        >>> query(target=target) == [ref]
        True
        >>> query(target=target, relationship='pooh') == [ref]
        True
        >>> query(relationship='pooh') == [ref]
        True

    Some queries that don't return our reference:

        >>> query(relationship='foo')
        []
        >>> query(source=target)
        []
        >>> query(target=source)
        []
        >>> query(source=source, target=source)
        []
        >>> query(source=target, target=target)
        []
        >>> query(source=source, target=target, relationship='foo')
        []
        
    Passing an unsupported kwarg should raise a TypeError:
        
        >>> query(source=source, target=target, blah='blah')
        Traceback (most recent call last):
        ...
        TypeError: Only source, target, and relationship are supported arguments.

    Let's add another reference and see how query now returns two
    references:

        >>> ref2 = archetypes.testing.FakeReference(
        ...     target, source, 'foo')

        >>> query.indexObject(ref2)
        >>> refs = query()
        >>> ref2 in refs
        True
        >>> ref in refs
        True
        >>> len(refs)
        2
        >>> query(relationship='foo') == [ref2]
        True

    After we unindex our references, we won't retrieve any references:

        >>> query.unindexObject(ref)
        >>> len(query())
        1
        >>> query.unindexObject(ref2)
        >>> len(query())
        0

    Teardown:

        >>> archetypes.testing.teardownPortal()
    """
    implements(IReferenceQuery)

    def __init__(self):
        # Make three "indexes" to query on
        self.sources = OOBTree()
        self.targets = OOBTree()
        self.relationships = OOBTree()

    def __call__(self, source=None, target=None, relationship=None, **kwargs):
        if kwargs: raise TypeError("Only source, target, and relationship are"
                                   " supported arguments.")
        source_match = target_match = rel_match = None

        if source is not None:
            source_match = self.sources.get(IUID(source)(), OOSet())

        if target is not None:
            target_match = self.targets.get(IUID(target)(), OOSet())

        if relationship is not None:
            rel_match = self.relationships.get(relationship, OOSet())

        rids = intersection(source_match, target_match)
        rids = intersection(rids, rel_match) or OOSet()

        if source is None and target is None and relationship is None:
            rid_sets = self.relationships.values()
            rids = OOSet()
            for set in rid_sets:
                rids = union(rids, set)

        uid_query = getUtility(IUIDQuery, context=self)
        
        return [uid_query.getObject(rid) for rid in rids]

    def indexObject(self, reference):
        """ Stores a reference object using it's UID, assumes we have
            UID providing IReference objects.
        """
        r_uid = IUID(reference)()
        s_uid = IUID(reference.source)()
        t_uid = IUID(reference.target)()
        relationship = reference.relationship

        sources = self.sources.get(s_uid, None)
        if sources is None:
            sources = self.sources[s_uid] = OOSet()
        sources.insert(r_uid)

        targets = self.targets.get(t_uid, None)
        if targets is None:
            targets = self.targets[t_uid] = OOSet()
        targets.insert(r_uid)

        relationships = self.relationships.get(relationship, None)
        if relationships is None:
            relationships = self.relationships[relationship] = OOSet()
        relationships.insert(r_uid)

    def unindexObject(self, reference):
        """ Remove the reference uid from our indexes """
        r_uid = IUID(reference)()
        s_uid = IUID(reference.source)()
        t_uid = IUID(reference.target)()

        # If the index entry corresponding to our source, target, or
        # relationship is empty, then we remove the entry.
        self.sources[s_uid].remove(r_uid)
        if not self.sources[s_uid]:
            del self.sources[s_uid]

        self.targets[t_uid].remove(r_uid)
        if not self.targets[t_uid]:
            del self.targets[t_uid]

        self.relationships[reference.relationship].remove(r_uid)
        if not self.relationships[reference.relationship]:
            del self.relationships[reference.relationship]