## ------------------------------------------------------------------------
## 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
## ------------------------------------------------------------------------

import logging
import wx, ws
import constants
import os, os.path, sys, string, tempfile, copy, time, traceback
import glob, pdb
import threading
import config.silk as config
import silk
from analysisresult import AnalysisResult, ResultRunContext
from graphresult import GraphResult
from tools.query import QueryTool
#from tools.filter import FilterTool
import tools.filter
import tools 
from tools.plugins import CreateGenericPlugin

def RunTool(runRemote, showDialog, toolClass, parentResult, name=None, **kargs):
    analysis = RunnableAnalysis(toolClass,  toolClass._GetToolName(), **kargs)
    print "name = %s" % name
    print "ToolName = %s" % toolClass._GetToolName()
    return analysis.Run(runRemote, showDialog, parentResult, name=name, **kargs)
    
# Run a function whose first argument is a context object
def RunFunction(runRemote, func, **kargs):
    analysis = RunnableFunction(func, **kargs)
    return analysis.Run(runRemote)
    
class ImportTool(object):
    def __init__(self, workerThread):
        self.workerThread = workerThread
    def _waitToComplete(self):
        while self.workerThread:
            print "Waiting to complete import"
            self.workerThread.join(5)
    
def ImportFile(parent, filename, name, resultClass):
    class _ImportThread (threading.Thread):
       def __init__(self, result, filename):
           threading.Thread.__init__(self)
           self.result = result
           self.filename = filename
       def run(self):
           _DoImport(self.result, self.filename)
           self.result.tool.workerThread = None
    parent.GetProblemSet().CreateDirectory()
    print "Importing %s" % filename
    r = resultClass(parent, name=name)
    if wx.GetApp():
        wx.GetApp().AddViews(r)
    workerThread = _ImportThread(r, filename)
    r.tool = ImportTool(workerThread)
    workerThread.start()
    return r

def _DoImport(r, filename):
    try:
        r._UpdateStatus(status="Importing remote file...")
        context = ResultRunContext(r, runRemote=True)
        silk.copy_remote_file(context,
            filename, r.GetRemotePath())
        if r.IsFilterResult():
            r._UpdateStatus(status="Running rwcut...")
            silk.cut(context, r.GetRemotePath(),
                r.GetRemotePath()+".asc")
        elif r.IsSetResult():
            r._UpdateStatus(status="Running rwsetcat...")
            silk.rwsetcat(context, r.GetRemotePath(),
                r.GetRemotePath()+".asc")
        r.fileinfo = silk.fileinfo(context,
            r.GetFileName())
        fileSize = silk.get_file_size(context,
            r.GetRemotePath() + ".asc")
        r.SetAscFileSize(fileSize)    
        r._UpdateStatus(constants.AVAILABLE_REMOTE)
        return r
    except Exception:
        wx.GetApp().ShowError("Import Error",
            "Unable to import remote file %s" % filename)
        context._UpdateStatus(constants.UNINITIALIZED,
             "Import failed")
    return None


class RunContext(object):
    def __init__(self, tool, runRemote=False):
        self.tool = tool
        self.result = tool.result
        self.sshOptions = self.result.GetSSHOptions()
        self.runRemote = runRemote
        if runRemote:
            self.runDir = self.result.GetRemoteDir()
        else:
            self.runDir = self.result.GetLocalDir()
    def GetSSHOptions(self):
        return self.sshOptions
    def _UpdateStatus(self, loc=None, status=None, percentComplete=None):
        self.result._UpdateStatus(loc, status, percentComplete)
        self.tool._UpdateStatus(loc, status, percentComplete)
    def _LogCommand(self, cmd):
        self.result._LogCommand(cmd)
        
        
class CreateRecipeDialog(wx.Dialog):
    def __init__(self, result, parent, *args, **kargs):
        wx.Dialog.__init__(self, parent, -1, *args, **kargs)
        self.result = result
        # Allow user to assign a name and see a preview
        self._createNameEditor()
        # Create buttons
        sz = wx.BoxSizer(wx.HORIZONTAL)
        sz.Add(wx.Button(self, wx.ID_CANCEL, "Cancel"), flag=wx.ALL, border=10)
        btnCreate = wx.Button(self, wx.ID_OK, "Create Recipe") #, (110,-1), (100,-1))
        sz.Add(btnCreate, flag=wx.ALL, border=10)
        # Create is default button
        btnCreate.SetDefault()
        # Escape key or close causes cancel
        self.SetEscapeId(wx.ID_CANCEL)
        szDlg = wx.BoxSizer(wx.VERTICAL)
        szDlg.Add(self.szName, flag=wx.ALL|wx.ALIGN_RIGHT, border=8)
        szDlg.Add(sz, flag=wx.ALL|wx.ALIGN_RIGHT, border=8)
        self.SetSizerAndFit(szDlg)

    def _createNameEditor(self):
        self.name = wx.TextCtrl(self, -1, "Untitled Recipe", 
            validator=ws.NameValidator(), size=(200,-1))
        self.szName = szName = wx.BoxSizer(wx.HORIZONTAL)
        szName.Add(ws.SmStaticText(self, -1, "Name"),
            flag=wx.RIGHT|wx.ALIGN_CENTER, border=6)
        szName.Add(self.name, flag=wx.ALIGN_CENTER|wx.RIGHT, border=10)
        
    def GetName(self):
        return self.name.GetValue()

def CreateRecipeFromResult(data, parentWindow):
    dlg = CreateRecipeDialog(data, parent=parentWindow,
        title="Create recipe")
    dlg.CenterOnParent()
    if dlg.ShowModal() == wx.ID_OK:
        return CompoundAnalysis(dlg.GetName(), [data])
    else:
        return None

class CompoundAnalysis(object):
    def __init__(self, name, results):
        self.name = name
        self.tree_list = []
        for result in results:
            self.tree_list.append(self._buildTree(result))
        
    def _str_(self, tree):
        (root, children) = tree
        str_val = str(root)
        for c in children:
            str_val += "\n%s" % self._str_(c)
        return str_val
        
    def __str__(self):
        str_val = ""
        for tree in self.tree_list:
            str_val += "\n%s" % self._str_(tree)
        return str_val
        
    def _buildTree(self, result):
        if result.GetTool().__class__.__name__ == "GenericTool":
            root = PluginAnalysis(result.GetTool().__class__,
                name=result.GetName(), **result.GetOptions())            
        else:
            root = RunnableAnalysis(result.GetTool().__class__,
                name=result.GetName(), **result.GetOptions())
        children = []
        for c in result.GetChildren():
            if isinstance(c, GraphResult):
                print "*** Skipping Graph Result ***"
                continue
            tree = self._buildTree(c)
            children.append(tree)
        return (root, children)
        
    def GetName(self):
        return self.name
        
    def SetName(self, name):
        name = name.replace(" ", "_")
        self.name = name
        
    def GetOpts(self):
        # FIMXE - proper behavior here?
        return {}
        
    def _run(self, tree, runRemote, parentResult, name=None):
        (root, children) = tree
        if name:
            result = root.Run(runRemote, False, parentResult, name=name)
        else:
            result = root.Run(runRemote, False, parentResult)
        for c in children:
            self._run(c, runRemote, result)

    def CanRun(self, result):
        #if result == None:
        #    return self.GetToolClass() == QueryTool
        for tree in self.tree_list:
            (root, children) = tree
            if not root.GetToolClass().CanRun(result):
                return False
        return True
            
    def GetToolClass(self):
        return self.tree_list[0][0].GetToolClass()
        
    def GetCategory(self):
        return "Compound %s" % self.GetToolClass().GetCategory()
            
    def Run(self, runRemote, showDialog, parentResult, name=None, **kargs):
        for tree in self.tree_list:
            self._run(tree, runRemote, parentResult, name)


class RunnableAnalysis(object):
    def __init__(self, toolClass, name = None, **kargs):
        if name == None:
            self.name = "Tool"
        else:
            self.name = name
        self.toolClass = toolClass
        self.opts = kargs
        
    def __str__(self):
        return "%s - %s\n%s" % (self.GetToolClass().GetCategory(),
            self.GetName(), 
            self.GetToolClass()._GetCommandLines("$input", "$output", **self.opts))
        
    def GetName(self):
        return self.name
        
    def SetName(self, name):
        name = name.replace(" ", "_")
        self.name = name
        
    def CanRun(self, result):
        if result == None:
            return self.toolClass == QueryTool
        return self.GetToolClass().CanRun(result)
        
    def GetToolClass(self):
        return self.toolClass
        
    def GetCategory(self):
        return self.GetToolClass().GetCategory()
        
    def GetOpts(self):
        return self.opts
        
    def Run(self, runRemote, showDialog, parentResult, name=None, **kargs):
        # FIXME: better way to do this?
        if self.GetToolClass() == QueryTool:
            parentResult = parentResult.GetProblemSet()
        # Override default options with specified options
        for k in kargs:
            self.opts[k] = kargs[k]
        # Get user-specified options
        if showDialog:
            # style=wx.RESIZE_BORDER|
            print "self.name = %s" % self.name
            t = "Running %s" % self.name
            dlg = self.GetToolClass().CreateDialog(parentResult, None, t) 
#                     STYLE=WX.DEFAULT_DIALOG_STYLE|WX.RESIZE_BORDER)
#                    style=wx.DEFAULT_DIALOG_STYLE)
            if dlg:
                dlg.SetOptions(self.opts)
                dlg.CenterOnParent()
                rv = dlg.ShowModal()
                if rv != wx.ID_OK: return
                self.opts = dlg.GetOptions()
                name = dlg.GetResultName()
        # Get input file name
        if isinstance(parentResult, AnalysisResult):
            inputFileName = parentResult.GetFileName()
        else:
            inputFileName = None
        # Create the tool instance we'll run
        tool = self.GetToolClass()(**self.opts)
        # Create the result associated with its output
        if not name:
            name = self.name
        if not name:
            resultName = tool._GetResultName()
        else:
            resultName = name
        # Instantiate a result
        print "Creating result instance with options %s" % self.opts
        tool.result = analysisResult = tool._GetResultClass()(parentResult,
            resultName, **self.opts)
        outputFileName = analysisResult.GetFileName()
        analysisResult.tool = tool
        # Add views in the user interface if we're running in application
        if wx.GetApp():
            wx.GetApp().AddViews(analysisResult)
        # Run the analysis within the context of the tool and result
        context = RunContext(tool, runRemote)
        # Start up the job thread
        _toolmanager._StartJob(context, inputFileName, outputFileName)
        # Return the result
        return analysisResult

# FIXME: is there a better way to do this?
# Hack to make recipes with plug in classes serializabe
class PluginAnalysis(RunnableAnalysis):
    def __init__(self, toolClass, name = None, **kargs):
        if name == None:
            self.name = "Tool"
        else:
            self.name = name
        
        self.str_toolclass = toolClass._SerializeToString()
        self.toolName = toolClass._GetToolName()
        self.defaultArgs = toolClass._GetDefaultArgs()
        self.category = toolClass.GetCategory()
        self.commandLine = toolClass._GetCommandLineTemplate()
        self.opts = kargs
        
    def GetToolClass(self):
        return CreateGenericPlugin(self.toolName, self.defaultArgs,
            self.category, self.commandLine)
        
class RunnableFunction(object):
    def __init__(self, func, **kargs):
        self.func = func
        self.opts = kargs

    def Run(self, runRemote, **kargs):
        for k in kargs:
            self.opts[k] = kargs[k]
        context = silk.SimpleRunContext(runRemote=True)
        print "Running function %s on args %s" % (self.func, self.opts)
        return self.func(context, **self.opts)


class ToolManager(object):
    def __init__(self):
        pass
    
    def _StartJob(self, context, inputFileName, outputFileName):
        class _JobThread (threading.Thread):
           def __init__(self, toolManager):
               threading.Thread.__init__(self)
               self.toolManager = toolManager
               self.context = context
           def run(self):
               try:
                   result = self.toolManager._runJob(context, inputFileName, outputFileName)
               except silk.CommandError, e:
                   context._UpdateStatus(constants.DELETED,
                        "Tool invocation failed")
                   if wx.GetApp():
                       wx.GetApp().ShowError(e.summary, e.details)
                   else:
                       print "(_JobThread) ERROR: %s\n\n%s" % (e.summary, e.details)
                   # Clean up the result
                   context.result.GetParent().DeleteChild(context.result)
               except Exception, e:
                   context._UpdateStatus(constants.DELETED,
                        "Tool invocation failed")
                   if wx.GetApp():
                       wx.GetApp().ShowError("Internal Error", str(e))
                   else:
                       print "INTERNAL ERROR: %s\n" % (str(e))
               self.context.tool.workerThread = None
        context.tool.workerThread = _JobThread(self)
        context.tool._PreAnalysis(context, inputFileName, outputFileName)
        context.tool.workerThread.start()
    
    def _runJob(self, context, inputFileName, outputFileName):
        tool = context.tool
        result = context.result
        parent = result.GetParent()
        # Is input file available to run on?
        available = (not context.runRemote and parent.IsLocalFileAvailable())
        available = available or (context.runRemote and parent.IsRemoteFileAvailable())        
        # If not, wait for the result to be available
        if not available:
            context._UpdateStatus(status="Waiting for '%s' to be available" % parent.GetName())
            logging.info("Waiting for %s to be available", parent.GetName())
            result.GetParent().GetTool()._waitToComplete()
        else:
            logging.info("Parent result '%s' is available locally", result.GetParent().GetName())
        status_msg = "Running tool"
        cmd = tool._GetCommandLines(inputFileName, outputFileName,
                **result.GetOptions())
        result.SetCurrentCommand(cmd)
        if context.runRemote:
            status_msg += " (Running remotely)"
        else:
            status_msg += " (Running locally)"
        context._UpdateStatus(loc=constants.RUNNING, status=status_msg)
        logging.info("%s: running command '%s'", tool.__class__.__name__, "<command info>")
        # Create directory if it doesn't already exist
        result.GetProblemSet().CreateDirectory()
        # Run the tool
        tool._RunAnalysis(context, inputFileName, outputFileName)
        # Get information about associated text file
        if result.IsToolResult():
            fileSize = silk.get_file_size(context, outputFileName)
            result.SetAscFileSize(fileSize)
        else:
            fileSize = silk.get_file_size(context, outputFileName + ".asc")
            result.SetAscFileSize(fileSize)            
        logging.info("%s: completed command '%s'", tool.__class__.__name__, "<command info>")        
        # Build command line information
#        cmdInfo = tool._GetCommandLines(inputFileName, outputFileName, **tool.opts)
        cmdInfo = tool.__class__._GetCommandLines(inputFileName, outputFileName, **tool.opts)        
        if result.IsFilterResult():
            result.fileinfo = silk.fileinfo(context, result.GetFileName())
            # Tool doesn't update command line info
            # FIXME: if/when the SiLK tools all write this information
            # use it.
            if not isinstance(tool, QueryTool) and not isinstance(tool, tools.filter.FilterTool):
                if cmdInfo:
                    result.fileinfo["command-lines"].append(cmdInfo)
        elif cmdInfo:
            if result.IsToolResult():
                result.fileinfo["count-records"] = silk.get_file_record_count(context, result.GetFileName())
            if result.IsSetResult():
                result.fileinfo["count-records"] = silk.get_file_record_count(context, result.GetFileName()+".asc")
            commandLines = [c for c in result.GetParent().GetFileInfo("command-lines")]
            commandLines = copy.copy(result.GetParent().GetCommandLines())
            result.fileinfo["command-lines"] = commandLines
            result.fileinfo["command-lines"].append("%s" % cmdInfo)
            for l in result.fileinfo["command-lines"]:
                print "  %s" % l
        if not context.runRemote:
            result._InitLocalCache()
        else:
            result._UpdateStatus(constants.AVAILABLE_REMOTE, status="Result available on remote host")
 
    def ExportCode(self, outputfile):
        outputfile.write("tool = %s(%s, parent, opts)" % \
                (self.__class__.__name__, self.result.__class__.__name__))
        ouptputfile.write("result_%s = tool.Run()" % self.result.uniqueId)
        
        
# First time module is loaded
if not "_toolmanager" in vars():
    _toolmanager = ToolManager()
