#More information http://www.plugincafe.com/forum/forum_posts.asp?TID=14102&PID=56287#56287
import c4d

# Be sure to use a unique ID obtained from http://www.plugincafe.com/.
PLUGIN_ID = 1000010 # TEST ID ONLY

# TreeView Column IDs.
ID_CHECKBOX = 1
ID_NAME = 2
ID_OTHER = 3

class PresetObject(object):
    """
    Class which represent a texture, aka an Item in our list
    """
    itemName = "Default"
    nodeType = "Preset"
    presetItems = []
    opened = True  
    _selected = False

    def __init__(self, itemName,presetItems):
        self.itemName = itemName
        self.presetItems = presetItems

    @property
    def IsSelected(self):
        return self._selected

    def Select(self):
        self._selected = True

    def Deselect(self):
        self._selected = False

    def __repr__(self):
        return str(self)

    def __str__(self):
        return self.itemName

class PresetItem(object):
    """
    Class which represent a item, aka an Item in our list
    """
    itemName = "Default"
    nodeType = "Item"
    # parentNode = None
    _selected = False

    def __init__(self, itemName, parentNode):
        self.itemName = itemName
        self.parentNode = parentNode
        # self.mikeData += itemName

    @property
    def IsSelected(self):
        return self._selected

    def Select(self):
        self._selected = True

    def Deselect(self):
        self._selected = False

    def __repr__(self):
        return str(self)

    def __str__(self):
        return self.itemName

class ListView(c4d.gui.TreeViewFunctions):

    def __init__(self):
        self.presetList = list() # Store all objects we need to display in this list

    def InputEvent(self, root, userdata, pArea, msg):
        
        if msg[c4d.BFM_INPUT_CHANNEL] == c4d.BFM_INPUT_MOUSELEFT: 
            print "Left Mouse Pressed" 

        if msg[c4d.BFM_INPUT_CHANNEL] == c4d.BFM_INPUT_MOUSERIGHT: 
            print "Right Mouse Pressed"

        return True

    def IsResizeColAllowed(self, root, userdata, lColID):
        return True

    def IsTristate(self, root, userdata):
        return False

    def GetColumnWidth(self, root, userdata, obj, col, area):
        return 80  # All have the same initial width

    def IsMoveColAllowed(self, root, userdata, lColID):
        # The user is allowed to move all columns.
        # TREEVIEW_MOVE_COLUMN must be set in the container of AddCustomGui.
        return True

    def Open(self, root, userdata, obj, opened):
        """
        Called when the user opens or closes the children of *obj* in
        the Tree View.
        """

        if obj.nodeType == "Preset":
            obj.opened = opened        # return None

    def IsOpened(self, root, userdata, obj):
        """
        Returns:
            (bool): True if *obj* is opened (expaneded), False if it is
                closed (collapsed).
        """

        if obj.nodeType == "Preset":
          return obj.opened
        return False

    def GetFirst(self, root, userdata):
        """
        Return the first element in the hierarchy, or None if there is no element.
        """
        rValue = None if not self.presetList else self.presetList[0]
        return rValue

    def GetDown(self, root, userdata, obj):
        """
        Return a child of a node, since we only want a list, we return None everytime
        """
        if obj.nodeType == "Preset" and obj.presetItems:
            return obj.presetItems[0]

        return None

    def GetNext(self, root, userdata, obj):
        """
        Returns the next Object to display after arg:'obj'
        """
        rValue = None

        pList = self.GetItemsList(obj)
        
        currentObjIndex = pList.index(obj)
        nextIndex = currentObjIndex + 1
        if nextIndex < len(pList):
            rValue = pList[nextIndex]

        return rValue

    def GetPred(self, root, userdata, obj):
        """
        Returns the previous Object to display before arg:'obj'
        """
        pList = self.GetItemsList(obj)
        
        currentObjIndex = pList.index(obj)
        nextIndex = currentObjIndex - 1
        if nextIndex < len(pList):
            rValue = pList[nextIndex]

    def GetId(self, root, userdata, obj):
        """
        Return a unique ID for the element in the TreeView.
        """
        return hash(obj)

    def Select(self, root, userdata, obj, mode):
        """
        Called when the user selects an element.
        """
        if mode == c4d.SELECTION_NEW:
            for preset in self.presetList:
                preset.Deselect()
                for item in preset.presetItems:
                    item.Deselect()
            obj.Select()
        elif mode == c4d.SELECTION_ADD:
            obj.Select()
        elif mode == c4d.SELECTION_SUB:
            obj.Deselect()

    def IsSelected(self, root, userdata, obj):
        """
        Returns: True if *obj* is selected, False if not.
        """
        return obj.IsSelected

    def SetCheck(self, root, userdata, obj, column, checked, msg):
        """
        Called when the user clicks on a checkbox for an object in a
        `c4d.LV_CHECKBOX` column.
        """
        if checked:
            obj.Select()
        else:
            obj.Deselect()

    def IsChecked(self, root, userdata, obj, column):
        """
        Returns: (int): Status of the checkbox in the specified *column* for *obj*.
        """
        if obj.IsSelected:
            return c4d.LV_CHECKBOX_CHECKED | c4d.LV_CHECKBOX_ENABLED
        else:
            return c4d.LV_CHECKBOX_ENABLED

    def GetName(self, root, userdata, obj):
        """
        Returns the name to display for arg:'obj', only called for column of type LV_TREE
        """
        return str(obj) # Or obj.presetType

    def SetName(self, root, userdata, obj, name):
        """
        Returns the name to display for arg:'obj', only called for column of type LV_TREE
        """
        # doc = c4d.documents.GetActiveDocument()
        # doc.StartUndo()
        # doc.AddUndo(c4d.UNDOTYPE_CHANGE_SMALL, obj)
        # doc.EndUndo()
        # obj.mikeData = name
        # obj.otherData = name
        # c4d.EventAdd()
        pass

    def CreateContextMenu(self, root, userdata, obj, lColumn, bc):
        """
        User clicked with the right mouse button on an entry so
        we can here enhance the menu
        """
        
        # bc = c4d.BaseContainer()
        # if obj and obj.nodeType == "Preset":
        #     bc.SetString(PRESBUT_CONT_MENU_DUP, "Duplicate")
            
        #     if self.presetList.index(obj):
        #         bc.SetString(PRESBUT_CONT_MENU_MOVEUP, "Move Up")
            
        #     if self.presetList.index(obj) != (len(self.presetList)-1):
        #         bc.SetString(PRESBUT_CONT_MENU_MOVEDOWN, "Move Down")
        
        # bc.RemoveIndex(c4d.ID_TREEVIEW_CONTEXT_REMOVE)
        # bc.SetString(c4d.ID_TREEVIEW_CONTEXT_RESET, "Remove All")

    def ContextMenuCall(self, root, userdata, obj, lColumn, lCommand):
        
        """
        The user executes an entry of the context menu.

        Returns:
            (bool): True if the event was handled, False if not.
        """
        if lCommand == c4d.ID_TREEVIEW_CONTEXT_RESET:
            print 'Call Reset'
            return True

        if lCommand == c4d.ID_TREEVIEW_CONTEXT_REMOVE:
            print 'Call Remove All'
            return True


    def DrawCell(self, root, userdata, obj, col, drawinfo, bgColor):
        """
        Draw into a Cell, only called for column of type LV_USER
        """
        # if col == ID_NAME:
        #     name = obj.itemName
        #     geUserArea = drawinfo["frame"]
        #     w = geUserArea.DrawGetTextWidth(name)
        #     h = geUserArea.DrawGetFontHeight()
        #     xpos = drawinfo["xpos"]
        #     ypos = drawinfo["ypos"] + drawinfo["height"]
        #     drawinfo["frame"].DrawText(name, xpos, ypos - h * 1.1)

        if col == ID_OTHER:
            if obj.nodeType == 'Preset':
                name = obj.nodeType
                geUserArea = drawinfo["frame"]
                w = geUserArea.DrawGetTextWidth(name)
                h = geUserArea.DrawGetFontHeight()
                xpos = drawinfo["xpos"]
                ypos = drawinfo["ypos"] + drawinfo["height"]
                drawinfo["frame"].DrawText(name, xpos, ypos - h * 1.1)

    def InsertObject(self, root, userdata, obj, dragtype, dragobject, insertmode, bCopy):
        """
        Called when a drag is dropped on the TreeView.
        """
        print obj,insertmode

        pList = self.GetItemsList(obj)
        currentObjIndex = pList.index(obj)
        newID = len(pList) + 1 

        if obj.nodeType == 'Preset': 
            insert_obj = PresetObject("T{}".format(newID),[])
            insert_obj.presetItems.append(PresetItem(dragobject[0].GetName(),insert_obj))
        
        elif obj.nodeType == 'Item': 
            insert_obj = PresetItem(dragobject[0].GetName(),obj.parentNode)

        if insertmode == c4d.INSERT_BEFORE:   
            nextIndex = currentObjIndex - 1
        if insertmode == c4d.INSERT_AFTER:
            nextIndex = currentObjIndex + 1

        if insertmode == c4d.INSERT_UNDER:
            insert_obj = PresetItem(dragobject[0].GetName(),obj)
            obj.presetItems.insert(0, insert_obj)
            return
        # c4d.INSERT_REPLACE


        # tex = PresetObject("T{}".format(newID),[])
        # tex.mikeData = dragobject[0].GetName()
        pList.insert(nextIndex, insert_obj)
        # self.Refresh()

        return

        # if dragtype != c4d.DRAGTYPE_ATOMARRAY:
        #     return # Shouldnt happen, we catched that in AcceptDragObject

        # for op in dragobject:
        #     op.Remove()
        #     if insertmode == c4d.INSERT_BEFORE:
        #         op.InsertBefore(obj)
        #     elif insertmode == c4d.INSERT_AFTER:
        #         op.InsertAfter(obj)
        #     elif insertmode == c4d.INSERT_UNDER:
        #         op.InsertUnder(obj)
        # return

    def GetDragType(self, root, userdata, obj):
        """
        Returns:
            (int): The drag datatype.
        """

        return c4d.DRAGTYPE_ATOMARRAY
    
    def GenerateDragArray(self, root, userdata, obj):
        """
        Return:
            (list of c4d.BaseList2D): Generate a list of objects that can be
                dragged from the TreeView for the `c4d.DRAGTYPE_ATOMARRAY` type.
        """

        if obj.GetBit(c4d.BIT_ACTIVE):
            return [obj, ]
            
    def DragStart(self, root, userdata, obj):
        """
        Returns:
            (int): Bitmask specifying options for the drag, whether it is
                allowed, etc.
        """

        return c4d.TREEVIEW_DRAGSTART_ALLOW | c4d.TREEVIEW_DRAGSTART_SELECT

    def AcceptDragObject(self, root, userdata, obj, dragtype, dragobject):
        """
        Called when a drag & drop operation hovers over the TreeView to check
        if the drag can be accepted.

        Returns:
            (int, bool)
        """
        print dragobject,dragtype

        if dragtype not in {c4d.DRAGTYPE_ATOMARRAY}:# or (dragobject not in self.presetList):
            return 0,False

        if obj.nodeType  == 'Item':
            return c4d.INSERT_BEFORE | c4d.INSERT_AFTER, True


        return c4d.INSERT_BEFORE | c4d.INSERT_AFTER | c4d.INSERT_UNDER, True

    def DoubleClick(self, root, userdata, obj, col, mouseinfo):
        """
        Called when the user double-clicks on an entry in the TreeView.

        Returns:
          (bool): True if the double-click was handled, False if the
            default action should kick in. The default action will invoke
            the rename procedure for the object, causing `SetName()` to be
            called.
        """
        # c4d.gui.MessageDialog("You clicked on " + str(obj))
        for o in obj.presetItems:
            print o.itemName

        return True

    def DeletePressed(self, root, userdata):
        """Called when a delete event is received."""
        
        for preset in reversed(self.presetList):
            
            if preset.IsSelected:
                self.presetList.remove(preset)
                continue
            
            for item in reversed(preset.presetItems):
                if item.IsSelected:
                    preset.presetItems.remove(item)

    def GetItemsList(self,obj):
        
        if obj.nodeType == "Preset":
            return self.presetList
        
        elif obj.nodeType == "Item":
            return obj.parentNode.presetItems

class TestDialog(c4d.gui.GeDialog):
    _treegui = None # Our CustomGui TreeView
    _listView = ListView() # Our Instance of c4d.gui.TreeViewFunctions

    def CreateLayout(self):
        # Create the TreeView GUI.
        self.GroupBegin(1001, c4d.BFH_SCALEFIT | c4d.BFV_SCALEFIT, 0, 0,'')
        customgui = c4d.BaseContainer()
        customgui.SetBool(c4d.TREEVIEW_BORDER, c4d.BORDER_THIN_IN)
        customgui.SetBool(c4d.TREEVIEW_HAS_HEADER, True) # True if the tree view may have a header line.
        customgui.SetBool(c4d.TREEVIEW_HIDE_LINES, False) # True if no lines should be drawn.
        customgui.SetBool(c4d.TREEVIEW_MOVE_COLUMN, True) # True if the user can move the columns.
        customgui.SetBool(c4d.TREEVIEW_RESIZE_HEADER, True) # True if the column width can be changed by the user.
        customgui.SetBool(c4d.TREEVIEW_FIXED_LAYOUT, True) # True if all lines have the same height.
        customgui.SetBool(c4d.TREEVIEW_ALTERNATE_BG, True) # Alternate background per line.
        customgui.SetBool(c4d.TREEVIEW_CURSORKEYS, True) # True if cursor keys should be processed.
        customgui.SetBool(c4d.TREEVIEW_NOENTERRENAME, False) # Suppresses the rename popup when the user presses enter.

        self._treegui = self.AddCustomGui(1000, c4d.CUSTOMGUI_TREEVIEW, "", c4d.BFH_SCALEFIT | c4d.BFV_SCALEFIT, 100, 100, customgui)
        if not self._treegui:
            print "[ERROR]: Could not create TreeView"
            return False

        self.GroupEnd()
        self.AddButton(1002, c4d.BFH_CENTER, name="Add")
        return True

    def InitValues(self):
        # Initialize the column layout for the TreeView.
        layout = c4d.BaseContainer()
        # layout.SetLong(ID_CHECKBOX, c4d.LV_CHECKBOX)
        layout.SetLong(ID_NAME, c4d.LV_TREE)
        layout.SetLong(ID_OTHER, c4d.LV_USER)
        self._treegui.SetLayout(3, layout)

        # Set the header titles.
        # self._treegui.SetHeaderText(ID_CHECKBOX, "Check")
        self._treegui.SetHeaderText(ID_NAME, "Preset")
        self._treegui.SetHeaderText(ID_OTHER, "Type")
        self._treegui.Refresh()

        # Set TreeViewFunctions instance used by our CUSTOMGUI_TREEVIEW
        self._treegui.SetRoot(self._treegui, self._listView, None)
        return True
    
    def Message(self, msg, result): 
        
        if msg.GetId()==c4d.BFM_DRAGRECEIVE:
            
    #         # Discard if lost drag or if it has been escaped
    #         if msg.GetInt32(c4d.BFM_DRAG_LOST) or msg.GetInt32(c4d.BFM_DRAG_ESC):
    #             # self.index_list = []
    #             return self.SetDragDestination(c4d.MOUSE_FORBIDDEN)
            
    #         # Check drop area and discard if not on the user area
            if not self.CheckDropArea(1000, msg, True, True):
                return self.SetDragDestination(c4d.MOUSE_FORBIDDEN)
    #         # else:
    #             # return True
    

    #         # if self.CheckDropArea(1001,msg, True, True):
    #             # print "Dragged"
    #             # self.SetDragDestination(c4d.MOUSE_INSERTCOPYDOWN)
        
    #         # return self.SetDragDestination(c4d.MOUSE_INSERTCOPY)

        return c4d.gui.GeDialog.Message(self, msg, result)

    def Command(self, id, msg):
        # Click on button
        if id == 1002:
            # Add data to our DataStructure (ListView)
            newID = len(self._listView.presetList) + 1
            
            tex = PresetObject("T{}".format(newID),[])
            
            # tex.presetItems.extend([PresetItem('Mike',tex),PresetItem('Udin',tex)])
            self._listView.presetList.append(tex)

            # Refresh the TreeView
            self._treegui.Refresh()

        return True

class MenuCommand(c4d.plugins.CommandData):
    dialog = None

    def Execute(self, doc):
        if self.dialog is None:
            self.dialog = TestDialog()
        return self.dialog.Open(c4d.DLG_TYPE_ASYNC, PLUGIN_ID, defaulth=600, defaultw=600)

    def RestoreLayout(self, sec_ref):
        if self.dialog is None:
            self.dialog = TestDialog()
        return self.dialog.Restore(PLUGIN_ID, secret=sec_ref)


def main():
    c4d.plugins.RegisterCommandPlugin(
        PLUGIN_ID, "Python TreeView Example", 0, None, "Python TreeView Example", MenuCommand())


if __name__ == "__main__":
    main()