"""
open/dulcinea/lib/partition.py
"""

from durus.btree import BTree
from durus.persistent import PersistentObject
from qp.lib.keep import Keyed, Counter
from qp.lib.spec import require, get_spec_problems, anything, mapping, spec
from qp.lib.spec import add_getters, Specified, Mixin

class PartKeyed (Keyed, Mixin):

    part_key_is = anything

    def __init__(self):
        assert isinstance(self, PersistentObject)
        Keyed.__init__(self)
        self.part_key = None

    def get_part_key(self):
        return self.part_key

class Partition (PersistentObject, Specified):

    value_spec_is = spec(
        anything,
        "Specifies the type of values allowed in mappings")
    part_key_spec_is = spec(
        anything,
        "Specifies the type of keys allowed in the partition mapping")
    parts_is = mapping({anything:BTree}, BTree)
    key_counter_is = Counter

    def __init__(self, value_spec=PartKeyed, part_key_spec=anything):
        self.parts = BTree()
        self.key_counter = Counter()
        self.value_spec = value_spec
        self.part_key_spec = part_key_spec

    def get(self, key, default=None):
        for part in self.parts.itervalues():
            value = part.get(key, None)
            if value is not None:
                return value
        return default

    def __nonzero__(self):
        return self.parts.__nonzero__()
    __bool__ = __nonzero__

    def get_max_item(self):
        """() -> (key:anything, value:anything)
        Return the part item with the maximal key.
        """
        maximal = (0, None)
        for part in self.parts.iteritems():
            part_max_item = part[1].get_max_item()
            if part_max_item[0] > maximal[0]:
                maximal = part_max_item
        assert maximal != (0, None), "empty Partition has no max item"
        return maximal

    def itervalues(self):
        for part in self.parts.itervalues():
            for value in part.itervalues():
                yield value

    def iterkeys(self):
        for part in self.parts.itervalues():
            for key in part.iterkeys():
                yield key

    def iteritems(self):
        for part in self.parts.itervalues():
            for item in part.iteritems():
                yield item

    def get_part(self, part_key, default=None):
        return self.parts.get(part_key, default)

    def remove(self, value):
        part = self.parts[value.part_key]
        del part[value.key]

    def discard(self, value):
        if value.part_key is not None:
            part = self.parts.get(value.part_key)
            if part and value.key in part:
                self.remove(value)

    def move(self, value, part_key):
        if value.get_part_key() != part_key:
            require(part_key, self.part_key_spec)
            self.discard(value)
            assert self.get(value.key) is None
            value.part_key = part_key
            if get_spec_problems(value):
                raise TypeError(''.join(get_spec_problems(value)))
            if part_key not in self.parts:
                self.parts[part_key] = BTree()
            self.parts[part_key][value.key] = value

    def add(self, value, part_key):
        require(value, self.value_spec)
        assert value.key is None
        assert value.part_key is None
        value.key = self.key_counter.next()
        self.move(value, part_key)

add_getters(Partition)