# -*- coding: utf8 -*-
'''
Keller Pipeline Tools
written by stefan ihringer (stefan@bildfehler.de)

version: 2012-06-01

Beschreibung:
diverse Hilfsfunktionen, z.B. um ein Script gemäß der Namenskonvention zu speichern
oder den Pfad einer Write-Node einzustellen.

Todo: Support für andere Departments, nicht (mehr) verfügbare Projekte
      create_autowrite(), update_autowrite()
      open script from read/write
      Filename-Preview im UpdateWritePanel
Bugs: Buchstaben in Shotnummer werden zwar erlaubt aber machen probleme bei update_writenode()
'''

import nuke, nukescripts
import re, os, glob

## dictionary mit gültigen Projektnamen, wie sie in Dateinamen vorkommen
project_list = {"VS": "Vampirschwestern",
                "VT": "Vatertage"}

## globale Settings:
default_format = "dpx"                          ## default für AutoWrite nodes
project_folder = r"\\calculon\o\_projekte"      ## kein Backslash am Ende der Pfade
scripts_folder = r"300_%s_Work2D\Comp"          ## %s wird durch Projektnamen ersetzt
render_folder  = "500_Renders"                  ## für finale Renderings
precomp_folder = "450_PreComp"                  ## für Vorgerendertes aus den Comps
footage_folder = "220_ProcessedFootage"         ## für Footage nach Degrain/Retusche


class KellerTools(KellerNukePlugin):

    def configurePlugin(self):
        menubar = nuke.menu("Nuke")
        m = menubar.addMenu("&Render")
        m.addCommand("Save Script By Convention...", 'kellertools.save_by_convention()')
        pass
        #m = menubar.addMenu("&Workgroup")
        #m.addCommand("Update Render Directory...", 'kellertools.update_writenodes()')

class FileSavePanel(nukescripts.PythonPanel) :

    def __init__(self, scriptinfo):
        '''
        Panel, das nach Projekt, Shot, Status, Name und Version fragt, bevor
        das Script mit einem gültigen Namen gespeichert werden kann. Scriptinfo
        kann ein dictionary sein, dass default-Werte enthält falls das aktuelle
        Script bereits mit einem gültigen Namen gespeichert ist.
        '''
        global project_list
        
        nukescripts.PythonPanel.__init__ (self, "Script speichern...", "com.keller.FileSavePanel", False)
        if scriptinfo is None:
            scriptinfo = {}

        last_artist = ""    ## todo: letzten Artist aus irgendwelchen Preferences laden?
        last_project = project_list.values()[0]

        self.projid = scriptinfo.get('projid', "")
        self.project = nuke.Enumeration_Knob("project", "Project", project_list.values())
        self.project.setValue(scriptinfo.get("project", last_project))
        ## todo: projekt richtig selektieren, unbekannte Projekte erlauben

        self.department = nuke.String_Knob("department", "Department", scriptinfo.get('department', "Comp"))
        ## todo: department als dropdown zeigen
        self.sequence = nuke.String_Knob("sequence", "Sequence/Shot", scriptinfo.get('sequence', ""))
        self.shot = nuke.String_Knob("shot", "/", scriptinfo.get('shot', ""))
        self.shot.clearFlag(nuke.STARTLINE)
        self.identifier = nuke.String_Knob("identifier", "Identifier", scriptinfo.get('identifier', "Comp"))
        self.version = nuke.String_Knob("version", "Version", scriptinfo.get('version', "01")+scriptinfo.get('subversion',""))
        self.specifier = nuke.String_Knob("opt_specifier", "Specifier (opt.)", scriptinfo.get('opt_specifier', ""))
        self.artist = nuke.String_Knob("artist", "Artist", scriptinfo.get('artist', last_artist))
        self.info = nuke.String_Knob("opt_info", "Info (opt.)", scriptinfo.get('opt_info', ""))
        self.preview = nuke.Text_Knob("preview", "Fileame: ", "...")
        for k in (self.project, self.sequence, self.shot, self.identifier, self.specifier, self.version, self.artist, self.info, self.preview):
            self.addKnob(k)

    def knobChanged (self, knob) :
        "Vorschau des Dateinamens updaten, wenn ein Feld geändert wurde."

        if knob.name() in ("project", "showPanel"):
            for key, value in project_list.items():
                if value == self.project.value():
                    self.projid = key
                    break
        if knob.name() in ("department", "sequence", "shot", "identifier", "opt_specifier", "artist", "opt_info"):
            ## Sonderzeichen löschen
            cleanstr = re.sub("(?i)[\W_]", "", knob.value())
            knob.setValue(cleanstr)
        if knob.name() == "shot" and knob.value().upper() == "GEN":
            ## "GEN" immer in Großbuchstaben
            knob.setValue("GEN")
        if knob.name() == "version":
            ## version: nur Zahlen und ein optionaler Buchstabe
            cleanstr = re.sub("(?i)[^\da-z]", "", knob.value())
            try:
                cleanstr = re.search("(?i)\d+[a-z]?", cleanstr).group(0)
            except:
                cleanstr = ""
            knob.setValue(cleanstr)

        ## Version aus version & subversion zusammensetzen
        version_str = ""
        v_exp = re.match("(\d+)([a-z])?", self.version.value())
        try:
            version_str = "%02d" % int(v_exp.group(1))
            self.version.setValue(version_str)
            version_str = version_str + v_exp.group(2)
            self.version.setValue(version_str)
        except:
            pass

        tokens = self.artist.value()
        if self.specifier.value():
            tokens = self.specifier.value() + "." + tokens
        if self.info.value():
            tokens = tokens + "." + self.info.value()

        newname = "%s_%s_%s_%s.%s_v%s.%s.nk" % (self.projid, self.department.value(), self.sequence.value(), self.shot.value(), self.identifier.value(), version_str, tokens)
        self.preview.setValue(newname)
        return

    def showModalDialog( self ):
        result = nukescripts.PythonPanel.showModalDialog( self )
        return result

    def getInfo( self ):
        '''
        Gibt ein dictionary mit den Werten der Eingabefelder zurück.
        Leere Eingabefelder werden ignoriert, so dass save_by_convention() dann den Fehler bemerkt.
        '''
        result = {}
        for k in (self.project, self.department, self.sequence, self.shot, self.identifier, self.specifier, self.version, self.artist, self.info):
            if k.value():
                result[k.name()] = re.sub("(?i)[\W_]", "", k.value())
        ## Projekt-ID ermitteln
        global project_list
        for key,value in project_list.items():
            if value == result['project']:
                result['projid'] = key
                break
        ## Version mit Buchstaben splitten in version & subversion
        v_exp = re.match("(\d+)([a-z])?", result['version'])
        if len(v_exp.groups()) == 2 and v_exp.group(2) is not None:
            result['version'] = v_exp.group(1)
            result['subversion'] = v_exp.group(2)
        return result


##-------------------------------------------------------------------------------------------------------

class UpdateWritePanel(nukescripts.PythonPanel) :

    def __init__(self, writename = "", default_target = 0, default_format = "exr", default_space = "", default_specifier = ""):
        '''
        Panel, das zum Updaten einer Write-Node angezeigt wird und nach Dateiformat
        und Colorspace fragt. Es unterstützt Haupt-Renderings, PreComps und ProcessedFootage.
        '''
        ## TODO: Defaults als dictionary übergeben wir im FileSavePanel
        fileformatlist = {"exr":'linear', "dpx":'logC', "tiff":'sRGB', "tga":'sRGB', "jpeg":'sRGB'}
        colorspacelist = ['linear', 'logC', 'sRGB', 'rec709']
        self.rendertargets =  ['Final Render', 'PreComp', 'Processed Footage']
        if not default_format:
            default_format = "exr"
        if not default_space:
            ## kein colorspace angegeben: auf Basis des Dateiformats ermitteln oder als letzte Alternative "linear"
            default_space = fileformatlist.get(default_format, "linear")

        nukescripts.PythonPanel.__init__ (self, "Set Render Path for '%s'" % writename, "com.keller.UpdateWritePanel", False)
        self.target = nuke.Enumeration_Knob("target", "Type", self.rendertargets)
        self.target.setValue(self.rendertargets[default_target])
        tmplist = fileformatlist.keys()
        tmplist.sort()
        self.fileformat = nuke.Enumeration_Knob("fileformat", "File Format / Color Space", tmplist)
        self.fileformat.setValue(default_format)
        self.colorspace = nuke.Enumeration_Knob("colorspace", " ", colorspacelist)
        self.colorspace.clearFlag(nuke.STARTLINE)
        self.colorspace.setValue(default_space)
        self.specifier = nuke.String_Knob("specifier", "PreComp Specifier", "")
        if default_target == 0:
            self.specifier.setFlag(nuke.DISABLED)
            self.specifier.setValue("n/a")
        else:
            self.specifier.setValue(default_specifier)
        ## TODO: Dateinamensvorschau wie im FileSavePanel
        for k in (self.target, self.fileformat, self.colorspace, self.specifier):
            self.addKnob(k)

    def knobChanged (self, knob) :
        if knob.name() == "target":
            i = self.rendertargets.index(knob.value())
            if i == 0:
                self.specifier.setFlag(nuke.DISABLED)
                if self.specifier.value() == "":
                    self.specifier.setValue("n/a")
            elif i == 1:
                self.fileformat.setValue("exr")
                self.colorspace.setValue("linear")
                self.specifier.clearFlag(nuke.DISABLED)
                if self.specifier.value() == "n/a":
                    self.specifier.setValue("")
            else:
                self.specifier.clearFlag(nuke.DISABLED)
                if self.specifier.value() == "n/a":
                    self.specifier.setValue("")
        if knob.name() == "specifier":
            ## Sonderzeichen löschen
            cleanstr = re.sub("(?i)[\W_]", "", knob.value())
            knob.setValue(cleanstr)

    def showModalDialog( self ):
        "Show the panel as a modal dialog."

        result = nukescripts.PythonPanel.showModalDialog( self )
        return result

    def getInfo( self ):
        "Gibt ein dictionary mit den Werten der Eingabefelder zurück."

        result = {}
        result['target'] = self.rendertargets.index(self.target.value())
        result['fileformat'] = self.fileformat.value()
        result['colorspace'] = self.colorspace.value()
        if result['target'] > 0:
            result['specifier'] = re.sub("(?i)[\W_]", "", self.specifier.value())
        return result


##-------------------------------------------------------------------------------------------------------

def parse_scriptname(s_name = "", quick = False):
    '''
    Analysiert den Dateinamen des angegebenen oder aktuellen Scripts, um Projekt, Shot, etc... zu ermitteln.
    Ist der Name ungültig, wird None zurück gegeben, ansonsten ein dictionary mit Stringstücken.
    Der Parameter quick kann verwendet werden, um nur schnell ein Verzeichnis zu scannen ohne Projektnamen aufzulösen.
    projid und project existieren dann nicht im zurückgegebenen Dictionary.
    '''
    ## Namenskonvention:
    ## [PROJECT]_[DEPARTMENT]_[SEQ]_[SHOT].[IDENTIFIER]_[VERSION][SUBVERSION].[OPT_SPECIFIER].[ARTISTNAME].[OPT_INFO].[EXT]
    ## z.B. VS_Comp_030_010.Main_v02b.sih.FixedKeyEdges.nk
    
    if not s_name:
        s_name = nuke.root().knob("name").value()
    if not s_name:
        return None

    s_path = os.path.dirname(s_name)
    s_name = os.path.splitext(os.path.basename(s_name))[0]
    s_exp = re.match(r"""(?P<projid>.+?)_        # VS_
                         (?P<department>.+?)_    # Comp_
                         (?P<sequence>\d+)_      # 001_
                         (?P<shot>\d+|GEN)       # 010  (oder "GEN")
                         \.                      # .
                         (?P<identifier>.+?)_    # Comp_
                         v(?P<version>\d+)       # v01
                         (?P<subversion>[a-z]?)  # a
                         \.                      # .
                         (?P<tokens>[\w.]+)      # x.artist.y  """, s_name, re.IGNORECASE|re.VERBOSE)
    if not s_exp:
        return None
    result = s_exp.groupdict()
    ## GEN statt Shotnummer immer in Großbuchstaben
    if result['shot'].lower() == "gen":
        result['shot'] = "GEN"
    if result['subversion'] == None:
        result['subversion'] = ""
    ## 'tokens' enthält den string nach der Versionsnummer, der 1-3 Felder enthalten kann.
    result['artist'] = ""
    result['opt_specifier'] = ""
    result['opt_info'] = ""
    tokens = result['tokens'].split('.', 3)
    if len(tokens) == 1:
        ## Nur Artistname (kein Check auf dessen Gültigkeit!)
        result['artist'] = tokens[0]
    elif len(tokens) == 3:
        ## Artist + beide optionale Tokens
        result.update(dict(zip( ['opt_specifier','opt_artist','opt_info'], tokens )))
    else:
        ## 2 Tokens. Herausfinden, welcher der beiden der Artistname (2 oder 3 Buchstaben) ist
        if re.match("[a-zA-Z][a-zA-Z][a-zA-Z]?", tokens[0]):
            result['artist'] = tokens[0]
            result['opt_info'] = tokens[1]
        elif re.match("[a-zA-Z][a-zA-Z][a-zA-Z]?", tokens[1]):
            result['opt_specifier'] = tokens[0]
            result['artist'] = tokens[1]
        else:
            return None

    if not quick:
        if result['projid'] in project_list:
            result['project'] = project_list[result['projid']]
        else:
            ## Projektnamen aus Dateipfad ermitteln
            global scripts_folder
            pattern = re.compile(scripts_folder % "(\w+)", re.IGNORECASE)
            p_exp = pattern.search(s_path)
            if p_exp:
                result['project'] = p_exp.group(1)
                print("Projektname aus Pfad ermittelt: %s", p_exp.group(1))

    return result


def scan_shot_directory(dirname, identifier = ""):
    '''
    Durchsucht ein Verzeichnis und erstellt eine Liste an vorhandenen Comps.
    Wird z.B. verwendet, um nach der nächsthöheren Versionsnummer zu suchen.
    Rückgabe ist ein Dict von Versionsnummern->Aristnamen sowie den Feldern
    "maxversion" (höchste gefundene Versionsnummer) und "nextversion"
    (nächsthöhere Version) bzw. None bei Fehlern / leeren Ordnern.
    Wenn Identifier angegeben, werden nur diese Files gesucht.
    '''
    ## TODO: Buchstaben hinter Version ermitteln
    if not os.path.isdir(dirname):
        return None
    result = {}
    maxversion = 0
    wildcardsearch = dirname + "/*" + identifier + "_v*.nk"
    filelist = glob.glob(wildcardsearch)
    ## Todo: exception handling
    for f in filelist:
        info = parse_scriptname(f, quick = True)
        if info:
            v = int(info['version'])
            if v > maxversion:
                maxversion = v
            result[v] = info['artist']
    if result:
        result['maxversion'] = maxversion
        result['nextversion'] = maxversion + 1
        if identifier:
            result['identifier'] = identifier
    return result
    

def build_scriptname(s = {}):
    '''
    Erstellt einen Dateinamen auf Basis der angegebenen Stringstücke.
    projid und project muss beides vorhanden sein.
    '''
    try:
        global project_folder
        global scripts_folder
        shot_or_gen = s['shot']
        if shot_or_gen.upper() == "GEN":
            ## gesondertes _GEN-Unterverzeichnis
            shot_or_gen = "_GEN"
        newpath = scripts_folder % s['project']
        newpath = os.path.join( project_folder, s['project'], newpath, s['sequence'], shot_or_gen )

        tokens = s['artist']
        if s.get('opt_specifier'):
            tokens = s.get('opt_specifier') + "." + tokens
        if s.get('opt_info'):
            tokens = tokens + "." + s.get('opt_info')

        newname = "%s_%s_%s_%s.%s_v%02d%s.%s.nk" % (s['projid'], s['department'], s['sequence'], s['shot'], s['identifier'], int(s['version']), s.get('subversion',''), tokens)
        return os.path.join(os.path.normpath(newpath), newname)
    except:
        ## nicht alle benötigten Stringstücke gefunden
        return None


def build_renderpath(s = {}, w = {}):
    '''
    Erstellt einen zum Scriptnamen passenden Dateinamen für write-nodes.
    s ist ein dictionary mit scriptinfo (von parse_scriptname)
    w ist ein dictionary mit den keys target, fileformat, colorspace und evtl. specifier (vom UpdateWritePanel)
    '''
    ## TODO: für ProcessedFootage sollte der Name eher dem Footagenamen entsprechen statt dem Comp-Namen.
    ## TODO: "GEN" unterstützen
    try:
        global project_folder
        global render_folder
        global precomp_folder
        global footage_folder
        global default_format

        fileformat = w['fileformat']
        colorspace = w['colorspace']
        targetfolders = (render_folder, precomp_folder, footage_folder)
        
        if not fileformat:
            fileformat = default_format
        if colorspace == "linear":
            ## "linear" wird nicht in den Pfad geschrieben. Hier ggf. "lin" setzen, wenn gewünscht
            colorspace = ""
        elif colorspace == "AlexaV3LogC":
            ## "logC" ist kürzer für Pfadnamen
            colorspace = "logC"

        dirname = "%s_%s_%s_%s.%s_v%02d%s.%s" % (s['projid'], s['department'], s['sequence'], s['shot'], s['identifier'], int(s['version']), s.get('subversion',''), s['artist'])
        if w['target'] > 0:
            dirname = dirname + "." + w['specifier']
        if colorspace:
            dirname = dirname + "." + colorspace
        filename = dirname + ".%04d." + fileformat
        renderpath = os.path.join( project_folder, s['project'], targetfolders[w['target']], s['sequence'], s['shot'], dirname, filename )
        return os.path.normpath(renderpath)
    except:
        ## nicht alle benötigten Stringstücke gefunden
        return None

    
def save_by_convention():
    "Fragt nach Dateinamen etc... und speichert das Script nach der Dateinamenskonvention."

    scriptinfo = parse_scriptname()
    dlg = FileSavePanel(scriptinfo)
    if dlg.showModalDialog():
        scriptinfo = dlg.getInfo()
        newname = build_scriptname(scriptinfo)
        if newname is None:
            nuke.message("Dateiname konnte nicht erstellt werden.\nWurden alle Felder ausgefüllt?")
        else:
            try:
                # existiert Datei schon? (ggf. von einem anderen Artist!)
                ## TODO: Buchstaben hinter Versionsnummer prüfen
                versionlist = scan_shot_directory(os.path.dirname(newname), scriptinfo['identifier'])
                ##print(versionlist)
                try:
                    if versionlist:
                        v = int(scriptinfo['version'])
                        ## wenn zu speichernde Version nicht existiert, wirft versionlist[v] eine Exception
                        if not nuke.ask("v%02d already exists (artist: %s).\nDo you want to save as v%02d (smallest available version) instead?" % (v, versionlist[v], versionlist['nextversion'])):
                            raise RuntimeError
                        # höhere Versionsnummer setzen
                        scriptinfo['version'] = "%02d" % versionlist['nextversion']
                        newname = build_scriptname(scriptinfo)
                except KeyError:
                    # zu speichernde Version existiert eh noch nicht
                    pass
                    
                ## create directory
                newpath = os.path.dirname(newname)
                if newpath and not os.path.exists(newpath):
                    os.makedirs(newpath)
                    print "kellertools: created directory " + newpath

                print "kellertools: saving as " + newname
                nuke.scriptSaveAs(newname)
                nuke.message("Script gespeichert: " + nuke.root().knob("name").value())
            except OSError:
                nuke.message("Error creating " + newpath)
            except RuntimeError:
                nuke.message("Script nicht gespeichert")
    

def update_writenode(scriptinfo, the_node):
    '''
    Stellt den Namen einer Write-Node ein, entweder für finale Renderings, preComps oder
    als Ausgabe nach ProcessedFootage.
    '''
    global default_format
    global render_folder
    global precomp_folder
    global footage_folder
    target = 0
    fileformat = ""
    colorspace = ""
    specifier = ""

    ## Defaults für Fileformat, Colorspace und Renderziel aus der Write-Node ermitteln
    print("kellertools: update_writenode für %s" % the_node.name())
    fileformat = the_node.knob("file_type").value()
    colorspace = the_node.knob("colorspace").value()
    if colorspace == "AlexaV3LogC":
        colorspace = "logC"
    elif colorspace not in ('linear', 'logC', 'sRGB', 'rec709'):
        colorspace = ""
    print("kellertools: ermitteltes fileformat: %s/%s" % (fileformat, colorspace))
    filename = the_node.knob("file").value()
    if precomp_folder.lower() in filename.lower():
        target = 1
    elif footage_folder.lower() in filename.lower():
        target = 2
    s_exp = re.search(r"_v\d+[a-z]?.[^.]+.(\w+)", os.path.basename(filename), re.IGNORECASE)
    try:
        specifier = s_exp.group(1)
    except:
        pass

    ## Dialogfenster anzeigen
    dlg = UpdateWritePanel(the_node.name(), target, fileformat, colorspace, specifier)
    if dlg.showModalDialog():
        writeinfo = dlg.getInfo()
        if writeinfo['colorspace'] == "logC":
            writeinfo['colorspace'] = "AlexaV3LogC"
        elif writeinfo['colorspace'] not in ('sRGB', 'rec709'):
            writeinfo['colorspace'] = "linear"
        writepath = build_renderpath(scriptinfo, writeinfo)
        if writepath:
            the_node.knob("file").fromUserText(writepath)
            the_node.knob("colorspace").setValue(writeinfo['colorspace'])
            the_node.knob("raw").setValue(False)
            print "kellertools: neuer Pfad gesetzt (%s/%s)" % (writeinfo['fileformat'], writeinfo['colorspace'])
        else:
            return False
    return True


def update_writenodes():
    "Ruft update_writenode() für alle selektierten Write-Nodes auf."
    
    scriptinfo = parse_scriptname()
    found_writes = False
    if scriptinfo:
        for the_node in nuke.selectedNodes():
            if the_node.Class() == "Write":
                found_writes = True
                if not update_writenode(scriptinfo, the_node):
                    nuke.message("Fehler beim Einstellen von %s.\nFür PreComps ist ein Specifier nötig." % the_node.name())
        if not found_writes:
            nuke.message("Keine Write-Nodes ausgewählt.")
    else:
        nuke.message("Scriptname entspricht leider nicht der Konvention.")


def update_autowrite():
    '''
    Wird vor dem rendern aufgerufen und erzwingt den richtigen Ausgabepfad, wenn eine
    Write-Node namens "AutoWrite" gefunden wurde.
    '''
    ## TODO
    pass


def create_autowrite():
    '''
    Erstellt eine Write-Node, die immer den richtigen Ausgabepfad beinhaltet.
    '''
    ## TODO
    pass

##fin
