'''
** 
**  File library 
**
**  Name    : submitToMuster.py
**  Author  : Leonardo Bernardini
**  Version : 2.0, Thu Feb 20th 2014
**
**  Copyright 2000-2017, Virtual Vertex
**  All Rights Reserved.
**
**  This file contains UNPUBLISHED PROPRIETARY SOURCE CODE.
**  The contents of this file may not be disclosed to third parties, copied
**  or duplicated in any form, in whole or in part, without the prior written
**  permission of the author.
**
'''

import nuke
import nukescripts
import os
import sys
import ConfigParser
import PySide2.QtCore as QtCore
import PySide2.QtWidgets as QtWidgets

def getMusterPath(installationPath):
	if os.name != "posix":
		return installationPath
	else:
		if sys.platform == "darwin":
			return (os.path.join(installationPath,"Contents"))
	return installationPath
	  		
def getMrtoolPath(installationPath):
	if os.name != "posix":
		return os.path.join(getMusterPath(installationPath),"mrtool.exe")
	return os.path.join(getMusterPath(installationPath),"mrtool")
	
def checkMusterPath(path):
	mrtool = getMrtoolPath(path)
	if not os.path.isfile(mrtool):
		return None
	return mrtool


def submitToMuster():
    
        rootNode = nuke.toNode("root")
        if not rootNode:
            nuke.message("Please save your job before submitting to Muster !")
            return None
        curScript = rootNode.name()
        job = os.path.splitext(os.path.basename(curScript))[0]
        
        homedir = os.path.expanduser("~")
        settingsDir = os.path.join(homedir,".Muster")
        if not os.path.isdir(settingsDir):
                os.mkdir(settingsDir)

        settingsFile = os.path.join(settingsDir,"nuke_connector.ini")
        config = ConfigParser.RawConfigParser()
        config.read(settingsFile)
        try:
                musterLocation = config.get('Muster', 'installationPath')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
                musterLocation = os.getenv("MUSTER")
                if not musterLocation:
                        nuke.message("MUSTER environmental variable not defined. Please select the location of Muster mrtool executable!")
        try:
                dispatcherAddress = config.get('Muster', 'address')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
                dispatcherAddress = "127.0.0.1"

        try:
            musterVersion = int(config.get('Muster', 'version'))
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            musterVersion = 0

        try:
            dispatcherPort = config.get('Muster', 'port')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            dispatcherPort = "9781"

        try:
                login = config.get('Muster', 'login')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
                login = "admin"
        try:
                packetSize = config.get('Muster', 'packetSize')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
                packetSize = "10"

        try:
            skipTemplates = config.get('Muster', 'skipTemplates')
            if skipTemplates == "1": skipTemplates = QtCore.Qt.Checked
            else: skipTemplates = QtCore.Qt.Unchecked

        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            skipTemplates = QtCore.Qt.Checked
        
        try:
            nukeTemplateId = config.get('Muster', 'nukeTemplateId')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            nukeTemplateId = "40"

        try:
            nukeTemplateVersion = config.get('Muster', 'nukeTemplateVersion')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            nukeTemplateVersion = ""

        try:
            includedPools = config.get('Muster', 'includedPools')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            includedPools = ""

        try:
            excludedPools = config.get('Muster', 'excludedPools')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            excludedPools = ""
        
        try:
            errorCheck = config.get('Muster', 'errorCheck')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            errorCheck = 0
          
        try:
            exitErrorCheck = config.get('Muster', 'exitErrorCheck')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            exitErrorCheck = 0
              
        try:
            additionalFlags = config.get('Muster', 'additionalFlags')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            additionalFlags = ""

        try:
            jobDepartment = config.get('Muster', 'jobDepartment')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            jobDepartment = ""

        try:
            jobProject = config.get('Muster', 'jobProject')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            jobProject = ""

        try:
            jobParent = config.get('Muster', 'jobParent')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            jobParent = ""

        try:
            jobMaximumInstances = config.get('Muster', 'jobMaximumInstances')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            jobMaximumInstances = "0"

        try:
            jobDependancies = config.get('Muster', 'jobDependancies')
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            jobDependancies = ""

        try:
            jobSubmitPaused = config.get('Muster', 'jobSubmitPaused')
            if jobSubmitPaused == "1": jobSubmitPaused = QtCore.Qt.Checked
            else: jobSubmitPaused = QtCore.Qt.Unchecked
        except (ConfigParser.NoSectionError,ConfigParser.NoOptionError):
            jobSubmitPaused = QtCore.Qt.Unchecked

        dlg = SubmitToMuster()

        try:
                nuke.scriptSave()
        except:
                return None

        dlg.jobName.setText(job)
        dlg.musterVersion.setCurrentIndex(musterVersion)
        dlg.dispatcherPort.setText(dispatcherPort)
        dlg.startFrame.setText("%s" % nuke.root().firstFrame())
        dlg.endFrame.setText("%s" % nuke.root().lastFrame())
        dlg.byFrame.setText("1")
        dlg.packetSize.setText("10")
        dlg.jobPriority.setText("1");
        dlg.installationPath.setText(musterLocation)
        dlg.dispatcherIpAddress.setText(dispatcherAddress)
        dlg.username.setText(login)
        dlg.packetSize.setText(packetSize)
        dlg.nukeTemplateID.setText(nukeTemplateId)
        dlg.nukeTemplateVersion.setText(nukeTemplateVersion)
        dlg.skipTemplates.setCheckState(skipTemplates)
        dlg.musterVersion.currentIndexChanged.connect(dlg.onMusterVersionChanged)

        dlg.includedPools.setText(includedPools)
        dlg.excludedPools.setText(excludedPools)
        dlg.additionalFlags.setText(additionalFlags)
        dlg.jobDepartment.setText(jobDepartment)
        dlg.jobProject.setText(jobProject)
        dlg.jobParent.setText(jobParent)
        dlg.errorCheck.setCurrentIndex(int(errorCheck))
        dlg.exitErrorCheck.setCurrentIndex(int(exitErrorCheck))
        dlg.jobMaximumInstances.setText(jobMaximumInstances)
        dlg.jobDependancies.setText(jobDependancies)
        dlg.jobSubmitPaused.setCheckState(jobSubmitPaused)

        res = dlg.exec_()
        if res != QtWidgets.QDialog.Accepted:
            return None

        try:
            config.add_section('Muster')
        except ConfigParser.DuplicateSectionError:
            pass

        config.set('Muster', 'installationPath', dlg.installationPath.text())
        config.set('Muster', 'address', dlg.dispatcherIpAddress.text())
        config.set('Muster', 'login', dlg.username.text())
        config.set('Muster', 'packetSize', dlg.packetSize.text())
        config.set('Muster', 'port', dlg.dispatcherPort.text())
        config.set('Muster', 'version', dlg.musterVersion.currentIndex())
        config.set('Muster', 'nukeTemplateId', dlg.nukeTemplateID.text())
        config.set('Muster', 'nukeTemplateVersion', dlg.nukeTemplateVersion.text())

        skipTemplates = "0"
        if dlg.skipTemplates.checkState() == QtCore.Qt.Checked: skipTemplates = "1"
        config.set('Muster', 'skipTemplates',skipTemplates)
        config.set('Muster', 'includedPools', dlg.includedPools.text())
        config.set('Muster', 'excludedPools', dlg.excludedPools.text())
        config.set('Muster', 'additionalFlags', dlg.additionalFlags.text())
        config.set('Muster', 'errorCheck', dlg.errorCheck.currentIndex())
        config.set('Muster', 'exitErrorCheck', dlg.exitErrorCheck.currentIndex())
        config.set('Muster', 'jobDepartment', dlg.jobDepartment.text())
        config.set('Muster', 'jobProject', dlg.jobProject.text())
        config.set('Muster', 'jobParent', dlg.jobParent.text())
        config.set('Muster', 'jobMaximumInstances', dlg.jobMaximumInstances.text())
        config.set('Muster', 'jobDependancies', dlg.jobDependancies.text())
        jobSubmitPaused = "0"
        if dlg.jobSubmitPaused.checkState() == QtCore.Qt.Checked: jobSubmitPaused = "1"
        config.set('Muster', 'jobSubmitPaused',jobSubmitPaused)

        with open(settingsFile, 'wb') as configfile:
                config.write(configfile)

class mrtoolMonitor(QtWidgets.QDialog):
    aborted = QtCore.Signal()
    
    def __init__(self, parent=None):
        QtWidgets.QDialog.__init__(self, parent)
        self.setWindowTitle("Mrtool output")
        self.setLayout(QtWidgets.QVBoxLayout())
        self.textView = QtWidgets.QTextEdit()
        self.abortButton = QtWidgets.QPushButton("Abort")
        self.layout().addWidget(self.textView)
        self.layout().addWidget(self.abortButton)
        self.abortButton.pressed.connect(self.onAbort)
        
    def addString(self,text):
        self.textView.append(text)

    def onAbort(self):
        self.aborted.emit()

    def closeEvent(self, event):
        event.ignore()

class SubmitToMuster(QtWidgets.QDialog):
    def onMusterVersionChanged(self,index):
        if self.musterVersion.currentIndex() == 0:
            self.dispatcherPort.setText("9881")
        elif self.musterVersion.currentIndex() == 1:
            self.dispatcherPort.setText("9781")
        else:
            self.dispatcherPort.setText("9681")


    def onInstallationPathPressed(self):
        path = QtWidgets.QFileDialog.getExistingDirectory (self,"Select installation path", self.installationPath.text());
        if len(path):
            if checkMusterPath(path):
                self.installationPath.setText(path)
            else:
                nuke.message("The specified path does not contains Muster mrtool command line application !")
   
    def onProcessFinished(self,exitCode):
        out = self.process.readAllStandardOutput()
        if self.mrtoolMon != None:
            self.mrtoolMon.addString(out.data())

        out = self.process.readAllStandardError()
        if self.mrtoolMon != None:
            self.mrtoolMon.addString(out.data())
            self.mrtoolMon.abortButton.setText("Close")
        
        self.submitButton.setEnabled(1)
        self.closeButton.setEnabled(1)
    
    def onProcessErrDataReady(self):
        out = self.process.readAllStandardError()
        self.mrtoolMon.addString(out.data())

    def onProcessOutDataReady(self):
        out = self.process.readAllStandardOutput()
        self.mrtoolMon.addString(out.data())

    def onMrtoolAbort(self):
        self.process.terminate()
        self.mrtoolMon.hide()
        self.mrtoolMon = None
    
    def fireMrtool(self,jobFile):
        
        mrtoolPath = checkMusterPath(self.installationPath.text())
        if not mrtoolPath:
            nuke.message("The specified path does not contains Muster mrtool command line application !")
            return None
        
        musterEnv = getMusterPath (self.installationPath.text())
        arguments = []
        arguments.append("-v")
        arguments.append("1")
        if self.skipTemplates.checkState() == QtCore.Qt.Checked:
            arguments.append("-usecache")
        arguments.append("-b")
        arguments.append("-n")
        arguments.append(self.jobName.text())
        arguments.append("-e")
        arguments.append(self.nukeTemplateID.text())
        arguments.append("-u")
        arguments.append(self.username.text())
        arguments.append("-p")
        arguments.append(self.password.text())
        arguments.append("-sf")
        arguments.append(self.startFrame.text())
        arguments.append("-ef")
        arguments.append(self.endFrame.text())
        arguments.append("-bf")
        arguments.append(self.byFrame.text())
        arguments.append("-pr")
        arguments.append(self.jobPriority.text())
        if len(self.nukeTemplateVersion.text()):
            arguments.append("-tv")
            arguments.append(self.nukeTemplateVersion.text())

        if len(self.includedPools.text()):
            arguments.append("-pool")
            arguments.append(self.includedPools.text())
        
        if len(self.excludedPools.text()):
            arguments.append("-excludepool")
            arguments.append(self.excludedPools.text())
        
        if len(self.jobDepartment.text()):
            arguments.append("-department")
            arguments.append(self.jobDepartment.text())
        
        if len(self.jobProject.text()):
            if self.musterVersion.currentIndex() == 0:
                arguments.append("-project")
            else:
                arguments.append("-group")
            arguments.append(self.jobProject.text())

        if len(self.jobParent.text()):
            arguments.append("-parent")
            arguments.append(self.jobParent.text())

        if len(self.jobMaximumInstances.text()):
            arguments.append("-max")
            arguments.append(self.jobMaximumInstances.text())

        if len(self.jobDependancies.text()):
            arguments.append("-wait")
            arguments.append(self.jobDependancies.text())

        if self.jobSubmitPaused.checkState() == QtCore.Qt.Checked:
            arguments.append("-ojs")
            arguments.append("2")

        if self.errorCheck.currentIndex() == 1:
            arguments.append("-logerrtype")
            arguments.append(str(self.errorCheck.currentIndex()))
        if self.errorCheck.currentIndex() == 2:
            arguments.append("-logerrtype")
            arguments.append(str(self.errorCheck.currentIndex()))
        if self.errorCheck.currentIndex() == 3:
            arguments.append("-logerrtype")
            arguments.append(str(self.errorCheck.currentIndex()))
            
        if self.exitErrorCheck.currentIndex() == 1:
            arguments.append("-ecerrtype")
            arguments.append(str(self.exitErrorCheck.currentIndex()))
        if self.exitErrorCheck.currentIndex() == 2:
            arguments.append("-ecerrtype")
            arguments.append(str(self.exitErrorCheck.currentIndex()))
        if self.exitErrorCheck.currentIndex() == 3:
            arguments.append("-ecerrtype")
            arguments.append(str(self.exitErrorCheck.currentIndex()))
            
        if len(self.additionalFlags.text()):
            arguments.append("-add")
            arguments.append(self.additionalFlags.text())
                
        arguments.append("-s")
        arguments.append(self.dispatcherIpAddress.text())
        arguments.append("-pk")
        arguments.append(self.packetSize.text())
        arguments.append("-f")
        arguments.append(jobFile)
        
        self.mrtoolMon = mrtoolMonitor()
        self.mrtoolMon.aborted.connect(self.onMrtoolAbort)
        self.mrtoolMon.setModal(1)
        self.mrtoolMon.show()
        self.mrtoolMon.addString("Starting " + mrtoolPath + " with the following arguments: " + '%s' % ' '.join(map(str, arguments)) + "\n")
        
        self.process = QtCore.QProcess()
        self.process.setProcessChannelMode(QtCore.QProcess.MergedChannels)
        processEnvironment = QtCore.QProcessEnvironment.systemEnvironment()
        processEnvironment.insert("MUSTER",musterEnv)
        self.process.setProcessEnvironment(processEnvironment)

        self.process.finished.connect(self.onProcessFinished)
        self.process.readyReadStandardError.connect(self.onProcessErrDataReady)
        self.process.readyReadStandardOutput.connect(self.onProcessOutDataReady)

        self.submitButton.setEnabled(0)
        self.closeButton.setEnabled(0)
        self.process.start(mrtoolPath , arguments)

        if not self.process.waitForStarted() :
            nuke.message("Failed to start Muster Mrtool process!")
            return None

    def onSubmit(self):
        if not checkMusterPath(self.installationPath.text()):
            nuke.message("The specified path does not contains Muster mrtool command line application !")
            return None

        rootNode = nuke.toNode("root")
        if not rootNode:
            nuke.message("Please save your job before submitting to Muster !")
            return None

        curScript = rootNode.name()
        job = os.path.splitext(os.path.basename(curScript))[0]

        self.fireMrtool(curScript)

    def onClose(self):
        self.accept()

    def onReset(self):
        rootNode = nuke.toNode("root")
        if rootNode:
            curScript = rootNode.name()
            job = os.path.splitext(os.path.basename(curScript))[0]
            self.jobName.setText(job)
            self.startFrame.setText("%s" % rootNode.firstFrame())
            self.endFrame.setText("%s" % rootNode.lastFrame())
        else:
            self.jobName.setText("Untitled")
            self.startFrame.setText("1")
            self.endFrame.setText("1")
        
        self.byFrame.setText("1")
        self.packetSize.setText("10")
        self.jobPriority.setText("1");
        self.nukeTemplateID.setText("40")
        self.nukeTemplateVersion.setText("")
        self.includedPools.setText("")
        self.excludedPools.setText("")
        self.additionalFlags.setText("")
        self.jobDepartment.setText("")
        self.jobProject.setText("")
        self.jobParent.setText("")
        self.jobMaximumInstances.setText("0")
        self.jobDependancies.setText("")
        self.jobSubmitPaused.setCheckState(QtCore.Qt.Unchecked)
        self.exitErrorCheck.setCurrentIndex(0)
        self.errorCheck.setCurrentIndex(0)
          
    def __init__(self, parent=None):
        QtWidgets.QDialog.__init__(self, parent)
        
        self.setWindowTitle("Submit to Muster")
        self.setLayout(QtWidgets.QVBoxLayout())
        self.tabWidget = QtWidgets.QTabWidget()
        self.jobTabWidget = QtWidgets.QWidget()
        self.dispatcherTabWidget = QtWidgets.QWidget()
        self.tabWidget.addTab(self.jobTabWidget,"Job")
        self.tabWidget.addTab(self.dispatcherTabWidget,"Dispatcher")
        self.layout().addWidget(self.tabWidget)
        
        self.dispatcherFormLayout = QtWidgets.QFormLayout()
        self.jobFormLayout = QtWidgets.QFormLayout()
        self.jobTabWidget.setLayout(self.jobFormLayout)
        self.dispatcherTabWidget.setLayout(self.dispatcherFormLayout)

        self.installationPath = QtWidgets.QLineEdit()
        self.installationPathLabel = QtWidgets.QLabel()
        self.installationPathLabel.setText("")
        self.pickInstallationPath = QtWidgets.QPushButton()
        self.pickInstallationPath.setText("...")
        self.installationLayout = QtWidgets.QHBoxLayout()
        self.installationLayout.addWidget(self.installationPath)
        self.installationLayout.addWidget(self.pickInstallationPath)
        
        self.musterVersion = QtWidgets.QComboBox()
        self.musterVersion.addItem("9.0 or greater")
        self.musterVersion.addItem("8.X")
        self.musterVersion.addItem("7.X")

        self.skipTemplates = QtWidgets.QCheckBox()
        self.dispatcherIpAddress = QtWidgets.QLineEdit()
        self.dispatcherPort = QtWidgets.QLineEdit()
        self.dispatcherPort.setInputMask("000000")
        self.username = QtWidgets.QLineEdit()

        self.jobName = QtWidgets.QLineEdit()
        self.password = QtWidgets.QLineEdit()
        self.password.setEchoMode(QtWidgets.QLineEdit.Password)
        self.startFrame = QtWidgets.QLineEdit()
        self.startFrame.setInputMask("0000000")
        self.endFrame = QtWidgets.QLineEdit()
        self.endFrame.setInputMask("0000000")
        self.byFrame = QtWidgets.QLineEdit()
        self.byFrame.setInputMask("0000000")
        self.packetSize = QtWidgets.QLineEdit()
        self.packetSize.setInputMask("0000")
        self.jobPriority = QtWidgets.QLineEdit()
        self.jobPriority.setInputMask("000");
        self.includedPools = QtWidgets.QLineEdit()
        self.excludedPools = QtWidgets.QLineEdit()
        self.errorCheck = QtWidgets.QComboBox()
        self.errorCheck.addItem("Default");
        self.errorCheck.addItem("Skip warnings");
        self.errorCheck.addItem("Skip errors");
        self.errorCheck.addItem("Skip both");
        self.exitErrorCheck = QtWidgets.QComboBox()
        self.exitErrorCheck.addItem("Default");
        self.exitErrorCheck.addItem("Skip warnings");
        self.exitErrorCheck.addItem("Skip errors");
        self.exitErrorCheck.addItem("Skip both");
        
        self.additionalFlags = QtWidgets.QLineEdit()
        self.nukeTemplateVersion = QtWidgets.QLineEdit()
        self.nukeTemplateID = QtWidgets.QLineEdit()
        self.nukeTemplateID.setInputMask("0000");
        self.jobDepartment = QtWidgets.QLineEdit()
        self.jobProject = QtWidgets.QLineEdit()
        self.jobParent = QtWidgets.QLineEdit()
        self.jobParent.setInputMask("00000000");
        self.jobMaximumInstances = QtWidgets.QLineEdit()
        self.jobMaximumInstances.setInputMask("00000000");
        self.jobDependancies = QtWidgets.QLineEdit()
        self.jobSubmitPaused = QtWidgets.QCheckBox()
        
        self.dispatcherFormLayout.addRow("Installation path:",self.installationLayout)
        self.dispatcherFormLayout.addRow("Muster version:",self.musterVersion)
        self.dispatcherFormLayout.addRow("Dispatcher address:",self.dispatcherIpAddress)
        self.dispatcherFormLayout.addRow("Dispatcher port:",self.dispatcherPort)
        self.dispatcherFormLayout.addRow("Login:",self.username)
        self.dispatcherFormLayout.addRow("Password:",self.password)
        self.dispatcherFormLayout.addRow("Skip templates sync:" ,self.skipTemplates)
        self.dispatcherFormLayout.addRow("Nuke template ID:",self.nukeTemplateID)

        self.jobFormLayout.addRow("Job name:",self.jobName)
        self.jobFormLayout.addRow("Job department:",self.jobDepartment)
        self.jobFormLayout.addRow("Job project:",self.jobProject)
        self.jobFormLayout.addRow("Start frame:",self.startFrame)
        self.jobFormLayout.addRow("End frame:",self.endFrame)
        self.jobFormLayout.addRow("By frame:",self.byFrame)
        self.jobFormLayout.addRow("Packet size:",self.packetSize)
        self.jobFormLayout.addRow("Job priority:",self.jobPriority)
        self.jobFormLayout.addRow("Job parent ID:",self.jobParent)
        self.jobFormLayout.addRow("Logs error check:",self.errorCheck)
        self.jobFormLayout.addRow("Exit codes error check:",self.exitErrorCheck)
        self.jobFormLayout.addRow("Job maximum instances:",self.jobMaximumInstances)
        self.jobFormLayout.addRow("Job dependancies:",self.jobDependancies)
        self.jobFormLayout.addRow("Destination pools:",self.includedPools)
        self.jobFormLayout.addRow("Excluded pools:",self.excludedPools)
        self.jobFormLayout.addRow("Submit paused:",self.jobSubmitPaused)
        self.jobFormLayout.addRow("Nuke template version:",self.nukeTemplateVersion)
        self.jobFormLayout.addRow("Additional flags:",self.additionalFlags)
        
        self.buttonsLayout = QtWidgets.QHBoxLayout()
        self.submitButton = QtWidgets.QPushButton("Submit")
        self.resetButton = QtWidgets.QPushButton("Reset")
        self.closeButton = QtWidgets.QPushButton("Close")
        self.buttonsLayout.addWidget(self.submitButton)
        self.buttonsLayout.addStretch()
        self.buttonsLayout.addWidget(self.resetButton)
        self.buttonsLayout.addWidget(self.closeButton)
        self.closeButton.pressed.connect(self.onClose)
        self.resetButton.pressed.connect(self.onReset)
        self.submitButton.pressed.connect(self.onSubmit)
        
        self.layout().addLayout(self.buttonsLayout)
        
        self.pickInstallationPath.pressed.connect(self.onInstallationPathPressed)
        self.setSizePolicy(QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Fixed))

