# -*- coding: utf-8 -*-
import json
from pathlib import Path
from PyQt5 import QtWidgets, QtCore, QtGui

class FileList(QtWidgets.QWidget):

    childrenPasteSignal = QtCore.pyqtSignal(list)
    childrenCopiedSignal = QtCore.pyqtSignal(list, list, list, list)
    childAppended = QtCore.pyqtSignal(QtCore.QModelIndex)
    tableClicked = QtCore.pyqtSignal(QtCore.QModelIndex)
    listUpdated = QtCore.pyqtSignal()
    fileDropped = QtCore.pyqtSignal(list)
    filePixmapNoneSignal = QtCore.pyqtSignal(QtCore.QModelIndex)

    def __init__(self, parent=None, path=None):
        super(FileList, self).__init__(parent)
        self.setLayout( QtWidgets.QVBoxLayout(self) )
        self.model = QtGui.QStandardItemModel(self)
        self.model.setHorizontalHeaderLabels(['', '', '', ''])
        self.view = TreeView(self)
        self.view.setModel(self.model)
        self.view.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
        self.view.setItemDelegate(Delegate())
        self.view.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        self.view.customContextMenuRequested.connect(self.customContextMenuEvent)

        self.layout().setContentsMargins(0, 0, 0, 0)
        self.layout().addWidget(self.view)

        self.view.clicked.connect( lambda index : self.tableClicked.emit(index) )
        self.view.keyPressed.connect(self.viewKeyPressed)
        self.setAcceptDrops(True)

    def appendChild(self, parentIndex, text, rectF, pen, penText):
        rectItem = QtGui.QStandardItem(text)
        rectItem.setData(rectF)
        rectItem.setCheckable(True)
        rectItem.setCheckState(QtCore.Qt.Checked)

        parentPixmap = self.model.item(parentIndex.row(), 1).data()
        if parentPixmap is None:
            self.filePixmapNoneSignal.emit(parentIndex)
            parentPixmap = self.model.item(parentIndex.row(), 1).data()
        cropPixmap = parentPixmap.copy(rectF.toRect())
        pixmapItem = QtGui.QStandardItem()
        pixmapItem.setData(cropPixmap)
        pixmapItem.setData(cropPixmap, QtCore.Qt.DecorationRole)
        pixmapItem.setEditable(False)

        penItem = QtGui.QStandardItem()
        penItem.setData(pen)
        penItem.setText(penText)

        parentItem = self.model.item( parentIndex.row() )
        parentItem.appendRow([rectItem, pixmapItem, QtGui.QStandardItem(), penItem])

    def appendFile(self, path):
        filename = QtGui.QStandardItem(str(path.name))
        filename.setData(path)
        image = QtGui.QStandardItem()
        edge = QtGui.QStandardItem()
        self.model.appendRow([ filename, image, edge ])

    def childrenCopied(self, selectedIndex):
        selectedItem = self.model.itemFromIndex(selectedIndex)
        datas, texts, pens, penTexts = [], [], [], []
        for row in range( selectedItem.rowCount() ):
            datas.append(    selectedItem.child(row, 0).data() )
            texts.append(    selectedItem.child(row, 0).text() )
            pens.append(     selectedItem.child(row, 3).data() )
            penTexts.append( selectedItem.child(row, 3).text() )
        self.childrenCopiedSignal.emit(datas, texts, pens, penTexts)

    def childrenExport(self, selectedIndex):
        selectedItem = self.model.itemFromIndex(selectedIndex)
        datas = []
        for row in range( selectedItem.rowCount() ):
            rectItem = selectedItem.child(row, 0)
            pen = selectedItem.child(row, 3)
            datas.append({
                'rect' : {
                    'x'      : rectItem.data().x(),
                    'y'      : rectItem.data().y(),
                    'width'  : rectItem.data().width(),
                    'height' : rectItem.data().height(),
                    'text'   : rectItem.text()
                },
                'pen' : {
                    'r'     : pen.data().brush().color().red(),
                    'g'     : pen.data().brush().color().green(),
                    'b'     : pen.data().brush().color().blue(),
                    'a'     : pen.data().brush().color().alpha(),
                    'width' : pen.data().width(),
                    'text'  : pen.text()
                }
            })
        filepath, _ = QtWidgets.QFileDialog.getSaveFileName(None, 'Save JSON file', '', 'JSON file(*.json)')
        if filepath == '':
            return
        with open(filepath, 'w') as f:
            json.dump(datas, f, indent=4)

    def childrenImport(self, selectedIndex):
        filepath, _ = QtWidgets.QFileDialog.getOpenFileName(None, 'Open JSON file', '', 'JSON file(*.json)')
        if filepath == '':
            return
        with open(filepath, 'r') as f:
            datas = json.load(f)
        for d in datas:
            self.appendChild(
                selectedIndex, 
                d['rect']['text'], 
                QtCore.QRectF(d['rect']['x'], d['rect']['y'], d['rect']['width'], d['rect']['height']), 
                QtGui.QPen( QtGui.QBrush( QtGui.QColor(d['pen']['r'], d['pen']['g'], d['pen']['b'], d['pen']['a']) ), d['pen']['width'] ), 
                d['pen']['text']
            )

    def childrenWithText(self, text, column=0, parentIndex=QtCore.QModelIndex()):
        children = []
        selectedIndexes = self.view.selectedIndexes()
        for item in self.items(column, parentIndex):
            if not item.checkState() == QtCore.Qt.Checked:
                continue
            if not item.text() == text:
                continue
            if self.model.index(item.row(), column) in selectedIndexes:
                continue
            children.append(item)
        return children

    def customContextMenuEvent(self, point):
        selectedChildren, selectedParents = self.selectedChildrenWithParentKey()
        selectedChildrenCount = sum([ len(selectedChildren[key]) for key in selectedChildren ])
        selectedChildrenParents = list( selectedChildren.keys() )
        selectedParentCount = len(selectedParents)

        menu = QtWidgets.QMenu(self.view)

        if selectedParentCount == 1 and selectedChildrenCount == 0 and len(selectedChildrenParents) == 0:
            selectedParent = selectedParents[0]
            menu.addAction('Copy children data', lambda : self.childrenCopied(selectedParent))
            menu.addAction('Export children data', lambda : self.childrenExport(selectedParent))
            menu.addAction('Import children data', lambda : self.childrenImport(selectedParent))

        if selectedParentCount > 0 and selectedChildrenCount == 0 and len(selectedChildrenParents) == 0:
            menu.addAction('Paste children data', lambda : self.childrenPasteSignal.emit(selectedParents))

        if selectedParentCount == 0 and selectedChildrenCount > 0 and len(selectedChildrenParents) > 0:
            menu.addAction('Remove selected row', lambda : self.removeSelectedRects())

        if len([ c.text() for c in menu.children() if not c.text() == '' ]) > 0:
            menu.exec(self.view.mapToGlobal(point))

    def dragEnterEvent(self, event):
        if event.mimeData().hasUrls():
            event.accept()
        else:
            event.ignore()

    def dropEvent(self, event):
        files = [u.toLocalFile() for u in event.mimeData().urls()]
        self.fileDropped.emit(files)

    def index(self, row, column=0, parentIndex=QtCore.QModelIndex()):
        return self.model.index(row, column, parentIndex)

    def item(self, row, column=0):
        return self.model.item(row, column)

    def items(self, column=0, parentIndex=QtCore.QModelIndex()):
        if parentIndex.isValid():
            item = self.model.item( parentIndex.row() )
            return [ item.child(row, column) for row in range( item.rowCount() ) ]
        else:
            return [ self.model.item(row, column) for row in range( self.model.rowCount() ) ]
        
    def viewKeyPressed(self, event):
        selectedChildren, selectedParents = self.selectedChildrenWithParentKey()
        selectedChildrenCount = sum([ len(selectedChildren[key]) for key in selectedChildren ])
        selectedParentCount = len(selectedParents)

        if (event.modifiers() & QtCore.Qt.ControlModifier):

            if event.key() == QtCore.Qt.Key_C:
                if selectedParentCount == 1 and selectedChildrenCount == 0:
                    selectedParent = list( selectedChildren.keys() )[0]
                    self.childrenCopied(selectedParent)
                    self.listUpdated.emit()
                return

            if event.key() == QtCore.Qt.Key_V:
                if selectedParentCount > 0 and selectedChildrenCount == 0:
                    lambda : self.childrenPasteSignal.emit(selectedParents)
                    self.listUpdated.emit()
                return
                
        if event.key() == QtCore.Qt.Key_Delete:
            for selectedParent in selectedChildren:
                for childIndex in selectedChildren[selectedParent]:
                    parentItem = self.model.item(selectedParent.row())
                    parentItem.removeRow(childIndex.row())
            for selectedParent in selectedParents:
                self.model.removeRow(selectedParent.row())
            self.listUpdated.emit()
            return
        
    def rectText(self, parentIndex, baseText='Rect'):
        numbers = []
        for child in self.items(0, parentIndex):
            numberText = ''.join([ s for s in child.text() if s in [ str(i) for i in range(10)] ])
            numbers.append(int(numberText))
        if len(numbers) == 0:
            return baseText + '0'
        return baseText + str( max(numbers) + 1 )

    def removeAllChildren(self, parentIndex=None):
        if parentIndex is None:
            self.model.removeRows(0, self.model.rowCount())
        else:
            parentItem = self.item(parentIndex.row())
            for childRow in list( range(parentItem.rowCount()) )[::-1]:
                parentItem.removeRow(childRow)

    def removeSelectedRects(self):
        removeList = {}
        removedRects = []
        for index in self.view.selectedIndexes():
            parentIndex = index.parent()
            if not parentIndex.isValid():
                continue
            if not parentIndex in removeList:
                removeList[parentIndex] = {}
            parentItem = self.model.item(parentIndex.row())
            item = parentItem.child(index.row())
            removeList[parentIndex][index.row()] = item.data()
        
        for parentIndex in removeList:
            deleteRowAndRects = sorted( removeList[parentIndex].items(), reverse=True )
            deleteRowAndRects = { key : data for key, data in deleteRowAndRects }
            for deleteRow in deleteRowAndRects:
                rect = deleteRowAndRects[deleteRow]
                self.model.removeRow(deleteRow, parentIndex)
                removedRects.append(rect)

        parentIndex = list(removeList.keys())[0]
        return removedRects, parentIndex
        
    def rowCount(self, parentIndex=QtCore.QModelIndex()):
        return self.model.rowCount(parentIndex)

    def selectedChildren(self, column=0, parentIndex=QtCore.QModelIndex()):
        children = []
        selectedIndexes = [ [i.row(), i.column(), i.parent()] for i in self.view.selectedIndexes() ]
        for item in self.items(column, parentIndex):
            if not item.checkState() == QtCore.Qt.Checked:
                continue
            if not [item.row(), 0, parentIndex] in selectedIndexes:
                continue
            children.append(item)
        return children

    def selectedIndexes(self):
        return self.view.selectedIndexes()

    def selectedChildrenWithParentKey(self):
        selectedChildrenIndexes = {}
        selectedParents = []
        for index in self.view.selectedIndexes():
            if not index.column() == 0:
                continue
            parentIndex = index.parent()
            if parentIndex.isValid():
                if parentIndex in selectedChildrenIndexes:
                    if not index.row() in [ i.row() for i in selectedChildrenIndexes[parentIndex] ]:
                        selectedChildrenIndexes[parentIndex].append(index)
                else:
                    selectedChildrenIndexes[parentIndex] = [index]
            else:
                if not index.row() in [ i.row() for i in selectedParents ]:
                    selectedParents.append(index)
        return selectedChildrenIndexes, selectedParents

    def setData(self, index, text, data, state=QtCore.Qt.Checked):
        item = QtGui.QStandardItem(text)
        item.setData(data)
        item.setCheckable(True)
        item.setCheckState(state)
        self.model.setItem( index.row(), index.column(), item )
        return item

    def setFileData(self, row, column, text, data, state=QtCore.Qt.Checked):
        item = QtGui.QStandardItem(text)
        item.setData(data)
        item.setCheckable(True)
        item.setCheckState(state)
        self.model.setItem( row, column, item )
        return item
        
class TreeView(QtWidgets.QTreeView):
    keyPressed = QtCore.pyqtSignal(QtGui.QKeyEvent)
    def __init__(self, parent=None):
        super(TreeView, self).__init__(parent)

    def keyPressEvent(self, event):
        super(TreeView, self).keyPressEvent(event)
        self.keyPressed.emit(event)

class Delegate(QtWidgets.QStyledItemDelegate):
    def __init__(self, parent=None):
        super(Delegate, self).__init__(parent)

    def paint(self, painter, option, index):
        item = index.model().itemFromIndex(index)
        if type( item.data() ) is QtGui.QPixmap and not item.parent() is None:
            pixmapWidth, pixmapHeight = item.data().size().width(), item.data().size().height()
            cellWidth, cellHeight = option.rect.width(), option.rect.height()
            if cellWidth / cellHeight < pixmapWidth / pixmapHeight:
                scaledPixmap = item.data().scaledToWidth(cellWidth)
            else:
                scaledPixmap = item.data().scaledToHeight(cellHeight)
            item.setData(scaledPixmap, QtCore.Qt.DecorationRole)

        super(Delegate, self).paint(painter, option, index)

if __name__ == '__main__':
    import sys
    
    def item(file):
        pixmapItem = QtGui.QStandardItem()
        pixmapItem.setData(QtGui.QPixmap(file))
        pixmapItem.setData(QtGui.QPixmap(file), QtCore.Qt.DecorationRole)
        return pixmapItem

    app = QtWidgets.QApplication(sys.argv)
    view = FileList()
    columnCount = view.model.columnCount()
    view.model.setHorizontalHeaderLabels([ 'column' + str(i) for i in range(columnCount) ])
    for row in range(3):
        items = [ QtGui.QStandardItem(str(row) + str(column)) for column in range(columnCount) ]
        view.model.appendRow(items)
        for child in range(3):
            items[0].appendRow([ QtGui.QStandardItem(str(row) + str(column) + str(child)) for column in range(columnCount) ])
    
    #view.model.appendRow([ item('cell00'), item('cell01'), QtGui.QStandardItem() ])
    #view.model.appendRow([ item('cell10'), item('cell11'), item('cell12') ])
    #view.model.appendRow([ item('cell20'), item('cell21'), QtGui.QStandardItem() ])

    view.resize(600, 300)
    view.show()
    app.exec()