import re, os, math
import nuke, nukescripts


class musterSubmit(KellerNukePlugin):

    def configurePlugin(self):
        menubar = nuke.menu("Nuke")
        self.m = menubar.addMenu("&Render")
        self.m.addCommand("MusterSubmit", 'musterSubmit.submit()', "ctrl+r")
        
    def unconfigurePlugin(self):
        self.m.removeItem("MusterSubmit")

#############################
# GLOBALS
#############################

mstr_server = "192.168.0.251"
submission_user ="admin"
password = ""

mrtool = "\\\\hades\\p\\env\\muster9\\Mrtool.exe"
username = "admin"
pool = "NukeWin"
template = "821"
#nukePath="\\\\hades\\p\\env\\nuke\\_workgroup\\%s.%s\\" %(nuke.NUKE_VERSION_MAJOR,nuke.NUKE_VERSION_MINOR)
nukePath="\\\\hades\\p\\env\\nuke\\_workgroup\\%s\\" %nuke.NUKE_VERSION_MAJOR
# Nuke Version
nukeVer = nuke.env['NukeVersionString']
nukeExe = "Nuke" + nukeVer[:-2] + ".exe"

# Build env dict for the farm and ignore these entries...
ignoreEnvEntries = ['PATH', 'NUKE_TEMP_DIR', 'TMP', 'TEMP', 'OCIO', 'ADSKFLEX_LICENSE_FILE', 'FPS', 'SOLIDANGLE_LICENSE', 'foundry_LICENSE',
                    'peregrinel_LICENSE', 'keller_LICENSE', 'QUALM_LICENSE', 'PROGRAMFILES', 'APPDATA', 'COMMONPROGRAMW6432', 'PATHEXT', 'ALLUSERSPROFILE',
                    'PSMODULEPATH', 'PROCESSOR_IDENTIFIER', 'PROCESSOR_ARCHITECTURE', 'LOCALAPPDATA', 'PROCESSOR', 'GUID', 'PYTHONPATH', 'IDEA_INITIAL_DIRECTORY',
                    'ONEDRIVE', 'PYCHARM', 'PYCHARM_HOSTED', 'NUMBER_OF_PROCESSORS', 'PROCESSOR', 'HOMEPATH' ,'COMPUTERNAME', 'USERDOMAIN', 'USERPROFILE','COMPILE_0_1DEV']

############################



def get_broken_write_nodes():
    """
    Returns list of all write nodes that refer to non-existing paths (or an empty list).
    If nodes are selected only they will be checked.
    """

    nodes = nuke.selectedNodes() 
    if not nodes: nodes = nuke.allNodes()
    n = [i for i in nodes if i.Class() == "Write"]

    result = []
    for curNode in n:
        fileKnob = curNode['file']
        try:
            if not os.path.exists(os.path.dirname(fileKnob.evaluate())):
                result.append(curNode)
        except:
            pass

    return result

def fix_broken_write_nodes():
    """
    Create paths for all (selected) write nodes in case they don't exist yet.
    """

    nodes = get_broken_write_nodes()
    for curNode in nodes:
        thePath = os.path.dirname(curNode['file'].evaluate())
        if thePath and not os.path.exists(thePath):
            try:
                os.makedirs(thePath)
                print ("musterSubmit: Info. Created directory "+thePath)

            except:
                nuke.message("musterSubmit: Error. Error creating "+thePath)

def _replaceKnobPath(fileknob):
    infile = os.path.normpath(fileknob.getValue()).strip()
    infile = infile.replace("\\", "/")
    infile = infile.replace("p:/", "//hades/p/")
    infile = infile.replace("P:/", "//hades/p/")
    infile = infile.replace("o:/", "//calculon/o/")
    infile = infile.replace("O:/", "//calculon/o/")

    return infile

def _replacePathInNodes(root):
    for node in root.nodes():
        if isinstance(node, nuke.Group):
            _replacePathInNodes(node)
        elif node.Class() in ["Read", "Write"]:

            newFile = _replaceKnobPath(node["file"])
            node["file"].setValue(newFile)
            print ("musterSubmit: Info. Checking file Path of node %s: %s" % (node.name(), newFile))

def replaceNodeFilePaths():
    r = nuke.root()
    _replacePathInNodes(r)



def submit(selectedNodes = None):
    """
    Submits a script to Muster. Cross-Platform (Win and MacOS)
    """
    replaceNodeFilePaths()
    # comp saved?
    if not nuke.knob("root.name") or nuke.modified():
        if nuke.ask("Script needs to be saved before rendering.\nSave and continue?"):
            try:
               nuke.scriptSave()
            except RuntimeError:
               # if save dialog was cancelled and script is not untitled we could try to render anyways
               if not nuke.knob("root.name") or not nuke.ask("Could not save script. Try to render anyways?"):
                  return
        else:
            return

    # get configs from scene
    # TODO: check with KEnv
    # ocio
    if nuke.root()['OCIO_config'].value() == 'custom':
        if nuke.root()['customOCIOConfigPath'].value() != "":
            ocioConfig = nuke.root()['customOCIOConfigPath'].value()
            ocioConfig = ocioConfig.replace("/","\\")
    else:
        ocioConfig = '\\env\\nuke\\_run\\Nuke\\' + nuke.NUKE_VERSION_STRING + '\\plugins\\OCIOConfigs\\configs\\' + \
                         nuke.root()['OCIO_config'].value() + '\\config.ocio'

    # Nuke Workgroup
    try:
        nukePath=os.environ['NUKE_PATH']
    except:
        # fallback to gen workgroup
        pass


    # get default render range
    a = nuke.knob("first_frame")
    b = nuke.knob("last_frame")
    start = int(a)
    end = int(b)
    incr = 1
    rangeStr = a+"-"+b

    # set packet size automatically to use all available nodes
    # set _max_instances to current size of render farm
    # OVERRIDING PACKET SIZE TO 1, OTHERWISE RENDER FAILS TOO OFTEN!
    #_max_instances = 6
    #_max_psize = 10
    psize = 4

    # find packet size so that there is a small remainder
    #_smallest_remainder = 100
    #for i in range(_max_psize+1, 1):
    #    _remainder = (end-start+1) % (i * _max_instances)
    #    if _remainder <= _smallest_remainder:
    #        _smallest_remainder = _remainder
    #        psize = i
    # packet size not smaller than 1/15th of frame range
    #psize = max( min(_max_psize, int((end-start+1)/15)), psize)

    #Priority:
    prio = 10

    #Threads
    threads = 16

    #Cache-Limit
    cachelimit = 8192

    #Default Folder NukeWin
    parentID = 147
    pool = "NukeWin"

    #One Job per Write
    parallel = False

    # build panel
    panel = nuke.Panel("Submit to Muster")
    panel.setWidth(500)
    panel.addSingleLineInput("Frames to render:", rangeStr)
    panel.addSingleLineInput("Packet size:", psize)
    panel.addSingleLineInput("Max. instances:", "")
    panel.addSingleLineInput("No. of threads:", threads)
    panel.addSingleLineInput("Ram-Cache Limit:", cachelimit)
    panel.addSingleLineInput("Priority:", prio)
    panel.addSingleLineInput("RenderPool:", pool)
    panel.addEnumerationPulldown("Muster Folder:","NukeWin HOT")
    panel.addSingleLineInput("Workgroup:", nukePath)
    panel.addBooleanCheckBox("Submit one job per Write Node", parallel)

    #panel.addSingleLineInput("OCIO Config:", ocioConfig)
    #

    panel.addButton("Cancel")
    panel.addButton("OK")

    if panel.show() == 0:
        return

    mfolder = panel.value("Muster Folder:")
    if mfolder == "HOT":
        parentID = 29037
    elif mfolder == "NukeWin":
        parentID = 147
    else:
        pass

    mpool = panel.value("RenderPool:")
    if mpool is not None:
        pool = mpool
    else:
        pass

    rangeStr = panel.value("Frames to render:")
    if rangeStr is not None:
        rangeStr = re.sub("[-|-|/]", " ", rangeStr)
        t = rangeStr.split()
    else:
        return
    if len(t) > 0:
        start = int(t[0])
        end = start
    else:
        return
    if len(t) > 1: end = int(t[1])
    if len(t) > 2: incr = int(t[2])

    psize = panel.value("Packet size:")
    if psize is not None:
        try:
            psize = int(psize)
            psize = max(1, psize)
        except ValueError:
            psize = 4
    else:
        psize = 4

    maxi = panel.value("Max. instances:")
    try:
        maxistr = "-max %i" % (max(0, int(maxi)))
    except ValueError:
        maxistr = ""

    prio = panel.value("Priority:")
    if prio is not None:
        try:
            prio = int(prio)
        except ValueError:
            prio = 10
    else:
        prio = 10

    threads = panel.value("No. of threads:")
    if threads is not None:
        try:
            threads = int(threads)
        except ValueError:
            threads = 16
    else:
        threads = 16
    
    cache = panel.value("Ram-Cache Limit:")
    if cache is not None:
        try:
            cache = int(cache)
        except ValueError:
            cache = 20000
    else:
        cache = 20000


    parallel = panel.value("Submit one job per Write Node")
    if parallel is None:
        parallel is False

    currNk = nuke.value("root.name")
    nukePath = panel.value("Workgroup:")

    if nuke.env["WIN32"]:
        currNk = re.sub(r"^[pP]:", r"//hades/p", currNk)
        currNk = re.sub(r"^[oO]:", r"//calculon/o", currNk)
        currNk = os.path.normpath(currNk)
    currNk = currNk.replace("\\","/")

    # -s  server
    # -n  job name
    # -u  user name
    # -sf start frame
    # -b  batch submit
    # -ef end frame
    # -pk packet size
    # -bf by frame (frame stepping)
    # -e  engine template ID
    # -st numbering step
    # -f  file name
    # -pool destination pool
    # -pr priority
    # -max  maximum nr of instances
    # -parent 39401 // folder ID

    # jobname for later use
    jobname = os.path.basename(nuke.value("root.name")).split(".nk")[0]
    jobtmp = "replaceMe"

    #envDict = {"horst":"hemke","peter":"wolf"}

    envDict = {key: val for key, val in os.environ.items() if key not in ignoreEnvEntries}

    cmd = '%s -s %s -u %s -b -e %s -n %s' % (mrtool, mstr_server, username, template, jobtmp)
    cmd += ' -sf %i -ef %i -bf 1 -st 1' % (int(start), int(end))
    cmd += ' -pk %i -pr %i %s -pool "%s" -parent %i' % (int(psize), int(prio), maxistr, pool, int(parentID))
    cmd += ' -attr "NUKEVERSION" %s 0 -attr "NUKEPATH" %s 0 -attr "NUKERENDERTHREADS" %s 0' % (nukeVer, nukePath, int(threads))
    cmd += ' -attr "NUKECACHELIMIT" %s 0 -attr "NUKEEXEVERSION" %s 0' % (int(cache), nukeExe)
    cmd += ' -attr "environmental_variables" "%s" 0' % (';'.join('{}={}'.format(key, val.replace('/','\\').replace('\\','\\\\')) for key, val in sorted(envDict.items())))
    cmd += ' -f "%s"' % currNk


    def appendWriteNodes(_nodes):
        writeNodes = list()
        if _nodes is not None:
            for n in _nodes:
                if n.Class() == "Write":
                    writeNodes.append(n)

                elif n.Class() == "Group":
                    writeNodes = writeNodes + appendWriteNodes(n.nodes())
        return writeNodes

    if selectedNodes is None:
        selectedNodes = nuke.selectedNodes()

    writeNodes = appendWriteNodes(selectedNodes)
    # sort
    selWrites = '__'+'_'.join([node.fullName() for node in sorted(writeNodes)])

    if writeNodes:
        if parallel == False:
            cmd += ' -add "-X '+','.join([node.fullName() for node in sorted(writeNodes)])
            rep = jobname+selWrites
            cmd = cmd.replace("replaceMe",str(rep))
            # submit
            print ("musterSubmit: Info. Submitting...")
            print (cmd)
            print (os.system(cmd))
        else:
            for no in sorted(writeNodes):
                cmd += ' -add "-X ' + no.fullName() + '"'
                # change job name
                cmd = cmd.replace("replaceMe", jobname + "__" + no.fullName())
                print ("musterSubmit: Info. Submitting...")
                print (cmd)

                # submit
                print (os.system(cmd))
                # change name back for next loop
                cmd = cmd.replace(jobname + "__" + no.fullName(),"replaceMe")