USE_TEXT = False  # Use TextDataObject if True, CustomDataObject if False

import wx

if not USE_TEXT:
    import pickle

class MainWindow(wx.Frame):
    """ This window displays the main Program form. """
    def __init__(self, parent, id, title):

        # Define the main Frame
        wx.Frame.__init__(self, parent, id, title, size = (600,600), style=wx.DEFAULT_FRAME_STYLE|wx.NO_FULL_REPAINT_ON_RESIZE)
        self.SetBackgroundColour(wx.WHITE)

        menuBar = wx.MenuBar()
        menu1 = wx.Menu()
        menu1.Append(wx.ID_CUT, "Cut")
        menu1.Append(wx.ID_COPY, "Copy")
        menu1.Append(wx.ID_PASTE, "Paste")
        menu1.Append(wx.ID_EXIT, "Exit")
        menuBar.Append(menu1, 'Edit')
        self.SetMenuBar(menuBar)

        self.Bind(wx.EVT_MENU, self.OnCutCopy, id=wx.ID_CUT)
        self.Bind(wx.EVT_MENU, self.OnCutCopy, id=wx.ID_COPY)
        self.Bind(wx.EVT_MENU, self.OnPaste, id=wx.ID_PASTE)
        self.Bind(wx.EVT_MENU, self.OnClose, id=wx.ID_EXIT)

        mainSizer = wx.BoxSizer(wx.VERTICAL)

        self.treeCtrl = DragAndDropTreeCtrl(self, wx.ID_ANY)
        mainSizer.Add(self.treeCtrl, 5, wx.EXPAND | wx.ALL, 5)

        self.txtCtrl = wx.TextCtrl(self, wx.ID_ANY, style=wx.TE_MULTILINE)
        mainSizer.Add(self.txtCtrl, 2, wx.EXPAND | wx.LEFT | wx.RIGHT | wx.BOTTOM, 5)

        if USE_TEXT:
            self.txtCtrl.AppendText('Using TextDataObject\n\n')
        else:
            self.txtCtrl.AppendText('Using CustomDataObject\n\n')
        
        self.SetSizer(mainSizer)
        self.SetAutoLayout(True)
        self.Layout()
        self.CenterOnScreen()
        self.Show(True)

    def OnCutCopy(self, event):
        self.treeCtrl.OnCutCopy(event)

    def OnPaste(self, event):

        self.treeCtrl.OnPaste(event)

    def OnClose(self, event):
        self.Close()


class DragAndDropTreeCtrl(wx.TreeCtrl):

    def __init__(self, parent, id):
        wx.TreeCtrl.__init__(self, parent, id, style=wx.TR_HAS_BUTTONS | wx.TR_EDIT_LABELS)
        self.parent = parent

        dt = DragAndDropTreeDropTarget(self.parent, self)
        self.SetDropTarget(dt)
        self.populateTree()

        self.cutCopy = None

        self.Bind(wx.EVT_TREE_BEGIN_DRAG, self.OnBeginDrag, id=self.GetId())

    def populateTree(self):

        self.DeleteAllItems()

        self.root = self.AddRoot('Main Node')

        self.AppendItem(self.root, 'Item 4')
        self.AppendItem(self.root, 'Item 7')
        self.AppendItem(self.root, 'Item 1')
        self.AppendItem(self.root, 'Item 5')
        self.AppendItem(self.root, 'Item 8')
        self.AppendItem(self.root, 'Item 6')
        self.AppendItem(self.root, 'Item 2')
        self.AppendItem(self.root, 'Item 3')

        self.Expand(self.root)

    def OnCutCopy(self, event):
        self.cutCopy = event.GetId()
        self.OnBeginDrag(event)

    def OnPaste(self, event):
        if event.GetId() == wx.ID_PASTE:
            if USE_TEXT:

                cdo = wx.TextDataObject()
                
            else:
                df = wx.DataFormat('application.Test.TreeCtrlDragData')
                cdo = wx.CustomDataObject(df)

            wx.TheClipboard.Open()
            success = wx.TheClipboard.GetData(cdo)
            if success:
                try:
                    if USE_TEXT:
                        data = cdo.GetText()
                    else:
                        data = pickle.loads(cdo.GetData())
                except:
                    data = None
            wx.TheClipboard.Close()
            if self.cutCopy == wx.ID_CUT:
                dragResult = wx.DragMove
            else:
                dragResult = wx.DragCopy

            self.UpdateTree(self.GetSelection(), data, dragResult)

    def UpdateTree(self, dropNode, dropData, dragResult):
        pos = 0
        (node, cookie) = self.GetFirstChild(self.root)
        while node.IsOk():
            if self.GetItemText(node) == dropData:

                if dragResult == wx.DragMove:
                    self.Delete(node)
                    txt = 'Mov'
                else:
                    txt = 'Copy'
                if dropNode == self.root:
                    self.parent.txtCtrl.AppendText('{1}ing {0} to 1st in node'.format(dropData, txt))
                    self.PrependItem(self.root, dropData)
                else:
                    self.parent.txtCtrl.AppendText('{2}ing {0} after {1}'.format(dropData, self.GetItemText(dropNode), txt))
                    self.InsertItem(self.root, dropNode, dropData)
                break
              
            (node, cookie) = self.GetNextChild(self.root, cookie)
            pos += 1

    def OnBeginDrag(self, event):

        event.Skip()

        dndList = []
        
        item = self.GetSelection()
        
        text = self.GetItemText(item)
           
        if USE_TEXT:
            # Create a TextDataObject to hold the data being dragged
            cdo = wx.TextDataObject()
            cdo.SetText(text)
        else:
            # Create a CustomDataObject to hold the data being dragged
            df = wx.DataFormat('application.Test.TreeCtrlDragData')
            cdo = wx.CustomDataObject(df)
            ptext = pickle.dumps(text, 1)
            cdo.SetData(ptext)
            
        if event.GetId() in [wx.ID_CUT, wx.ID_COPY]:
            # Open the Clipboard
            wx.TheClipboard.Open()
            # ... put the data in the clipboard ...
            wx.TheClipboard.SetData(cdo)
            # Close the Clipboard
            wx.TheClipboard.Close()
        else:
            tds = DragAndDropTreeDropSource(self.parent, self)
            tds.SetData(cdo)

            try:
                
                print('About to call DoDragDrop()')
                
                dragResult = tds.DoDragDrop(wx.Drag_AllowMove)

                print('DoDragDrop completed with result {0}'.format(dragResult))
                
            except:

                self.parent.txtCtrl.AppendText('DragAndDropTreeCtrl.OnBeginDrag():  Exception in DoDragDrop()')
                pass
        

class DragAndDropTreeDropSource(wx.DropSource):

    def __init__(self, parent, tree):
        self.parent = parent

        wx.DropSource.__init__(self, tree)
        self.tree = tree

    def SetData(self, obj):

        wx.DropSource.SetData(self, obj)

        if USE_TEXT:
            self.data = obj.GetText()
            self.parent.txtCtrl.AppendText('DragAndDropTreeDropSource.SetData(): {0}\n\n'.format(self.data.encode('utf8')))
        else:
            self.data = obj.GetData()
            self.parent.txtCtrl.AppendText('DragAndDropTreeDropSource.SetData(): {0}\n\n'.format(pickle.loads(self.data)))
        
    def GiveFeedback(self, effect):

        (winx, winy) = wx.GetMousePosition()
        (x, y) = self.tree.ScreenToClient(winx, winy)
        (dropNode, flag) = self.tree.HitTest((x, y))

        if dropNode.IsOk():
            return True
        else:
            return False


class DragAndDropTreeDropTarget(wx.PyDropTarget):

    def __init__(self, parent, tree):
        self.parent = parent

        wx.DropTarget.__init__(self)
        self.tree = tree

        if USE_TEXT:
            self.data = wx.TextDataObject()
        else:
            df = wx.DataFormat('application.Test.TreeCtrlDragData')
            self.data = wx.CustomDataObject(df)
        self.SetDataObject(self.data)

    def OnDrop(self, x, y):

        (self.dropNode, flag) = self.tree.HitTest((x, y))

        if self.dropNode.IsOk():
            self.parent.txtCtrl.AppendText('DragAndDropTreeDropTarget.OnDrop(): dropNode is {0}\n'.format(self.tree.GetItemText(self.dropNode)))
            return True
        else:
            self.dropNode = None
            return False
            

    def OnData(self, x, y, dragResult):

        if self.dropNode is not None:
            self.parent.txtCtrl.AppendText('DragAndDropTreeDropTarget.OnData(): dropNode is {0}\n'.format(self.tree.GetItemText(self.dropNode)))

            if self.GetData():
                if USE_TEXT:
                    data = self.data.GetText()
                else:
                    pdata = self.data.GetData()
                    data = pickle.loads(pdata)
              
                self.parent.txtCtrl.AppendText('DragAndDropTreeDropTarget.OnData(): data is {0}\n\n'.format(data))
            self.tree.UpdateTree(self.dropNode, data, dragResult)

        return dragResult

class MiniApp(wx.App):

    def OnInit(self):
        frame = MainWindow(None, -1, "Drag-and-Drop Demo")
        self.SetTopWindow(frame)
        return True
    

app = MiniApp()
app.MainLoop()