#
#	Asset Manager for Nuke v1.0.0
#	Jorge Sanchez Salcedo (2018)
#
#	www.jorgesanchez-da.com
#	jorgesanchez.da@gmail.com
#

import os
import re
from glob import glob
import nuke
import nukescripts

### DEFINE ROOT ###
def rootDir():
    return 'D:\Projects'


print '   '
print ' > Root: ' + rootDir()
print '   '
print ' > Loading Asset Manager for Nuke v1.0.0'
print ' > by Jorge Sanchez Salcedo (2018)'
print ' > www.jorgesanchez-da.com'
print ' > jorgesanchez.da@gmail.com'
print '   '


### PROJECTS ###
def getDirs():
    projectsDir = rootDir()
    projects = []
    sequences = {}
    shots = {}

    for i in os.listdir(projectsDir):
        path = os.path.join(projectsDir, i)
        if os.path.isdir(path):
            projects.append(i)

    for t in projects:
        sequences[t] = []
        sequencesDir = os.path.join(projectsDir, t, 'seqs')

        for d in os.listdir(sequencesDir):
            path = os.path.join(sequencesDir, d)

            if os.path.isdir(path):
                sequences[t].append(d)
                shots[d] = []
                shotsDir = os.path.join(sequencesDir, d, 'shots')

                for i in os.listdir(shotsDir):
                    path = os.path.join(shotsDir, i)
                    if os.path.isdir(path):
                        shots[d].append(i)
    return projects, sequences, shots


def createProjectTab():
    projects, sequences, shots = getDirs()
    departament = ['comp', 'lighting', 'lookdev']
    n = nuke.thisNode()
    js_tab = nuke.Tab_Knob('JS','JS')
    js_prj = nuke.Enumeration_Knob('prj', 'Project', projects)
    js_seq = nuke.Enumeration_Knob('seq', 'Sequence', [])
    js_sht = nuke.Enumeration_Knob('sht', 'Shot', [])
    js_dpt = nuke.Enumeration_Knob('dpt', 'Departament', departament)

    for k in (js_tab, js_prj, js_seq, js_sht, js_dpt):
        n.addKnob(k)

    updateProjectTab()


def updateProjectTab():
    projects, sequences, shots = getDirs()
    n = nuke.thisNode()
    k = nuke.thisKnob()

    if not k or k.name() in ['prj', 'seq', 'showPanel']:
        n['seq'].setValues(sequences[n['prj'].value()])
        n['sht'].setValues(shots[n['seq'].value()])
        n['seq'].setValue(0)
        n['sht'].setValue(0)
        os.environ['PRJ'] = n['prj'].value()
        os.environ['SEQ'] = n['seq'].value()
        os.environ['SHT'] = n['sht'].value()
        os.environ['DPT'] = n['dpt'].value()


### DEFINE SHOT ###
def getNkDir():
    nkDir = os.path.join(rootDir(), os.getenv('PRJ'), 'seqs', os.getenv('SEQ'), 'shots', os.getenv('SHT'), os.getenv('DPT'))
    if not os.path.isdir(nkDir):
        raise ValueError, 'NUKE directory does not exist'
    return nkDir


### GET NK FROM CURRENT SHOT ###
def getNkFiles():
    nkFiles = glob(os.path.join(getNkDir(), '*.nknc'))
    return nkFiles


### UI ###
class NkPanel(nukescripts.PythonPanel):
    def __init__(self, nkScripts):
        nukescripts.PythonPanel.__init__(self, 'Open Nuke Script')
        self.checkboxes = []
        self.nkScripts = nkScripts
        self.selectedScript = ''

        ### INDEX ###
        for i, n in enumerate(self.nkScripts):
            nkDivider = nuke.Text_Knob('divider', '')
            nk = nuke.Boolean_Knob('nknc_%s' % i, os.path.basename(n))
            self.addKnob(nk)
            self.addKnob(nkDivider)
            nk.setFlag(nuke.STARTLINE)
            self.checkboxes.append(nk)

    def knobChanged(self, knob):
        if knob in self.checkboxes:
            for cb in self.checkboxes:
                if knob == cb:
                    index = int(knob.name().split('_')[-1])
                    self.selectedScript = self.nkScripts[index]
                    continue
                cb.setValue(False)


### EASY SAVE ###
def easySave():
    nkDir = getNkDir()
    description = nuke.getInput('script description', 'bashComp').replace(' ', '')
    fileSaved = False
    version = 1

    ### NAME CONVENTION ###
    while not fileSaved:
        nkName = '%s_%s_%s_%s_v%02d.nknc' % (os.getenv('PRJ'), os.getenv('SEQ'), os.getenv('SHT'), description, version)

        ### JOIN DIRECTORY ###
        nkPath = os.path.join(nkDir, nkName)

        ### VERSION ###
        if os.path.isfile(nkPath):
            version += 1
            continue

        ### SAVE NK ###
        nuke.scriptSaveAs(nkPath)
        fileSaved = True
    return


### CHECK NK VERSION ###
def checkScriptName():
    if not re.search(r'[vV]\d+', nuke.root().name()):
        raise NameError, 'Please, include a version number and save script again.'


def createWriteDir():
    wrDir = os.path.dirname(nuke.filename(nuke.thisNode()))
    if not os.path.isdir(wrDir):
        os.makedirs(wrDir)


def getVersions():
    dpt = ['comp', 'lighting', 'lookdev']
    dptDict = {}


    shotDir = os.path.join(rootDir(), os.getenv('PRJ'), 'seqs', os.getenv('SEQ'),'shots',os.getenv('SHT'))
    for t in dpt:
        dptDict[t] = []
        dptDir = os.path.join(shotDir, t, 'render')
        for d in os.listdir(dptDir):
            path = os.path.join(dptDir, d)
            if os.path.isdir(path):
                dptDict[t].append(d)
    return dptDict


def createVersionKnob():
    node = nuke.thisNode()
    tabKnob = nuke.Tab_Knob('JS', 'JS')
    typeKnob = nuke.Enumeration_Knob('versionType', 'type', ['comp', 'lighting', 'lookdev'])
    versionKnob = nuke.Enumeration_Knob('_version', 'version',[])

    for k in (tabKnob, typeKnob, versionKnob):
        node.addKnob(k)
    updateVersionKnob()


def updateVersionKnob():
    node = nuke.thisNode()
    knob = nuke.thisKnob()

    if not knob or knob.name() in ['versionType', 'showPanel']:
        versionDict = getVersions()
        node['_version'].setValues(versionDict[node['versionType'].value()])
        node['_version'].setValue(0)