"""
Adds a menu item to the Render menu in Nuke, allowing the user to submit a job to OpenCue providing 
it has been set up according to the Distributed CopyCat Training document. The nukePath, nukeScript, 
cuePyPath, cueModuleDirs and CUEBOT_HOSTS environment variable in this script should be updated 
accordingly.
"""
import os
import nuke
import nukescripts
from sys import stderr, path
import ipaddress

# set up environment variable for cuebot server
os.environ["CUEBOT_HOSTS"] = #insert_ip_address_here

# path to the Nuke executable on worker machine
nukePath = "/opt/Nuke14.1v1/Nuke14.1"

# nuke script to use
nukeScript = "/shared/copycat_training.nk"

# cuePyPath below should point to the directory where you installed all the 
# OpenCue python modules. This may be in a virtual environment directory.
cuePyPath = "/opencue/venv/lib64/python3.9/site-packages"
cueModuleDirs = [
  cuePyPath, 
  os.path.join(cuePyPath, "pycue-0.21.13-py3.9.egg"),
  os.path.join(cuePyPath, "pyoutline-0.21.13-py3.9.egg")
]
path.extend(cueModuleDirs)

import opencue
import outline
import outline.cuerun
import outline.modules.shell


def tagFromLayerNum(layerNum):
  return "Worker" + str(layerNum + 1)


def getWorkerIpFromTag(tagToFind, workers):
  for worker in workers:
    if tagToFind in worker.tags():
      workerName = worker.name()
      break

  # test that it is an ip address
  try:
    ipaddress.ip_address(workerName)
  except ValueError:
    print("Warning: worker name {} not an IP!".format(workerName), file=stderr)
  return workerName


def buildCopyCatCmd(openCueDialog, mainCopycatWorkerIp, layerNum, workers):
  cmd = "COPYCAT_MAIN_ADDR={} ".format(mainCopycatWorkerIp)
  cmd += "COPYCAT_MAIN_PORT=30000 "
  cmd += "COPYCAT_RANK={} ".format(layerNum)
  cmd += "COPYCAT_WORLD_SIZE={} ".format(openCueDialog.numWorkers())
  if layerNum > 0:
    # find the ip address of the worker machine with the tag for this layer
    tag = tagFromLayerNum(layerNum)
    workerIp = getWorkerIpFromTag(tag, workers)
    cmd += "COPYCAT_LOCAL_ADDR={} ".format(workerIp)

  cmd += "{nukeExe} -F #IFRAME# -X {copyCatNode} --gpu ".format(nukeExe=nukePath, copyCatNode=openCueDialog.getNodeName())
  cmd += nukeScript
  return cmd


def buildLayer(openCueDialog, mainCopycatWorkerIp, layerNum, workers):
  layer = outline.modules.shell.Shell(
    "copycat-training-layer-{}".format(str(layerNum)), 
    command=buildCopyCatCmd(openCueDialog, mainCopycatWorkerIp, layerNum, workers).split(), 
    chunk='1',
    threads=float(0),
    range=str(openCueDialog.frame()),
    threadable=False
  )
  # in this example, Services have the same name as the tags
  service = tagFromLayerNum(layerNum)
  layer.set_service(service)
  return layer


def submitJob(openCueDialog):
  """ Creates an OpenCue job and submits the job using PyOutline, with information 
  from the dialog.
  """
  if not nuke.toNode(openCueDialog.getNodeName()):
    print("Warning: Can't find node {} in your Nuke session".format(openCueDialog.getNodeName()), file=stderr)

  workers = opencue.api.getHosts()
  if not 1 <= openCueDialog.numWorkers() <= len(workers):
    print("Error: Number of workers not in range {}-{}".format(1, len(workers)), file=stderr)
    return None

  ol = outline.Outline(
    "copycat-training-job",
    shot="copycat-training-shot",
    show="testing",
    user="purple"
  )
  mainCopycatWorkerIp = getWorkerIpFromTag(tagFromLayerNum(0), workers)
  for layerNum in range(openCueDialog.numWorkers()):
    layer = buildLayer(openCueDialog, mainCopycatWorkerIp, layerNum, workers)
    ol.add_layer(layer)

  ol.set_facility("local")
  return outline.cuerun.launch(ol, use_pycuerun=False)


class ModalOpenCueDialog( nukescripts.PythonPanel ):
  """ Simple dialog for submitting a CopyCat training job to OpenCue.
  """
  def __init__( self ):
    nukescripts.PythonPanel.__init__( self, "CopyCat training", "uk.co.thefoundry.OpenCueDialog" )

    maxMachines = len(opencue.api.getHosts())
    self._numWorkers = nuke.Int_Knob("numWorkers", "Workers to use (max {}):".format(maxMachines))
    self._numWorkers.setValue(maxMachines)
    self._copycatNodeName = nuke.String_Knob("copycatNodeName", "CopyCat node to train:")
    if nuke.selectedNodes():
      self._copycatNodeName.setValue(nuke.selectedNodes()[0].name())

    self._frame = nuke.Int_Knob("frame", "Frame:")
    self._frame.setValue(1)

    self.addKnob( self._numWorkers )
    self.addKnob( self._copycatNodeName )
    self.addKnob( self._frame )

  def numWorkers( self ):
    return self._numWorkers.value()

  def getNodeName( self ):
    return self._copycatNodeName.value()

  def frame( self ):
    return self._frame.value()

  def showModalDialog( self ):
    result = nukescripts.PythonPanel.showModalDialog( self )
    if result:
      jobs = submitJob( self )
      if jobs:
        jobNames = ", ".join([job.name() for job in jobs])
        nuke.message("job " + jobNames + " submitted!")


def openModalDialog():
  return ModalOpenCueDialog().showModalDialog()


mainMenu = nuke.menu("Nuke")
renderMenu = mainMenu.menu("Render")
renderMenu.addCommand("Train CopyCat on render farm", "openModalDialog()")