## ------------------------------------------------------------------------
## iSiLK - A python-based graphical user interface for SiLK
## ------------------------------------------------------------------------
## Copyright (c) 2007-2008 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
## ------------------------------------------------------------------------

# FIXME: Odd things happen if I adjust the window size before
# adding the graph window--if I make the window smaller after that
# the contents do not resize properly.
#

import wx, ws
import logging, re
import traceback, pdb
from wx.lib.buttonpanel import ButtonPanel, ButtonInfo
import  wx.lib.buttons as buttons
import datetime, os, string, tempfile, threading
import views
from datacache import formatFieldAsString
from constants import *
from datacache import SimpleDataCache
from analysisresult import ProblemSet, ChangeEventSource
from analysisresult import ResultRunContext
from analysistool import ImportTool
from tools.filter import FilterTool
from tools.uniq import UniqTool
import tools.settool
import tools.uniq
from tools.count import CountTool
from resultset import QueryResult, SetResult
from resultset import UniqResult, CountResult, SimpleSplitResult
from resultset import FilterResult
from graphresult import GraphResult, CountGraph
from infoframe import InfoFrame
import silk
import xmlutils
import config.silk
import analysistool
from config.editor import RecordViewOptionsDialog
import config.viewoptions
from resultbook import ResultViewBook
import library
from filemanager import FileManager
from tools.plugins import ShowCreatePluginDialog, ManageCustomPlugins
from tools.plugins import CreateCustomPlugin
from tools import RegisterPlugin

class ChooseFieldsDialog(wx.MultiChoiceDialog):
    def __init__(self, parent, title, fields):
        self.fields = fields
        wx.MultiChoiceDialog.__init__(self, parent, "Choose fields",
                title, self.fields)

    def GetResult(self):
        selections = self.GetSelections()
        strings = [self.fields[x] for x in selections]
        return strings


class ChooseFieldDialog(wx.SingleChoiceDialog):
    def __init__(self, parent, title, fields):
        self.fields = fields
        wx.SingleChoiceDialog.__init__(self, parent, "Choose fields",
                title, self.fields)

    def GetResult(self):
        selections = [self.GetSelection()]
        strings = [self.fields[x] for x in selections]
        return strings

    
class ExplorerStatusBar(wx.StatusBar):
    
    def __init__(self, parent):
        wx.StatusBar.__init__(self, parent)
        self.SetFieldsCount(1)
 
    # reposition the checkbox
    def Reposition(self):
         rect = self.GetFieldRect(1)
         self.cb.SetPosition((rect.x+2, rect.y+2))
         self.cb.SetSize((rect.width-4, rect.height-4))
         self.sizeChanged = False


class CommandLogger(wx.PyLog):
    def __init__(self, textCtrl, logTime=0):
         wx.PyLog.__init__(self)
         self.textCtrl = textCtrl
         self.logTime = logTime

    def DoLogString(self, message, timeStamp):
         self.textCtrl.AppendText(message + '\n')


class AnalysisTool(object):
    def __init__(self, id, name, bitmap_file):
        self.id = id
        self.name = name
        self.bitmap_file = bitmap_file


class ProblemSetExplorer(wx.Frame):
    
    def GetLastActive():
        return ProblemSetExplorer._lastActive
    GetLastActive = staticmethod(GetLastActive)
    
    def _buildToolbar(self):
        # Build up a list of tools to put on the toolbar
        allTools = []
        basicTools = [
            AnalysisTool(ID_QUERY_BUILDER, "Query", "icons/cog-query.gif"),
            None,
            AnalysisTool(ID_SHOW_INFO, "Info", "icons/info.gif")]
        if wx.Platform in ("__WXGTK__", "__WXMAC__", "__WXMSW__"):
            basicTools.append(
                AnalysisTool(ID_SHOW_FILES, "Local Files", "icons/folder.gif"))
        # FIXME: this icon doesn't display properly on Windows
        # Disabling this for now
        if wx.Platform == "__WXMAC__":
            basicTools.append(AnalysisTool(ID_OPEN_EXCEL,
                "Excel", "icons/excel.gif"))
        if wx.Platform == "__WXMSW__":
            basicTools.append(AnalysisTool(ID_OPEN_EXCEL,
                "Excel", "icons/excel_windows.gif"))                    
        analysisTools = [
            AnalysisTool(ID_ANALYSIS_REFINE, "Filter", "icons/cog-f.gif"),
            AnalysisTool(ID_ANALYSIS_COUNT, "Uniq", "icons/cog-u.gif"),
            AnalysisTool(ID_ANALYSIS_TREND, "Count", "icons/cog-c.gif"),
            None,
            AnalysisTool(ID_ANALYSIS_SET, "Set", "icons/cog-s.gif"),
        ]
        
        graphTools = [
            AnalysisTool(ID_ANALYSIS_SMART_GRAPH, "Quick Graph", "icons/graph.gif"),
        ]
        
        # Add plugin tools
        pluginTools = []
        #for (i, plugin) in enumerate(tools.GetPlugins()):
        #    pluginTools.append(AnalysisTool(ID_ANALYSIS_CUSTOM+i, 
        #                    plugin.toolName, plugin.toolIcon))
                            
        allTools.extend(basicTools)
        allTools.append(None)
        allTools.extend(analysisTools)
        allTools.append(None)
        allTools.extend(graphTools)
        #allTools.append(None)
        #allTools.extend(pluginTools)
        #allTools.append(None)

        # Create the tool bar and populate it
        self.CreateToolBar(style=wx.TB_TEXT|wx.TB_FLAT)
        #self.GetToolBar().ClearTools()
        self.GetToolBar().SetToolBitmapSize((32,32))
        self.GetToolBar().SetToolSeparation(1)
        self.GetToolBar().SetToolPacking(1)
        self.GetToolBar().SetMargins((18,6))
        for t in allTools: 
            if t: 
                label = t.name
                self.GetToolBar().AddLabelTool(t.id, label=label,
                    shortHelp=label,
                    bitmap=wx.Bitmap(t.bitmap_file, wx.BITMAP_TYPE_GIF))
                if t.id == ID_QUERY_BUILDER:
                    self.queryToolMenu = self.AddToolMenu("Query")
                elif t.id == ID_ANALYSIS_REFINE:
                    self.filterToolMenu = self.AddToolMenu("Filter")
                elif t.id == ID_ANALYSIS_COUNT:
                    self.uniqToolMenu = self.AddToolMenu("Uniq")
                elif t.id == ID_ANALYSIS_TREND:
                    self.countToolMenu = self.AddToolMenu("Count")
                elif t.id == ID_ANALYSIS_SET:
                    self.setToolMenu = self.AddToolMenu("Set")
            else: 
                self.GetToolBar().AddSeparator()
        self.GetToolBar().Realize()
        
    def AddToolMenu(self, category):
        bmp = wx.Bitmap("icons/sort_down.png", wx.BITMAP_TYPE_PNG)
        b = wx.BitmapButton(self.GetToolBar(), -1, bmp, # (128, 128),
            style = wx.NO_BORDER)
        self.GetToolBar().AddControl(b)
        if category == "Query":
            self.Bind(wx.EVT_BUTTON, self.OnQueryPlugin, b)
        elif category == "Filter":
            self.Bind(wx.EVT_BUTTON, self.OnFilterPlugin, b)
        elif category == "Uniq":
            self.Bind(wx.EVT_BUTTON, self.OnUniqPlugin, b)            
        elif category == "Count":
            self.Bind(wx.EVT_BUTTON, self.OnCountPlugin, b)                        
        elif category == "Set":
            self.Bind(wx.EVT_BUTTON, self.OnSetPlugin, b)                                    
        return b
        
    def _buildPluginPopup(self, tool_icon="icons/cog-u.gif", category=None):
        print "BUilding plugin in popu menu"
        print "plugin list len = %d" % len(tools.GetPlugins())
        menu = wx.Menu()
        plugins = tools.GetPlugins(category)
        if category=="Query":
            menu.Append(ID_QUERY_BUILDER, "Query Builder...")
        elif category=="Filter":
            menu.Append(ID_ANALYSIS_REFINE, "Filter...")
        elif category=="Uniq":
            menu.Append(ID_ANALYSIS_COUNT, "rwuniq...")
        elif category=="Count":
            menu.Append(ID_ANALYSIS_COUNT, "rwcount...")
        elif category=="Set":
            menu.Append(ID_ANALYSIS_SET_BUILD, "rwsetbuild...")                        
            menu.Append(ID_ANALYSIS_SET, "rwset...")            
            menu.Append(ID_ANALYSIS_SET_TOOL, "rwsettool...")                                    
        menu.AppendSeparator()                        
        
        if len(plugins) == 0:
            menu.Append(ID_NO_COMMAND, "(No plug-ins defined)")
        for (i, plugin) in enumerate(plugins):
            print "Adding tool %s (%s)" % (plugin.toolName, category)
            tool_id = plugin.toolId + ID_ANALYSIS_CUSTOM
            tool_name = plugin.toolName
            item = wx.MenuItem (menu, tool_id, " %s" % tool_name)
            bmp = wx.Image(tool_icon, wx.BITMAP_TYPE_GIF)
            bmp = bmp.Rescale(16,16).ConvertToBitmap()
            item.SetBitmap(bmp)
            menu.AppendItem(item) 
        return menu
        
    def OnQueryPlugin(self, evt):
        menu = self._buildPluginPopup("icons/cog-query.gif", "Query")
        pos = self.queryToolMenu.PopupMenu(menu)
        
    def OnFilterPlugin(self, evt):
        menu = self._buildPluginPopup("icons/cog-f.gif", "Filter")        
        pos = self.filterToolMenu.PopupMenu(menu)       
        
    def OnUniqPlugin(self, evt):
        menu = self._buildPluginPopup("icons/cog-u.gif", "Uniq")
        pos = self.uniqToolMenu.PopupMenu(menu)
        
    def OnCountPlugin(self, evt):
        menu = self._buildPluginPopup("icons/cog-c.gif", "Count")
        pos = self.countToolMenu.PopupMenu(menu)        
        
    def OnSetPlugin(self, evt):
        menu = self._buildPluginPopup("icons/cog-s.gif", "Set")
        pos = self.countToolMenu.PopupMenu(menu)                 
        
    def RefreshTools(self):
        #self._buildToolbar()
        self.SetMenuBar(wx.GetApp().CreateMenuBar())

    def __init__(self, app, problemSet, p=(0,0)):
        wx.Frame.__init__(self, None, -1, 
            "Explorer", p, style=wx.DEFAULT_FRAME_STYLE)
        ProblemSetExplorer._lastActive = self
        self._last_log_message = None
        self.app = app
        self.problemSet = problemSet
        # Explorer processes log messages
        self.problemSet.AddLogListener(self)
        self.problemSet.AddChangeListener(self)
        self.SetMenuBar(wx.GetApp().CreateMenuBar())    
        # Create the status bar
        self.statusbar = ExplorerStatusBar(self)
        self.SetStatusBar(self.statusbar)
       
        self._buildToolbar()
        
        class MySplitterWindow(wx.SplitterWindow):
            def __init__(self, *args, **kargs):
                wx.SplitterWindow.__init__(self, *args, **kargs)
            def OnViewChange(self):
                self.GetParent().OnViewChange()

        # Create top-level splitter: problem set view and library
        self.top_splitter = wx.SplitterWindow(self,
            style=wx.SP_LIVE_UPDATE)

        # Create splitter separating viewbook and log panels
        self.splitter = MySplitterWindow(self.top_splitter,
            style=wx.SP_LIVE_UPDATE)
        self.splitter.SetMinimumPaneSize(75)
        if wx.Platform == "__WXMAC__": 
            self.top_splitter.SetSashSize(7)
        
        # Initialize splitter separating problem set view from library
        # FIXME: re-enable once recipe stuff works
        self.library = library.LibraryPanel(self.top_splitter)
        #self.library = wx.Window(self.top_splitter)
        self.top_splitter.SplitVertically(self.splitter, self.library)
        self.top_splitter.SetSashGravity(1.0)
        self.top_splitter.SetMinimumPaneSize(75)
        self.top_splitter.SetSashPosition(700)
        (w,h) = self.top_splitter.GetSize()
        self.top_splitter.SetSashPosition(w - 300)
        
        if wx.Platform == "__WXMAC__": 
            self.top_splitter.SetSashSize(7)
        
        # Create results view book and the problem set log
        self.viewbook = lb = ResultViewBook(self.splitter, -1, problemSet)# views.ViewBook(self.splitter, -1, problemSet)        
        
        # Log control
        self.logCtrl = ws.SmTextCtrl(self.splitter, -1,
            style = wx.TE_MULTILINE|wx.TE_READONLY|wx.TE_RICH2) # doesn't work on mac: wx.HSCROLL)

        self.logCtrl.SetFont(ws.SmFont(wx.MODERN, wx.NORMAL, wx.NORMAL))
        if wx.Platform == "__WXMAC__": self.logCtrl.MacCheckSpelling(False)
        self.logCtrl.AppendText("\n".join(self.problemSet.GetCommandLog()))
        
        self.splitter.SplitHorizontally(self.viewbook, self.logCtrl)
        (w,h) = self.splitter.GetSize()
        self.splitter.SetSashPosition(w - 300)
        self.splitter.SetSashGravity(1.0)
        (w,h) = self.splitter.GetSize()
        self.SetSize((720,500))
        self.OnViewChange()#        self.SetMinSize((1000,500))

        # Hide the library and log to start (?)
        self.OnShowLibrary(None)
        self.OnShowCommandLog(None)
        
        # Handle viewbook page changing
        self.Bind(ws.EVT_RESULTBOOK_VIEW_CHANGED, self.OnViewChanged)
        
        # Handle operations on the current selection
        self.Bind(wx.EVT_MENU, self.OnCopy, id=ID_EDIT_COPY_RECS_AS_TEXT)
        self.Bind(wx.EVT_MENU, self.OnCopyAsLatex, id=ID_EDIT_COPY_RECS_AS_LATEX)
        self.Bind(wx.EVT_MENU, self.OnQuery, id=ID_QUERY_BUILDER)
        self.Bind(wx.EVT_MENU, self.OnImport, id=ID_IMPORT_RESULT)
        self.Bind(wx.EVT_MENU, self.OnRefine, id=ID_ANALYSIS_REFINE)
        self.Bind(wx.EVT_MENU, self.OnCount, id=ID_ANALYSIS_COUNT)
        self.Bind(wx.EVT_MENU, self.OnTrend, id=ID_ANALYSIS_TREND)
        self.Bind(wx.EVT_MENU, self.OnSet, id=ID_ANALYSIS_SET)
        self.Bind(wx.EVT_MENU, self.OnSetTool, id=ID_ANALYSIS_SET_TOOL)        
        self.Bind(wx.EVT_MENU, self.OnSetBuild, id=ID_ANALYSIS_SET_BUILD)     
        self.Bind(wx.EVT_MENU, self.OnShellTool, id=ID_ANALYSIS_SHELL_TOOL)
        self.Bind(wx.EVT_MENU, self.OnShowLibrary, id=ID_VIEW_SHOW_LIBRARY)
        self.Bind(wx.EVT_MENU, self.OnShowCommandLog, id=ID_VIEW_SHOW_COMMAND_LOG)
        self.Bind(wx.EVT_MENU, self.OnSaveInLibrary, id=ID_ANALYSIS_SAVE_IN_LIBRARY)
        self.Bind(wx.EVT_MENU, self.OnAddSetToLibrary, id=ID_ANALYSIS_ADD_SET_TO_LIBRARY)
        # FIXME: application should handle these
        self.Bind(wx.EVT_MENU, self.OnManageLibrary, id=ID_ANALYSIS_MANAGE_SET_LIBRARY)
        self.Bind(wx.EVT_MENU, self.OnManageCustomTools, id=ID_ANALYSIS_MANAGE_CUSTOM_TOOLS)
        
        # FIXME: real way of loading these
        for i in range(0,50):
            self.Bind(wx.EVT_MENU, self.OnCustomAnalysis, id=ID_ANALYSIS_CUSTOM+i)
        self.Bind(wx.EVT_MENU, self.OnSmartGraph, id=ID_ANALYSIS_SMART_GRAPH)
        self.Bind(wx.EVT_MENU, self.OnInfo, id=ID_SHOW_INFO)
        self.Bind(wx.EVT_MENU, self.OnViewOptions, id=ID_VIEW_OPTIONS)
        self.Bind(wx.EVT_MENU, self.OnOpenExcel, id=ID_OPEN_EXCEL)
        self.Bind(wx.EVT_MENU, self.OnShowFiles, id=ID_SHOW_FILES)
        self.Bind(wx.EVT_MENU, self.OnSave, id=ID_SAVE_PROBLEMSET)
        self.Bind(wx.EVT_MENU, self.OnSaveAs, id=ID_SAVE_AS)
        self.Bind(wx.EVT_MENU, self.OnCleanup, id=ID_CLEANUP)
        # Enable/disable toolbar and menu items
        for i in range(0,50):
            self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateCustomAnalysis, id=ID_ANALYSIS_CUSTOM+i)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateCopy, id=ID_EDIT_COPY_RECS_AS_TEXT)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateCopyAsLatex, id=ID_EDIT_COPY_RECS_AS_LATEX)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateSave, id=ID_SAVE_PROBLEMSET)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateAnalysis, id=ID_ANALYSIS_REFINE)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateAnalysis, id=ID_ANALYSIS_COUNT)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateAnalysis, id=ID_ANALYSIS_TREND)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateAnalysis, id=ID_ANALYSIS_SET)  
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateSetTool, id=ID_ANALYSIS_SET_TOOL)                
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateAnalysis, id=ID_ANALYSIS_SPLIT)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateAnalysis, id=ID_ANALYSIS_MERGE)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateGraph, id=ID_ANALYSIS_SMART_GRAPH)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateSaveInLibrary, id=ID_ANALYSIS_SAVE_IN_LIBRARY)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateAddSetToLibrary, id=ID_ANALYSIS_ADD_SET_TO_LIBRARY)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateOpenExcel, id=ID_OPEN_EXCEL)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateShowFiles, id=ID_SHOW_FILES)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateInfo, id=ID_SHOW_INFO)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateShowCommandLog, id=ID_VIEW_SHOW_COMMAND_LOG)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateShowLibrary, id=ID_VIEW_SHOW_LIBRARY)        
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateViewOptions, id=ID_VIEW_OPTIONS)
        self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateNoCommand, id=ID_NO_COMMAND)
        # Real way of loading these
        for i in range(0,10):
            self.Bind(wx.EVT_UPDATE_UI, self.OnUpdateCustomAnalysis, id=ID_ANALYSIS_CUSTOM+i)
        # Handle close event
        self.Bind(wx.EVT_SPLITTER_DCLICK, self.OnSplitterDoubleClick, self.splitter)
        self.Bind(wx.EVT_SPLITTER_UNSPLIT, self.OnSplitterUnSplit, self.splitter)
        self.Bind(wx.EVT_SPLITTER_SASH_POS_CHANGING,
            self.OnLibrarySplitterChanging, self.top_splitter)
        self.Bind(wx.EVT_ACTIVATE, self.OnActivate)
        self.Bind(wx.EVT_CLOSE, self.OnClose)
        # Added explicit handler so we get the familiar Command-W on Mac
        self.Bind(wx.EVT_MENU, self.OnCloseWindow, id=ID_CLOSE_WINDOW)
        # Status message from data set while it's loading
        self.Bind(ws.EVT_UPDATE_STATUS, self.OnResultChange)
        self.Bind(ws.EVT_RESULT_CHANGED, self.OnResultChange)
        # Log message
        self.Bind(ws.EVT_LOG_COMMAND, self.OnLogCommand)
        # Background image
        #self.Bind(wx.EVT_ERASE_BACKGROUND, self.OnEraseBackground)
        #self.bg_bmp = wx.Bitmap("icons/background.gif", wx.BITMAP_TYPE_GIF)
        self.RefreshContents()
        
    def OnUpdateNoCommand(self, evt):
        evt.Enable(False)
        
    def RefreshContents(self):
        title =  "iSiLK %s - %s - %s@%s:%s" % (wx.GetApp().GetVersion(),
            self.problemSet.GetName(),
            #os.path.split(self.problemSet.GetLocalDir())[1],
            self.problemSet.GetRemoteUser(),
            self.problemSet.GetRemoteHost(),
            self.problemSet.GetRemoteDir())
        self.SetTitle(title)
        self.viewbook.RefreshContents()

    def OnUpdateCopy(self, evt):
        view = self.viewbook.GetSelectedView()
        enable = self.GetSelection().IsLocalFileAvailable() \
            and view.IsClipboardFormatAvailable(CLIPBOARD_FORMAT_DEFAULT)
        evt.Enable(enable)

    def OnUpdateCopyAsLatex(self, evt):
        view = self.viewbook.GetSelectedView()
        enable = self.GetSelection().IsLocalFileAvailable() \
            and view.IsClipboardFormatAvailable(CLIPBOARD_FORMAT_LATEX_TABLE)
        print "enable=%s" % enable
        evt.Enable(enable)

    def OnCopy(self, evt):
        print "Explorer on copy"
        view = self.viewbook.GetSelectedView()
        data = view.GetClipboardObject()
        if wx.TheClipboard.Open():
            wx.TheClipboard.SetData(data)
            wx.TheClipboard.Close()
            
    def OnCopyAsLatex(self, evt):
        view = self.viewbook.GetSelectedView()
        data = view.GetClipboardObject(CLIPBOARD_FORMAT_LATEX_TABLE)
        if wx.TheClipboard.Open():
            wx.TheClipboard.SetData(data)
            wx.TheClipboard.Close()

    def OnSplitterDoubleClick(self, evt):
        print "Double click on splitter"
        evt.Skip()
    
    def OnLibrarySplitterChanging(self, evt):
        evt.Skip(True)
        
    def OnSplitterUnSplit(self, evt):
        print "Splitter unsplit"
        self.logCtrl.Show(False)
        evt.Skip(True)
        
    def OnEraseBackground(self, evt):
        dc = evt.GetDC()
        if not dc:
            dc = wx.ClientDC(self)
            rect = self.GetUpdateRegion().GetBox()
        sz = self.GetClientSize()
        w = self.bg_bmp.GetWidth()
        h = self.bg_bmp.GetHeight()
        x = 0
        while x < sz.width:
            y = 0
            while y < sz.height:
                dc.DrawBitmap(self.bg_bmp, x, y)
                y = y + h
            x = x + w
        
    def GetLogCtrl(self):
        return self.logCtrl

    def GetProblemSet(self):
        return self.problemSet
        
    def GetSelection(self):
        return self.viewbook.GetSelectedData() 
        
    def OnLogCommand(self, evt):
        self.DoLogMessage(evt.cmd)
        
    def DoLogMessage(self, msg):
        print "Logging '%s'" % msg
        if msg == self._last_log_message:
            return
        self._last_log_message = msg
        msg = msg.rstrip("\n")
        msg = silk.format_log_message(self.problemSet.GetRemoteUser(),
            self.problemSet.GetRemoteHost(), msg)
        # FIXME: add syntax highlighting back?
        if False:
            start_cmd = self.logCtrl.GetInsertionPoint()
        self.logCtrl.AppendText(msg + '\n')
        if False:
            # FIXME: do some sort of syntax highlighting?
            end_cmd = msg.find(" ")
            if end_cmd == -1:
                end_cmd = start_cmd + len(msg)-1
            else:
                end_cmd = start_cmd + end_cmd
            cmd_str=msg.split(" ")[0]
            if cmd_str in ["mapsid"] or cmd_str[0:2] == "rw":
                f = ws.SmFont(wx.MODERN, wx.NORMAL, wx.BOLD)
                style_cmd = wx.TextAttr("BLUE", font=f)
            else:
                f = ws.SmFont(wx.MODERN, wx.NORMAL, wx.BOLD)
                style_cmd = wx.TextAttr("BLACK", font=f)            

            self.logCtrl.SetStyle(start_cmd, end_cmd, style_cmd)
            
    def OnUpdateViewOptions(self, evt):
        # FIXME: The view should really be deciding whether it's got view options and how to display them
        if self.GetSelection().IsFilterResult() or self.GetSelection().IsToolResult() or self.GetSelection().IsSetResult():
            evt.Enable(True)
        else:
            evt.Enable(False)

    def OnUpdateSave(self, evt):
        evt.Enable(self.problemSet.IsModified())
        
    def OnUpdateAnalysis(self, evt):
        data = self.viewbook.GetSelectedData() 
        state = True 
        if data != None:
            state = data.IsFilterResult()
        evt.Enable(state)        
        if evt.GetId() == ID_ANALYSIS_REFINE:
            self.filterToolMenu.Enable(state)
        elif evt.GetId() == ID_ANALYSIS_TREND:
            self.countToolMenu.Enable(state)
        elif evt.GetId() == ID_ANALYSIS_COUNT:
            self.uniqToolMenu.Enable(state)
        elif evt.GetId() == ID_ANALYSIS_SET:
            self.setToolMenu.Enable(True)
        
    def OnUpdateSetTool(self, evt):
         data = self.viewbook.GetSelectedData() 
         evt.Enable((data!=None) and data.IsSetResult())
         
    def OnUpdateSaveInLibrary(self, evt):
        data = self.viewbook.GetSelectedData()
        evt.Enable((data!=None) and not isinstance(data, ProblemSet) and not isinstance(data.tool, ImportTool))

    def OnUpdateAddSetToLibrary(self, evt):
        data = self.viewbook.GetSelectedData()
        evt.Enable((data!=None) and data.IsSetResult())
        
    def OnUpdateOpenExcel(self, evt):
        data = self.viewbook.GetSelectedData()
        if data:
            if data.GetProblemSet() != data and not data.IsGraphResult():
                evt.Enable(data.IsLocalFileAvailable())
            else:
                evt.Enable(False)
        else:
            evt.Enable(False)
        
    def OnUpdateShowFiles(self, evt):
        data = self.viewbook.GetSelectedData()
        evt.Enable((data!=None))
        
    def OnUpdateGraph(self, evt):
        data = self.viewbook.GetSelectedData() 
        evt.Enable((data!=None) and data.IsToolResult() and data.IsLocalFileAvailable())
        
    def OnUpdateInfo(self, evt):
        evt.Enable(True)
        
    def OnImport(self, evt):
        from filemanager import FilePicker
        dlg = FilePicker(self, cancel_label="Cancel", 
            ext=[".set",".rwf",".bin"],
            showLocation = False,
            libraryOnly = True,
            title="Choose a file to import (flow data or IP set file)")
        if dlg.ShowModal() == wx.ID_OK:
            filename = dlg.GetResult()
            # FIXME: add silk.filetype method and use that
            if os.path.splitext(filename)[1] == ".rwf":
                q = analysistool.ImportFile(self.problemSet, filename,
                    "Imported Query Result", QueryResult)
            elif os.path.splitext(filename)[1] == ".set":
                q = analysistool.ImportFile(self.problemSet, filename,
                    "Imported IP Set", SetResult)
            return
        else:
            return
        
        dlg = wx.TextEntryDialog(self, "Enter a path on the remote analysis server", "Enter path")
        dlg.CenterOnParent()
        filename = None
        if dlg.ShowModal() == wx.ID_OK:
            filename = dlg.GetValue()
        else:
            return
        print "filename = %s" % filename
        context = ResultRunContext(self.problemSet, True)
        f = silk.exec_command(context, "ls %s" % filename)
        while True:
            print "Reading line"
            l = f.readline()
            if not l: break
            print "Read line: %s" % l
            
        q = analysistool.ImportFile(self.problemSet, filename,
            "Imported Query", QueryResult)
        
    def OldImportFile(self, filename):
        if os.path.splitext(filename)[1] == ".rwf":
            r = FilterResult(self.problemSet)
            r.SetName("Imported Filter Result")            
        elif os.path.splitext(filename)[1] == ".asc":
            tempCache = SimpleDataCache()
            tempCache._cacheData(filename)
            if "date" in tempCache.fields:
                r = CountResult(self.problemSet)
                r.SetName("Imported Trend Result")                            
            else:
                r = UniqResult(self.problemSet)
                r.SetName("Imported Count Result")            
        else:
            raise Exception, "Can't import file"    
        if wx.GetApp():
            wx.GetApp().AddViews(r)
        context = ResultRunContext(r, runRemote=False)
        if os.path.splitext(filename)[1] == ".rwf":
            r.fileinfo = silk.fileinfo(context, filename)
        r.BindFile(filename, copyFile=True)
        
    def OnCleanup(self, evt):
        msg = '''Delete any orphaned files in the problem set directory. '''
        d = wx.MessageDialog(self, msg, "Feature not yet implemented", wx.ICON_INFORMATION)
        d.ShowModal()
    
    def OnSave(self, evt):
        self.problemSet.Save()
        
    def OnSaveAs(self, evt):
        print "In OnSaveAs"
        class EditNameDialog(wx.Dialog):
            def __init__(self, parent, label, value, *args, **kargs):
                wx.Dialog.__init__(self, parent, *args, **kargs)
                self.label = wx.StaticText(self, -1, label)
                self.edit = wx.TextCtrl(self, -1, "", validator=ws.NameValidator(), size=(125,-1))
                self.edit.SetValue(value)
                szDlg = wx.BoxSizer(wx.VERTICAL)
                szBtns = self.CreateButtonSizer(wx.OK | wx.CANCEL)
                self.SetEscapeId(wx.ID_CANCEL)
                self.FindWindowById(wx.ID_OK).SetDefault()
                self.FindWindowById(wx.ID_OK).SetLabel("Set Name")
                szDlg.Add(self.label, flag=wx.EXPAND|wx.TOP|wx.LEFT|wx.RIGHT, border=8)        
                szDlg.Add(self.edit, flag=wx.EXPAND|wx.ALL, border=8)
                szDlg.Add(szBtns, flag=wx.EXPAND|wx.ALL, border=8)
                self.edit.SetFocus()
                self.SetSizerAndFit(szDlg)
            def GetValue(self):
                return self.edit.GetValue()
        print "self = %s" % self
        dlg = EditNameDialog(self, "Problem Set Name", self.problemSet.GetName(),
            title="Rename problem set")
        if dlg.ShowModal() == wx.ID_OK:
            newName = dlg.GetValue()
            self.problemSet.SetName(newName)
    
    def OnActivate(self, evt):
        ProblemSetExplorer._lastActive = self
        evt.Skip()
                
    def OnCloseWindow(self, evt):
        self.Close()
        
    def CanClose(self):
        print "In CanClose"
        return not (self.problemSet.HasActiveChildren() \
                    or self.problemSet.IsModified())                    
        
    def OnClose(self, evt):
        if self.problemSet.HasActiveChildren():
            d = wx.MessageDialog(self, 
                "There are tools currently running.\
Are you sure you want to exit?", "Confirm",
                 wx.YES_NO | wx.ICON_EXCLAMATION)
            result = d.ShowModal()
            if result != wx.ID_YES:
                print "NOT EXITING"
                return
        if self.problemSet.IsModified():
            d = wx.MessageDialog(self, 
                '''Some of your changes may be lost if you don't save them.
Do you want to save the changes you made the problem set?''', "Confirm",
                 wx.YES_NO | wx.CANCEL | wx.ICON_EXCLAMATION)
            result = d.ShowModal()
            if not result == wx.ID_CANCEL:
                if result == wx.ID_YES:
                    self.problemSet.Save()
                wx.GetApp().ExplorerClosed(self)
                evt.Skip(True)
            return
        print "Closing window"    
        wx.GetApp().ExplorerClosed(self)
        evt.Skip(True)
        print "At the end of OnClose"

    def NewView(self, dataset):
        if not isinstance(dataset.GetParent(), ProblemSet):
            try:
                self.viewbook.SelectDataSet(dataset.GetParent())
            except:
                raise RuntimeError, \
                    "Internal Error: can't create view for '%s'" % dataset.GetName() + \
                    "  Can't find view of parent '%s'." % dataset.GetParent().GetName()
        self.viewbook.NewView(dataset)
        (w,h) = self.GetSize()
        self.SetSize((w+10,h+10))
        self.SetSize((w,h))
        self.OnViewChange()
        
    def SelectDataSet(self, dataset):
        self.viewbook.SelectDataSet(dataset)
    
    def PopulateTree(self, dataset):
        # FIXME: For some reason I'm not getting refinements of 
        # refinements... probably a parsing issue
        self.viewbook.NewView(dataset)
        self.viewbook.SelectDataSet(dataset)
        for resultset in dataset.children:
            if isinstance(resultset, FilterResult):
                self.viewbook.NewView(resultset)
                self.PopulateTree(resultset)
            if isinstance(resultset, CountResult):
                self.viewbook.NewView(resultset)
            if isinstance(resultset, UniqResult):
                self.viewbook.NewView(resultset)
                
    def OnViewChanged(self, event):
        self.OnViewChange()
        data = self.viewbook.GetSelectedData() 
        #self.library.UpdateContents(data)
        
        return
        if wx.GetApp().library:
            data = self.viewbook.GetSelectedData() 
            wx.GetApp().library.UpdateContents(data)


    # FIXME: should be real event handler not callback from view
    def OnViewChange(self):
        data = self.viewbook.GetSelectedData() 
        view = self.viewbook.GetSelectedView()
        # Not yet initialized
        if self.viewbook.GetViewCount() == 0:
            self.viewbook.Show(False)
            self.splitter.Show(False)
            self.logCtrl.Show(False)
            self.statusbar.SetStatusText("Empty problem set")
            return
        self.splitter.Show(True)
        # Only show the log window if the user hasn't hidden it
        if self.splitter.GetWindow2() != None:
            self.logCtrl.Show(True)
        self.viewbook.Show(True)
        # Update the status bar to reflect new view selected
        if view:
            status_str = view.GetStatusString()
            self.statusbar.SetStatusText(status_str)
              
    def OnResultChange(self, evt):
        if evt.data == self.problemSet:
            self.RefreshContents()
        if evt.data.IsDeleted():
            self.viewbook.DeleteResult(evt.data, showConfirmation=False)
        self.viewbook.SetPageText(self.viewbook.GetPageIndex(evt.data), 
            evt.data.GetName())
        print "OnResultChange - %s" % evt
        self.OnViewChange()

    def OnExportAsScript(self, evt):
        pass

    def OnExportAsTool(self, evt):
        dataset = self.viewbook.GetSelectedData() 
        dataset.ExportAsTool("test.py")
                        
    def OnQuery(self, evt):
        wx.GetApp().DoShowQueryBuilder()
        
    def ShowDialogAndRun(self, analysisClass, **kargs):
        dataset = self.viewbook.GetSelectedData()
        if False:
            dlg = analysisClass.CreateDialog(dataset, self,
                style=wx.RESIZE_BORDER|wx.DEFAULT_DIALOG_STYLE)
            opts = kargs
            if dlg:
                dlg.CenterOnParent()
                rv = dlg.ShowModal()
                if rv != wx.ID_OK: return
                opts = dlg.GetOptions()
        result = analysistool.RunTool(config.silk.workRemotely(),
               analysisClass, dataset) # **opts
        return result
        
    def OnRefine(self, evt):
        opts = {}
        data = self.viewbook.GetSelectedData() 
        result = analysistool.RunTool(config.silk.workRemotely(), True,
                    tools.filter.FilterTool, data, **opts)
              
    def OnTrend(self, evt):
        opts = { "bin-size" : "300"}
        data = self.viewbook.GetSelectedData() 
        result = analysistool.RunTool(config.silk.workRemotely(), True,
                    tools.count.CountTool, data, **opts)
        
    def OnSet(self, evt):
        opts = { "sip-file": "$output"}
        data = self.viewbook.GetSelectedData() 
        result = analysistool.RunTool(config.silk.workRemotely(), True,
                    tools.set.SetTool, data, **opts)

    def OnSetTool(self, evt):
        opts = { "union": True}
        data = self.viewbook.GetSelectedData() 
        result = analysistool.RunTool(config.silk.workRemotely(), True,
                    tools.settool.SetToolTool, data, **opts)                  

    def OnSetBuild(self, evt):
        opts = { "union": True}
        data = self.viewbook.GetSelectedData() 
        result = analysistool.RunTool(config.silk.workRemotely(), True,
                    tools.settool.SetBuildTool, data, **opts)                  

    def OnCount(self, evt):
        opts = { "fields" : "sport", "bytes" : True}
        data = self.viewbook.GetSelectedData() 
        result = analysistool.RunTool(config.silk.workRemotely(), True,
               tools.uniq.UniqTool, data, **opts)
               
    def OnSaveInLibrary(self, evt):
        data = self.viewbook.GetSelectedData()
        recipe = analysistool.CreateRecipeFromResult(data, self)
        if recipe:
            library.GetLibrary().addAnalysis(recipe)
                
    def OnAddSetToLibrary(self, evt):
        data = self.viewbook.GetSelectedData()
        dlg = wx.SingleChoiceDialog(self, "Choose a directory in your library path",
            "Add a Set", library.GetLibrary().getLibraryPath())
        rv = dlg.ShowModal()
        if rv == wx.ID_OK:
            directory = dlg.GetStringSelection()
            library.GetLibrary().addSet(data, directory)
            
    def OnCreateCustomTool(self, evt):
        data = self.viewbook.GetSelectedData()
        cls = ShowCreatePluginDialog(self)
        if cls:
            RegisterPlugin(cls._GetToolName(),
                "icons/tool.gif", cls, cls.GetCategory())
                
    def OnShellTool(self, evt):
        parentResult = self.viewbook.GetSelectedData()
        cls = CreateCustomPlugin(self, parentResult=parentResult)
        if cls != None:
            opts = cls._GetDefaultArgs()
            analysistool.RunTool(config.silk.workRemotely(), False,
                        cls, parentResult, **opts)
                
    def OnManageCustomTools(self, evt):
        ManageCustomPlugins(self)
        
    def OnManageLibrary(self, evt):
        ext=[".pmap", ".set", ".bin"]
        dlg = FileManager(self,  ext=ext, title="Manage Library")
        rv = dlg.ShowModal()
        
    def OnShowLibrary(self, evt):
        if self.top_splitter.GetWindow2() == None:
            self.top_splitter.SplitVertically(self.splitter, self.library)
            self.library.Show(True)
            (w,h) = self.top_splitter.GetSize()
            self.top_splitter.SetSashPosition(w - 300)
        else:
            self.library.Show(False)
            self.top_splitter.Unsplit(self.library)
            
    def OnShowCommandLog(self, evt):
        if self.splitter.GetWindow2() == None:
            self.splitter.SplitHorizontally(self.viewbook, self.logCtrl)
            self.logCtrl.Show(True)
            (w,h) = self.splitter.GetSize()
            self.splitter.SetSashPosition(w - 300)
        else:
            self.logCtrl.Show(False)
            self.splitter.Unsplit(self.logCtrl)
            
    def OnUpdateShowLibrary(self, evt):
        evt.Check(self.top_splitter.GetWindow2() != None)        
                    
    def OnUpdateShowCommandLog(self, evt):
        evt.Check(self.splitter.GetWindow2() != None)        

    def OnUpdateCustomAnalysis(self, evt):
        data = self.viewbook.GetSelectedData() 
        plugin_id = evt.GetId() - ID_ANALYSIS_CUSTOM
        tool = tools.FindPluginById(plugin_id)
        plugin_class = tool.toolClass
        evt.Enable(plugin_class.CanRun(data))
        return

    def OnCustomAnalysis(self, evt):
        tb = self.GetToolBar().FindById(ID_ANALYSIS_CUSTOM)
        print "evt = %s" % evt
        print "tb = %s" % tb
        #print "ctrl = %s" % tb.GetControl()
        data = self.viewbook.GetSelectedData() 
        plugin_id = evt.GetId() - ID_ANALYSIS_CUSTOM
        tool = tools.FindPluginById(plugin_id)
#        tool = tools.GetPlugins()[plugin_index]
        plugin_class = tool.toolClass
        print plugin_class
        result = analysistool.RunTool(config.silk.workRemotely(), True,
               plugin_class, data)
        
    def OnInfo(self, evt):
        print evt
        print dir(evt)
        self.SetFocus()
        dataset = self.viewbook.GetSelectedData()
        pos = self.GetScreenPosition()
        pos = (pos.x + 50, pos.y + 50)
#        info = InfoFrame(dataset, self, "Info for %s" % dataset.GetName(), pos = pos)
        info = InfoFrame.ShowPanel(dataset, self, pos=pos)
#        info.Show(True)
        self.viewbook.splitter.GetWindow2().Enable(True)
        self.viewbook.splitter.GetWindow2().Show(True)
        
    def OnViewOptions(self, evt):
        dataset = self.viewbook.GetSelectedData()        
        d = RecordViewOptionsDialog(opts=dataset.GetViewOptions(),
            parent=self, title="View Options")
        d.CenterOnParent()
        d.ShowModal()
        self.viewbook.RebuildView()
        
    def OnShowFiles(self, evt):
        wx.GetApp().ShowFile(self.viewbook.GetSelectedData().GetLocalDir())

    def CreateGraph(self, cls, parent, name=None, **kargs):
        resultset = self.viewbook.GetSelectedData()
        f = resultset.GetFields()
        if False and isinstance(parent, UniqResult):
            if "bytes" in f: title = "Volume in Bytes"
            elif "packets" in f: title = "Volume in Packets"            
            elif "records" in f: title = "Volume in Flow Counts"
            elif "unique_sip" in f: title = "Unique Source IP Addresses"
            elif "unique_dip" in f: title = "Unique Destination IP Addresses"
            else: title = name
        else:
            title = name
        graph = cls(parent, name=title, **kargs)
        graph.AddChangeListener(self)
        graph.RunAnalysis()
        wx.GetApp().AddViews(graph)
        #self.viewbook.NewView(graph)
        return graph
        
    def OnSmartGraph(self, evt):
        resultset = self.viewbook.GetSelectedData()

        if isinstance(resultset, UniqResult):
            yscale = 'log'
            orientation = 'horizontal'
            max_rows = 10
            title_prefix = "Top %d talkers" % max_rows
            if "sip" in resultset.GetFields() or "dip" in resultset.GetFields() or "nhIP" in resultset.GetFields():
                orientation = 'vertical'
            graph = []
            if "unique_sip" in resultset.GetFields():
                graph.append(self.CreateGraph(CountGraph, resultset,
                    name="Unique Source IPs", y_fields="unique_sip",max_rows = max_rows,
                    orientation=orientation, yscale=yscale, title_prefix=title_prefix))
            if "unique_dip" in resultset.GetFields():                    
                graph.append(self.CreateGraph(CountGraph, resultset,
                    name="Unique Dest IPs", y_fields="unique_dip",max_rows = max_rows,
                    orientation=orientation, yscale=yscale, title_prefix=title_prefix))
            if "bytes" in resultset.GetFields():
                graph.append(self.CreateGraph(CountGraph, resultset,
                    name="Bytes", y_fields="bytes",max_rows = max_rows,
                    orientation=orientation, yscale=yscale, title_prefix=title_prefix))
            if "packets" in resultset.GetFields():
                graph.append(self.CreateGraph(CountGraph, resultset,
                    orientation=orientation, max_rows = max_rows,
                    name="Packets", y_fields="packets", yscale=yscale, title_prefix=title_prefix))
            if "records" in resultset.GetFields():
                graph.append(self.CreateGraph(CountGraph, resultset, 
                    orientation=orientation, max_rows = max_rows,
                    name="Flow Records", y_fields="records", yscale=yscale, title_prefix=title_prefix))
            self.SelectDataSet(graph[0])
        elif isinstance(resultset, CountResult):
            graph = []
            max_rows = resultset.GetRecordCount()
            title_prefix = "Time series graph (%d bins)" % max_rows
            for f in ["bytes", "packets", "records"]:
                if f in resultset.GetFields():
                    graph.append(self.CreateGraph(CountGraph,  resultset, 
                        name = "Time Series Graph - %s" % f, y_fields=f,
                        orientation='horizontal', yscale='linear', max_rows=max_rows,
                        sort_data=False, bar_thickness=1.0, title_prefix = title_prefix))
            self.SelectDataSet(graph[0])
        else:
            raise RuntimeError("Unexpected argument to OnSmartGraph")
        
    def OnOpenExcel(self, evt):
        sel = self.viewbook.GetSelectedData()
        path_in = sel.GetLocalAscPath()
        path_out = sel.GetLocalPath() + ".csv"
        f = open(path_in)
        fout = open(path_out, "w")
        while 1:
            l = f.readline()
            if not l: break
            l = re.sub("\|", ",", l)
            l = re.sub(" ", "", l)
            fout.write(l)
        f.close()
        fout.close()
        print "Output path is %s" % path_out
        wx.GetApp().OpenExcel(path_out)
