## ------------------------------------------------------------------------
## iSiLK - A python-based graphical user interface for SiLK
## ------------------------------------------------------------------------
## Copyright (C) 2006-2007 Carnegie Mellon University. All Rights Reserved.
## ------------------------------------------------------------------------
## Authors: Andrew Kompanek (ajk@cert.org)
## ------------------------------------------------------------------------
## GNU General Public License (GPL) Rights pursuant to Version 2, June 1991
## Government Purpose License Rights (GPLR) pursuant to DFARS 252.225-7013
## ------------------------------------------------------------------------
## See the LICENSE file included in this distribution
## ------------------------------------------------------------------------

import wx
import pdb
import sys, os.path, copy
import ws
import config.silk, silk
import wx.lib.mixins.listctrl as listmix
import analysisresult
import remote

def get_file_list(lib_path, ext=[".set",".bin"], libraryOnly = False):
    file_list = []
    # Remote problem set
    if not libraryOnly:
        print "Listing files in %s" % wx.GetApp().GetActiveProblemSet().GetRemoteDir()
        if wx.GetApp().GetActiveProblemSet().DirectoryExists():
            print "Directory exists"
            for f in silk.list_files(wx.GetApp().GetActiveProblemSet().GetRemoteDir()):
                (path, filename) = os.path.split(f)
                full_path = remote.join(wx.GetApp().GetActiveProblemSet().GetRemoteDir(), f)
                print "File %s" % full_path
                if os.path.splitext(filename)[1] in ext:
                    s = FileReference(FileReference.REMOTE_PROBLEM_SET, full_path, "")
                    file_list.append(s) 
    # Remote library
    print "Listing files in %s" % lib_path
    for p in lib_path:
        for f in silk.list_files(p, recurse=True):
            (path, filename) = os.path.split(f)
            file_ext = os.path.splitext(filename)[1]
            if file_ext in ext:
                if f[0:2] == "./":
                    f = f[2:len(f)]
                path = remote.join(p, f)
                print "path = %s" % path
                s = FileReference(FileReference.REMOTE_LIBRARY, path, "")
                file_list.append(s)
    return file_list

class FileReference(object):
    (LOCAL_LIBRARY, REMOTE_LIBRARY, 
        LOCAL_PROBLEM_SET, REMOTE_PROBLEM_SET) = range(0,4)
    def __init__(self, where, filename, notes):
        print "Creating FileReference for filename = %s" % filename
        self.filename = filename
        self.where = where
        self.notes = notes
    def locstr(self):
        locs = ["Local Library", "Library",
                "Current Problem Set (Local)", "Problem Set"]
        return locs[self.where]

class FileListCtrl(wx.ListCtrl, 
        listmix.ListCtrlAutoWidthMixin, 
        listmix.ColumnSorterMixin
        ):
    def __init__(self, parent, sets=None, showLocation=True):
        wx.ListCtrl.__init__(self, parent, wx.NewId(), 
                 style = wx.LC_SINGLE_SEL|wx.LC_REPORT|wx.SIMPLE_BORDER|
                        wx.LC_HRULES|wx.LC_SORT_ASCENDING)
        self.SetFont(ws.XSmFont(wx.SWISS, wx.NORMAL, wx.NORMAL))
        self.sets = sets
        self.showLocation = showLocation
        column = 0
        self.InsertColumn(column, "Name", width=100); column += 1
        if self.showLocation:
            self.InsertColumn(column, "Where", width=100); column += 1
        self.InsertColumn(column, "Type"); column += 1
        self.InsertColumn(column, "Path"); column += 1
        if sets:
            self.AddSets(sets)
        listmix.ListCtrlAutoWidthMixin.__init__(self)
        self.itemDataMap = {}

        self.il = wx.ImageList(8, 8)

        self.il.Add(wx.Bitmap("icons/sort_down.png", wx.BITMAP_TYPE_PNG))
        self.il.Add(wx.Bitmap("icons/sort_up.png", wx.BITMAP_TYPE_PNG))
        self.sm_dn = 0
        self.sm_up = 1

        self.SetImageList(self.il, wx.IMAGE_LIST_SMALL)
        
    def AddSets(self, sets):
        self.DeleteAllItems()
        for (i,s) in enumerate(sets):
            path = s.filename
            (d, filename) = os.path.split(path)
            #ignoreUniqueId = s.where in (FileReference.LOCAL_LIBRARY, FileReference.REMOTE_LIBRARY)
            try:
                name = analysisresult._extractNameFromFileName(filename)#, ignoreUniqueId)
            except:
                try:
                    name = analysisresult._extractNameFromFileName(filename, True)
                except:
                    name = filename
            file_type = "Unknown"
            if os.path.splitext(path)[1] in [".set", ".bin"]:
                file_type = "IP Set"
            elif os.path.splitext(path)[1] in [".pmap"]:
                file_type = "Prefix Map"
            elif os.path.splitext(path)[1] in [".rwf"]:
                file_type = "Flow Data"
            index = self.InsertStringItem(sys.maxint, name)
            column = 1
            if self.showLocation:
                self.SetStringItem(index, column, s.locstr()); column += 1
            self.SetStringItem(index, column, file_type); column += 1
            self.SetStringItem(index, column, path); column += 1
            self.itemDataMap[index] = (name, s.locstr, file_type, path)
            self.SetItemData(index, index)
        
        listmix.ColumnSorterMixin.__init__(self, 4)

    def GetListCtrl(self):
        return self

    def GetSortImages(self):
        return (self.sm_dn, self.sm_up)

class FilePicker(wx.Dialog):
    def __init__(self, parent, ok_label = "Choose",
                 cancel_label = "No Set",
                 lib_path=None,
                 ext=[".set",".bin"],
                 showLocation = True,
                 libraryOnly = False,
                 **kargs):
        wx.Dialog.__init__(self, parent, **kargs)

        sz = wx.GridBagSizer()
        self.list = FileListCtrl(self, showLocation=showLocation)
        self.libraryOnly = libraryOnly
        sz.Add(self.list, (0,0), (1,1), flag=wx.EXPAND|wx.ALL, border=16)
        sz.Add(self.CreateButtonSizer(wx.OK|wx.CANCEL), (1,0), (1,1),
            flag=wx.EXPAND|wx.ALL, border=16)
        sz.AddGrowableCol(0)
        sz.AddGrowableRow(0)
        self.SetSize((720,480))
        self.SetSizer(sz)
        self.Layout()
        self.Bind(wx.EVT_SIZE, self.OnSize) 
        self.Bind(wx.EVT_INIT_DIALOG, self.OnInitDialog)
        self.Bind(wx.EVT_LIST_ITEM_ACTIVATED, self.OnActivateItem)
        self.Bind(wx.EVT_LIST_ITEM_SELECTED, self.OnSelChange)
        self.Bind(wx.EVT_LIST_ITEM_DESELECTED, self.OnSelChange)        
        self.list.Select(0)
        self.FindWindowById(wx.ID_OK).SetDefault()
        self.FindWindowById(wx.ID_OK).SetLabel(ok_label)
        self.FindWindowById(wx.ID_CANCEL).SetLabel(cancel_label)
        if self.list.GetSelectedItemCount() == 0:
            self.FindWindowById(wx.ID_OK).Enable(False)
        else:
            self.FindWindowById(wx.ID_OK).Enable(True)
            
        
        self.lib_path = lib_path
        self.ext = ext
        if self.lib_path == None:
            self.lib_path = config.silk.opts["REMOTE_LIBRARY_PATH"].split(":")
            

    def OnInitDialog(self, evt):
        busy = wx.BusyInfo("Getting contents of library", self)
        self.items = get_file_list(self.lib_path, self.ext, libraryOnly=self.libraryOnly)
        print "Got items: %s" % self.items
        if len(self.items) != 0:
            self.list.AddSets(self.items)
        else:
            self.FindWindowById(wx.ID_OK).Enable(False)
            
    def OnSelChange(self, evt):
        self.FindWindowById(wx.ID_OK).Enable(self.GetResult()!=None)
            
    def OnActivateItem(self, evt):
        self.EndModal(wx.ID_OK)
        
    def GetResult(self):
        sel = self.list.GetFirstSelected()
        if sel==-1: 
            return None
        return self.items[sel].filename
    
    def OnSize(self, evt):
        self.Layout()
        
        
class FileManager(wx.Dialog):
    def __init__(self, parent,
                 lib_path=None,
                 ext=[".set",".bin"], **kargs):
        wx.Dialog.__init__(self, parent, **kargs)

        sz = wx.GridBagSizer()
        self.list = FileListCtrl(self)
        self.szops = wx.BoxSizer(wx.HORIZONTAL)
        self.btnrefresh = ws.SmButton(self, -1, "Refresh library")
        self.btndel = ws.SmButton(self, -1, "Delete")
        self.szops.Add(self.btnrefresh, flag=wx.ALL, border=8)
        self.szops.Add(self.btndel, flag=wx.ALL, border=8)        
        sz.Add(self.list, (0,0), (1,2), flag=wx.EXPAND|wx.ALL, border=8)
        sz.Add(self.szops, (1,0), (1,2), flag=wx.LEFT|wx.BOTTOM|wx.ALIGN_LEFT, border=4)
        sz.Add(self.CreateButtonSizer(wx.OK), (3,0), (1,2),
            flag=wx.EXPAND|wx.BOTTOM|wx.RIGHT, border=16)
        sz.AddGrowableCol(0)
        sz.AddGrowableRow(0)
        self.SetSize((640,512))
        self.SetSizer(sz)
        self.Layout()
        self.Bind(wx.EVT_SIZE, self.OnSize) 
        self.Bind(wx.EVT_INIT_DIALOG, self.OnInitDialog)
        self.Bind(wx.EVT_BUTTON, self.OnDelete, self.btndel)
        self.Bind(wx.EVT_BUTTON, self.OnRefresh, self.btnrefresh)
        self.Bind(wx.EVT_LIST_ITEM_SELECTED, self.OnSelItem, self.list)
        self.Bind(wx.EVT_LIST_ITEM_DESELECTED, self.OnSelItem, self.list)        
        self.list.Bind(wx.EVT_CHAR, self.OnChar)
        self.FindWindowById(wx.ID_OK).SetDefault()
        self.FindWindowById(wx.ID_OK).SetLabel("Close")
        self.OnSelItem(None)
        self.ext = ext
        self.lib_path = lib_path
        if self.lib_path == None:
            self.lib_path = config.silk.opts["REMOTE_LIBRARY_PATH"].split(":")
    
    def DeleteSelection(self):
        sel = self.list.GetFirstSelected()
        if sel == -1: 
              return
        d = wx.MessageDialog(self,
            "This will permanently delete the set from the library.  Are you sure?",
             "Confirm", wx.YES_NO | wx.ICON_QUESTION)
        res = d.ShowModal()
        if res != wx.ID_YES:
            return
        filename = self.items[sel].filename      
        print "Deleting file = %s" % filename
        context = silk.SimpleRunContext(runRemote=True)
        silk.remove_remote_file(context, filename)
        self.RefreshList()
        
    def OnSelItem(self, evt):
        if self.list.GetFirstSelected() == -1:
            self.btndel.Enable(False)
        else:
            self.btndel.Enable(True)
        
    def OnRefresh(self, evt):
        self.RefreshList()
        
    def OnDelete(self, evt):
        self.DeleteSelection()
    
    def OnChar(self, evt):
        if evt.GetKeyCode() in [wx.WXK_DELETE, wx.WXK_BACK]:
            self.DeleteSelection()
        else:
            evt.Skip(True)
            
    def OnInitDialog(self, evt):
        self.RefreshList()

    def RefreshList(self):
        busy = wx.BusyInfo("Refreshing library...", self)
        self.items = get_file_list(self.lib_path, self.ext, libraryOnly=True)
        print "Got items: %s" % self.items
        if len(self.items) != 0:
            self.list.AddSets(self.items)
        else:
            self.FindWindowById(wx.ID_OK).Enable(False)

    def OnSize(self, evt):
        self.Layout()
