# -*- coding: utf-8 -*-

class TreeIter:

    def __init__(self, root, collapsed=[]):

        self._node = root

        self._index = 0

        self._stack = []

        self._collapsed = collapsed

    def __next__(self):

        node = None

        while node is None:

            try:

                if self._node in self._collapsed:

                    raise IndexError()

                node = self._node[self._index]

                if node is None:

                    raise Exception('Bad node: None')

            except IndexError:

                if len(self._stack):

                    self._node, self._index = self._stack.pop()

                else:

                    raise StopIteration

        level = len(self._stack) + 1

        index = self._index + 1

        count = len(self._node)

        self._index += 1

        self._stack.append((self._node, self._index))

        self._node = node

        self._index = 0

        return (level, index, count, node)

class TreeNode(list):

    def __init__(self, parent=None, name=''):

        list.__init__(self)

        self.parent = parent

        self.name = name

    def __eq__(self, other):

        # do not compare by list content

        return self is other

class TreeModel(EventSource):

    def __init__(self):

        EventSource.__init__(self)

        self.addevents('change')

        self.root = TreeNode()

    def __iter__(self):

        return TreeIter(self.root)

    def find(self, path):

        if isinstance(path, str):

            path = path.split('/')

        # strip empty strings from both ends

        while path and path[0] == '':

            del path[0]

        while path and path[-1] == '':

            del path[-1]

        node = self.root

        for item in path:

            if isinstance(item, int):

                node = node[item]

            else:

                found = False

                for subnode in node:

                    if subnode.name == item:

                        node = subnode

                        found = True

                        break

                if not found:

                    item = int(item)

                    node = node[item]

        return node

    def add(self, path, names):

        node = self.find(path)

        if isinstance(names, str):

            names = (names,)

        for name in names:

            node.append(TreeNode(node, name))

        self.emit('change')

class TreeView(Widget):

    def __init__(self, model=None, width=20, height=20):

        Widget.__init__(self, width, height)

        # cursor

        self.cnode = None

        # model

        self._model = None

        self.setmodel(model)

        self.collapsed = []

        self.connect('draw', self.on_draw)

        self.connect('keypress', self.on_keypress)

    def __iter__(self):

        return TreeIter(self._model.root, self.collapsed)

    def getmodel(self):

        '''TreeModel in use by this TreeView.'''

        return self._model

    def setmodel(self, value):

        if self._model:

            self._model.disconnect('change', self.redraw)

        self._model = value

        if self._model:

            self._model.connect('change', self.redraw)

            try:

                self.cnode = self._model.root[0]

            except IndexError:

                pass

    model = property(getmodel, setmodel)

    def collapse(self, path, collapse=True):

        node = self._model.find(path)

        self.collapse_node(node, collapse)

    def collapse_node(self, node, collapse=True):

        if collapse:

            if not node in self.collapsed and len(node) > 0:

                self.collapsed.append(node)

        else:

            try:

                self.collapsed.remove(node)

            except ValueError:

                pass

    def on_draw(self, screen, x, y):

        screen.pushcolor('normal')

        lines = 0  # bit array, bit 0 - draw vertical line on first column, etc. 

        for level, index, count, node in self:

            # prepare string with vertical lines where they should be

            head = []

            for l in range(level-1):

                if lines & (1 << l):

                    head.append(screen.VLINE + ' ')

                else:

                    head.append('  ')

            # add vertical line if needed

            if index < count:

                head.append(screen.LTEE)

                lines |= 1 << level-1

            else:

                head.append(screen.LLCORNER)

                lines &= ~(1 << level-1)

            # draw lines and name

            head = ''.join(head)

            if node in self.collapsed:

                sep = '+'

            else:

                sep = ' '

            screen.puts(x, y, head + sep + node.name)

            if node is self.cnode:

                screen.pushcolor('active')

                screen.puts(x + len(head), y, sep + node.name + ' ')

                screen.popcolor()

            y += 1

        screen.popcolor()        

    def on_keypress(self, keyname, char):

        if keyname:

            if keyname == 'up':    self.move_up()

            if keyname == 'down':  self.move_down()

            if keyname == 'left':  self.move_left()

            if keyname == 'right': self.move_right()

        self.redraw()

    def prev_node(self, node):

        # previous sibling

        parent = node.parent

        i = parent.index(node)

        if i > 0:

            node = parent[i-1]

            while node not in self.collapsed and len(node) > 0:

                node = node[-1]

            return node

        else:

            if parent.parent is None:

                return None

            return parent

    def next_node(self, node):

        if node in self.collapsed or len(node) == 0:

            # next sibling

            parent = node.parent

            while parent is not None:

                i = parent.index(node)

                try:

                    return parent[i+1]

                except IndexError:

                    node = parent

                    parent = node.parent

            return None

        else:

            # first child

            return node[0]

    def move_up(self):

        prev = self.prev_node(self.cnode)

        if prev is not None:

            self.cnode = prev

            return True

        return False

    def move_down(self):

        next = self.next_node(self.cnode)

        if next is not None:

            self.cnode = next

            return True

        return False

    def move_left(self):

        self.collapse_node(self.cnode, True)

    def move_right(self):

        self.collapse_node(self.cnode, False)