# -*- coding: utf-8 -*-`
"""
EXPLANATION PART :

1- Defines an abstract tree qmodel based on a Node class
2- Defines a ElementQTreeModel bases on an Element class
3- Defines a MyQModel bases on an MyElement class

Now the user can add subelements and the indexes will be automtically created
assuming the ElementQTreeModel calls `add_elements` from QTreeModel.

A - Define an alphabet ['a', 'b', 'c', 'd']
B - Set root elements to be each letter
C - When an elements is expanded, add each letter as subelements using `fetchMore`

This works and can be tested using `python abstract_tree.py`

PROBLEM PART:

For conveniance, we add a `set_selection_and_expand_to` to the tree.
If a string param is given (as `ababc`) the tree will expand until element data
matches.
If at any moment a string part cannot be found, the expand process stops.
If it is found it will select the matching element
This process is done using the `rowsInserted` signal.

This works and can be tested using `python abstract_tree.py abdc`

It works. BUT, when enabling the last item selection :
    - the next 2 root elements after the first one expanded are not displayed.
    Yet, if you expand an item at the last level or one or two upper, they
    are now shown
    - if the first thing you do is collapsing an item 3 level or higher than
    the last one expanded, Qt crashes

If you comment selection line #439, there is no issue

LEADS :
    - there seems to be an issue with the selection model. When using right
    and left arrow keys after an expansion the collapsed items are not the
    selected one.
    - Search for keyword BUG in the code for problems variations
"""
import sys

from PySide2 import QtCore, QtWidgets


# Abstract tree model based on Node class

class Node(object):

    def __init__(self, parent=None):
        """
        :param parent: the parent node
        :type parent: Node
        """
        self.parent = parent

    @property
    def row(self):
        if self.parent is not None:
            return self.parent.subnodes.index(self)
        return None

    @property
    def subnodes(self):
        raise NotImplementedError()


class AbstractQTreeModel(QtCore.QAbstractItemModel):
    """
    from https://www.hardcoded.net/articles/using_qtreeview_with_qabstractitemmodel  # noqa
    """

    def __init__(self):
        QtCore.QAbstractItemModel.__init__(self)
        self.root_nodes = self._get_root_nodes()
        self.__column_count = 0

    def _get_root_nodes(self):
        raise NotImplementedError()

    def columnCount(self, parent=QtCore.QModelIndex()):
        return self.__column_count

    def data(self, index, role=QtCore.Qt.DisplayRole):
        raise NotImplementedError()

    def index(self, row, column, parent):
        try:
            if not parent.isValid():
                return self.createIndex(row, column, self.root_nodes[row])
            # return normal behaviour
            parent_node = parent.internalPointer()
            return self.createIndex(row, column, parent_node.subnodes[row])
        except IndexError:
            return QtCore.QModelIndex()

    def parent(self, index):
        if not index.isValid():
            return QtCore.QModelIndex()

        node = index.internalPointer()
        if node.parent is None:
            return QtCore.QModelIndex()
        else:
            try:
                if node.parent.parent is None:
                    parent_row = self.root_nodes.index(node.parent)
                else:
                    parent_row = node.parent.row
            except ValueError:
                return QtCore.QModelIndex()

            return self.createIndex(parent_row, 0, node.parent)

    def rowCount(self, parent):
        if not parent.isValid():
            return len(self.root_nodes)

        node = parent.internalPointer()
        return len(node.subnodes)

    def update_root_nodes(self):
        self.root_nodes = self._get_root_nodes()


# Element QTreeModel based on AbstractQTreeModel (using nodes)

class Element(Node):
    """
    The internal structure of a UnitModel.
    """

    class _SubElements(list):
        """
        Override a list
        """

        def __init__(self, iterable, parent):
            """
            """
            self.__parent = parent
            list.__init__(self, iterable)

        def append(self, element):
            """
            set parent when an item is appended
            """
            element.parent = self.__parent
            list.append(self, element)

        def extend(self, elements):
            """
            set parent when an items are appended
            """
            for element in elements:
                element.parent = self.__parent
                if element not in self:
                    self.append(element)

        def __setitem__(self, key, element):
            """
            set parent when an item is directly set
            """
            element.parent = self.parent

    def __init__(
            self,
            parent=None,
            subelements=None):
        """
        :param parent: the parent node
        :type parent: Node
        :param subelements: the children elements
        :type subelements: list
        """
        self._subelements = Element._SubElements(subelements or list(), self)
        super(Element, self).__init__(parent=parent)

    @property
    def subelements(self):
        return self._subelements

    @subelements.setter
    def subelements(self, value):
        assert isinstance(value, list)
        self._subelements = Element._SubElements(value, self)

    @property
    def subelement_count(self):
        return len(self._subelements)

    @property
    def subnodes(self):
        """
        Reimplementing to compatible with our tree QModel
        """
        return self.subelements


class ElementQTreeModel(AbstractQTreeModel):
    """
    The model for a standard tree based on elements
    Elements and their subelements will be represented as nodes of
    the tree.
    """

    def __init__(self, root_elements=None):
        """
        :param root_elements: the root elements of the tree
        :type root_elements: list
        """
        self._root_elements = root_elements or list()
        super(ElementQTreeModel, self).__init__()

    def _get_root_nodes(self):
        return self.root_elements

    def add_elements(self, parent_index, elements):
        """
        Adds elements to the parent at the given index.

        :param parent_index: the index of the parent
        :type parent_index: QModelIndex
        :param elements: the elements to add
        :type elements: list
        """
        current_row_count = self.rowCount(parent=parent_index)
        self.beginInsertRows(
            parent_index,
            current_row_count,
            current_row_count + len(elements) - 1)

        if not parent_index.isValid():
            self.root_elements.extend(elements)
            self.update_root_nodes()
        else:
            parent_element = self.get_element_at_index(parent_index)
            parent_element.subelements.extend(elements)

        self.endInsertRows()

        return None

    def data(self, index, role=QtCore.Qt.DisplayRole):
        return None

    def get_element_at_index(self, index):
        """
        Retrieves the element for the item at the given index

        :param index: the index of the item
        :type index: QModelIndex
        :returns: the tree element
        :rtype: Element
        """
        return index.internalPointer() if index.isValid() else None

    @property
    def root_elements(self):
        """
        :param root_elements: root tree elements
        :type root_elements: list
        """
        return self._root_elements

    @root_elements.setter
    def root_elements(self, value):
        """
        - Clear all current elements
        - Sets the root elements of the tree
        - Update the model internal data

        :param root_elements: root tree elements
        :type root_elements: list
        """
        self._root_elements = value
        self.update_root_nodes()

# IMPLEMENTATION


ALPHABET = ['a', 'b', 'c', 'd']


class MyElement(Element):
    """
    The internal structure of my tree.
    """
    def __init__(
            self,
            data,
            parent=None,
            subelements=None):
        """
        :param data: the element infos
        :type data: any
        :param parent: the element parent
        :type parent: Element
        :param subelements: the children elements
        :type subelements: list
        """
        self._data = data
        super(MyElement, self).__init__(
            parent=parent,
            subelements=subelements)

    @property
    def data(self):
        value = ''

        # add parent
        if self.parent:
            value += self.parent.data

        value += self._data
        return value


class MyQModel(ElementQTreeModel):

    def __init__(self, root_elements=None):
        """
        :param root_elements: the root elements of the tree
        :type root_elements: list
        """
        super(MyQModel, self).__init__(root_elements=root_elements or list())

    def canFetchMore(self, index):
        if not index.isValid():
            return False

        element = self.get_element_at_index(index)
        return element.subelement_count < 4

    def columnCount(self, parent=QtCore.QModelIndex()):
        return 1

    def data(self, index, role):
        if not index.isValid():
            return None

        element = self.get_element_at_index(index)

        if role == QtCore.Qt.DisplayRole and index.column() == 0:
            return element.data

        return super(MyQModel, self).data(index, role)

    def fetchMore(self, parent_index):
        parent_element = self.get_element_at_index(parent_index)

        # either elements one by one or as as a whole
        # As a whole seems unstable

        # VISUAL BUG WHEN UNCOMMENTED AND LOWERT PART UNCOMMENTED
        for char in ALPHABET:
            self.add_elements(
                parent_index,
                [MyElement(
                    char,
                    parent=parent_element,
                    subelements=list()), ])

        # VISUAL BUG WHEN UNCOMMENTED AND UPPER PART COMMENTED
        # elements_to_add = list()
        # for char in ALPHABET:
        #     elements_to_add.append(MyElement(
        #         char,
        #         parent=parent_element))
        # self.add_elements(parent_index, elements_to_add)

    def hasChildren(self, index):
        """
        Overrides to show the arrow if has possible children
        """
        return True


class Tree(QtWidgets.QWidget):

    def __init__(self, elements=None, parent=None):
        """
        :param parent: the parent of the widget
        :type parent: QtWidgets.QWidget
        """
        super(Tree, self).__init__(parent)
        self.elements = elements or list()
        self.qmodel = MyQModel(root_elements=self.elements)

        # target utils
        self.is_reaching_target = False
        self.target_part_index = None

        self.__init_ui()
        self.__init_signals()

        self.selection_model = self.view.selectionModel()

    def __init_ui(self):
        """Initializes the widget interface"""
        layout = QtWidgets.QVBoxLayout()
        self.view = QtWidgets.QTreeView()
        self.view.setModel(self.qmodel)

        layout.addWidget(self.view)
        self.setLayout(layout)

    def __init_signals(self):
        self.qmodel.rowsInserted.connect(self._on_rows_inserted)

    def _expand_to_next_target_part(self, parent_index):
        """
        :param parent_index: the index of the parent of the items inserted
        :type index: QModelIndex
        """
        print('search next target part %d:%s' % (
            self.target_part_index,
            self.target[self.target_part_index]))

        for row in range(self.qmodel.rowCount(parent_index)):
            index = self.qmodel.index(row, 0, parent=parent_index)
            element = self.qmodel.get_element_at_index(index)

            if element.data != self.target[0:self.target_part_index + 1]:
                continue

            self.target_part_index += 1

            # is this the last part ? if so, stop searching
            if self.target_part_index == len(self.target):
                print("target found")
                self.is_reaching_target = False

                # NO BUG WHEN COMMENTING THIS
                self.selection_model.select(
                    index, QtCore.QItemSelectionModel.Select)
            else:
                self.view.expand(index)
            break

    def _on_rows_inserted(self, index, *_):
        if self.is_reaching_target:

            # COLLAPSE BUG WHEN UNCOMMENTED
            # expand next item only when all item are of the alaphabet are added
            if self.qmodel.rowCount(parent=index) != len(ALPHABET):
                return

            self._expand_to_next_target_part(index)

    def set_selection_and_expand_to(self, value):
        """
        :param value: the value to match
        :type value: str
        """
        if not value:
            return

        self.target = value
        self.is_reaching_target = True
        self.target_part_index = 0

        # expand
        self._expand_to_next_target_part(QtCore.QModelIndex())


app = QtWidgets.QApplication(sys.argv)

# build root elements
elements = []
for char in ALPHABET:
    elements.append(MyElement(char))

# build tree
widget = Tree(elements=elements)
widget.setMinimumHeight(400)
widget.show()

# expand is param is given
if len(sys.argv) == 2 and isinstance(sys.argv[-1], str):
    widget.set_selection_and_expand_to(sys.argv[-1])

sys.exit(app.exec_())