# DazV5 Bodygroup System by LordAardvark
# V2 - A complete rewrite
# Now known as "Material Toggler", due to its generalized nature of just adding/removing $alpha to materials.
# ##

import vs
import os
from PySide import QtCore
from PySide import QtGui

# Globals
MatTog_GameInfoPath = "usermod/gameinfo.txt"
MatTog_GameModFolderNames = []
MatTog_CurAnimset = None

# Utilities
# Make all of the relevant directories in a filepath; Boo not using Python 3.2, means I have to do it this ugly way.
def mkdir_p(path): # http://stackoverflow.com/questions/600268/mkdir-p-functionality-in-python/600612#600612
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if os.path.isdir(path):
            pass
        else:
            raise
            
# Given a filepath, just return the actual filename.
def GetRawFileName(fName):
    return os.path.splitext(os.path.basename(fName))[0]
    
def GetFilePath(filePath):
    strSplit = filePath.replace("/","\\").split("\\")
    
    fileName = strSplit[-1]
    
    return filePath.replace(fileName,"")
    
def GetFileName(filePath):
    strSplit = filePath.replace("/","\\").split("\\")
    
    fileName = strSplit[-1]

    return fileName
    
# A more convenient file-write
def fwrite(f, sTxt, *args):
    f.write((str(sTxt) + "\n").format(*args))
    
# Following two functions are ported from my DMX Packager 2019 tool.
def ParseGameInfo():
    global MatTog_GameModFolderNames
    
    fullGameInfoPath = os.path.join(os.getcwd(), MatTog_GameInfoPath)
    
    del MatTog_GameModFolderNames[:] # Because apparently SFM's implementation of Python doesn't have list.clear() ....
    
    InSearchPaths = False
    
    if not os.path.isfile(fullGameInfoPath):
        print "CRITICAL ERROR :: NO GAMEINFO.TXT FOUND"
        return
        
    with open(fullGameInfoPath,'r') as f:
        for line in f:
            # First just chew through until we find SearchPaths. We don't care about anything else
            if (not InSearchPaths) and ("SearchPaths" in line):
                InSearchPaths = True
                
            # Once we've found SearchPaths, the first "}" means we've dropped out. Stop scanning the file at that point.
            if (InSearchPaths) and ("}" in line):
                break
                
            # Once we've found SearchPaths, we only care about "Game" paths.
            if (not InSearchPaths) or ("Game" not in line):
                continue
                
            # At this point, we've found a Game line. Process it!
            workLine = line.replace("\"","").replace("Game","", 1).strip().replace("\"","") # Strip the first instance of "Game". It might be in the folder name, too. Don't want to replace that. But it will ALWAYS come before the folder name, as a key. MatTogo strip quotes.
            
            if workLine == "|gameinfo_path|.":  # Special path that means "name of the folder containing gameinfo.txt"
                workLine = GetFilePath(MatTog_GameInfoPath)[:-1] # GetFilePath keeps the slash. We don't want it, so leave out the last character.
                
            MatTog_GameModFolderNames.append(workLine)
            
def GetAbsolutePath(relativePath):
    absPath = relativePath
    
    for modName in MatTog_GameModFolderNames:
        absPath = os.path.join(os.getcwd(),modName,relativePath).lower()
        
        if os.path.isfile(absPath):
            return absPath
            
    return None
    
def GetOrCreateMaterial(animset, materialName):
    materialName = str(materialName)
    
    if hasattr(animset.gameModel,"materials"):
        for mat in animset.gameModel.materials:
            if mat is None:
                continue
                
            if str(mat.name) == materialName:
                return mat
                   
    newMat = vs.CreateElement("DmeMaterial", materialName, 0)
        
    if not hasattr(animset.gameModel,"materials"):
        animset.gameModel.AddAttribute('materials', vs.AT_ELEMENT_ARRAY)
        
    animset.gameModel.materials.AddToTail(newMat)
        
    return newMat
    
def EnsureMaterialOverrideExists(animset, matName):
    lastMat = None # Should only ever be one of these per material name, but just in case, return the last one.
        
    shdr = animset.gameModel.GetStudioHdr()
    for j in range(shdr.numcdtextures):
        texpath = shdr.pCdtexture(j)
        fullname = texpath + matName
        
        relPath = os.path.join("materials",fullname + ".vmt")
        absPath = GetAbsolutePath(relPath)
        
        # Only apply this material if it exists. This is to deal with multiple $cdmaterial's on models.
        if not absPath:
            continue
                
        mat = GetOrCreateMaterial(animset,matName)
        mat.SetValue("mtlName",str(fullname))
        mat.SetValue("groupName","Material Overrides")
        
        lastMat = mat
        
    return lastMat
    
def GetMaterialNames(animset):
    names = []
    
    shdr = animset.gameModel.GetStudioHdr()
    for i in range(shdr.numtextures):
        matname = shdr.pTexture(i).pszName()
        
        names.append(matname)
        
    return names
    
def IsMaterialVisible(animset, matName):
    if not hasattr(animset.gameModel,"materials"): # No material overrides at all? Assume it's on.
        return True
        
    matObj = None
    
    for mat in animset.gameModel.materials:
        if mat is None:
            continue
            
        if str(mat.name) == matName:
            matObj = mat
            
    if matObj is None: # No material override for this material? Assume it's on.
        return True
        
    alphaValue = matObj.GetValue("$alpha")
    
    if alphaValue is None: # No $alpha attribute on the material override? Assume it's on.
        return True
        
    alphaValue = float(str(alphaValue))
    
    return alphaValue >= 0.1 # If alphaValue is 0.0, then it is off, and return False. Any other value, assume it's on, and return True.
    
def SetMaterialVisibility(animset, matName, isVisible):    
    if not isVisible and not IsMaterialVisible(animset, matName): # We don't want it to be visible, and it currently isn't visible anyways.... So don't do anything.
        return
        
    if isVisible and IsMaterialVisible(animset, matName): # Similar, we want it to be visible, and it already is visible... So don't do anything.
        return
        
    # Okay, now we have to actually do work
    alphaValue = 1.0
    if not isVisible:
        alphaValue = 0.0
        
    EnsureMaterialOverrideExists(animset, matName)
    
    mat = GetOrCreateMaterial(animset, matName)
    mat.SetValue("$alpha",alphaValue)
    
# UI Utilities
def GetListItems(list):
    r = []
    
    for i in range(list.count()):
        r.append(list.item(i))
        
    return r
    
def ListStripeItems(lstItems):
    # Two different background colors, to stripe with.
    bg1 = QtGui.QColor()
    bg2 = QtGui.QColor()
    selClr = QtGui.QColor()
    
    bg1.setRgb(64,64,64)
    bg2.setRgb(80,80,80)
    selClr.setRgb(80,128,64)
    
    useBg = bg1
    
    for item in lstItems:
        isSel = item.isSelected()
        
        overrideSel = item.data(QtCore.Qt.UserRole+1) # Because selecting it manually isn't working on load, I am using this hacky workaround to get the first item selected on load.
        if overrideSel is not None:
            item.setData(QtCore.Qt.UserRole+1,None)
            isSel = True
            
        if isSel:
            item.setBackground(selClr)
        else:
            item.setBackground(useBg)
            
        if useBg == bg1:
            useBg = bg2
        else:
            useBg = bg1
     
def GetSelectedItems(list):
    selItems = []
    
    uiList = GetListItems(list)
    
    for i in range(0, len(uiList)):
        uiItem = uiList[i]
        
        if uiItem.checkState() is QtCore.Qt.CheckState.Checked:
            selItems.append(uiItem.data(QtCore.Qt.UserRole))
            
    return selItems
    
# UI logic
def GetCandidateAnimsets():
    shot = sfmClipEditor.GetSelectedShots()[0]
    candidates = []
    
    for candSet in shot.animationSets:
        candSetName = str(candSet.name).lower().strip()

        if hasattr(candSet,"gameModel"):
            candidates.append(candSet)
            
    return candidates
    
def AddCandidateAnimsetsToList(list, lstMaterials):
    global MatTog_CurAnimset
    
    candidates = GetCandidateAnimsets()
    
    selItem = None
    for candSet in candidates:
        candSetName = str(candSet.name).lower().strip()
        
        # Size hint
        hint = QtCore.QSize()
        hint.setHeight(24)
        
        # Create the new item
        item = QtGui.QListWidgetItem()
        item.setData(QtCore.Qt.UserRole,candSet)
        item.setText(candSetName)
        item.setSizeHint(hint)
        
        if selItem is None:
            selItem = item
            item.setSelected(True)
            item.setData(QtCore.Qt.UserRole+1,True)
            
            AddCandidateMaterialsToList(lstMaterials, candSet)
            ListStripeItems(GetListItems(lstMaterials))
            
            MatTog_CurAnimset = candSet
            
        list.addItem(item)   
            
def AddCandidateMaterialsToList(list, animset):
    matNames = GetMaterialNames(animset)
    
    for matName in matNames:
        # Size hint
        hint = QtCore.QSize()
        hint.setHeight(24)
        
        checkState = QtCore.Qt.Unchecked
        
        if IsMaterialVisible(animset, matName):
            checkState = QtCore.Qt.Checked
            
        # Create the new item
        item = QtGui.QListWidgetItem()
        item.setText(matName)
        item.setData(QtCore.Qt.UserRole, matName)
        item.setFlags(item.flags() | QtCore.Qt.ItemIsUserCheckable)
        item.setCheckState(checkState)
        item.setSizeHint(hint)
        
        list.addItem(item)

# Hooks
def onSelectAnimset(lstModels, lstMaterials, currentItem, previousItem, window):
    global MatTog_CurAnimset
    
    if currentItem is None:
        return
        
    if previousItem is None:
        return
        
    currentItem.setSelected(True)
    previousItem.setSelected(False)
    
    ListStripeItems(GetListItems(lstModels)) 
    
    lstMaterials.clear()
    
    selAnimset = currentItem.data(QtCore.Qt.UserRole)
    
    MatTog_CurAnimset = selAnimset
    
    AddCandidateMaterialsToList(lstMaterials, selAnimset)
    ListStripeItems(GetListItems(lstMaterials))
    
def onCheckMaterial(lstModels, lstMaterials, changedItem, window):
    global MatTog_CurAnimset
    
    matName = changedItem.data(QtCore.Qt.UserRole)
    animset = MatTog_CurAnimset
    
    if animset is None:
        return
        
    checkState = changedItem.checkState()
    
    isVisible = True
    if checkState == QtCore.Qt.Unchecked:
        isVisible = False
        
    SetMaterialVisibility(animset, matName, isVisible)
        
    
# UI 
def CreateUI():
    # Create the UI window
    window = QtGui.QDialog()
    window.resize(300,400)
    window.setWindowTitle("Material Toggler")

    grid = QtGui.QGridLayout()
    grid.setSpacing(12)
    grid.setColumnMinimumWidth(0, 128)
    grid.setColumnMinimumWidth(1, 128)
    
    # Labels
    lblInstructions = QtGui.QLabel()
    lblInstructions.setText("Select a model to show its materials. If a material is checked,\nit is visible. Toggle its check to toggle its visiblity.")
    grid.addWidget(lblInstructions, 0, 0, 1, 2)
    
    lblModels = QtGui.QLabel()
    lblModels.setText("Models:")
    grid.addWidget(lblModels, 1, 0, 1, 1)
    
    lblMaterials = QtGui.QLabel()
    lblMaterials.setText("Materials:")
    grid.addWidget(lblMaterials, 1, 1, 1, 1)
    
    # Lists
    lstModels = QtGui.QListWidget()
    lstModels.setSpacing(1)
    lstModels.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
    grid.addWidget(lstModels, 2, 0, 5, 1)
    
    lstMaterials = QtGui.QListWidget()
    lstMaterials.setSpacing(1)
    lstMaterials.setSelectionMode(QtGui.QAbstractItemView.ExtendedSelection)
    grid.addWidget(lstMaterials, 2, 1, 5, 1)
    
    lstModels.currentItemChanged.connect(lambda current,previous: onSelectAnimset(lstModels, lstMaterials, current, previous, window))
    lstMaterials.itemChanged.connect(lambda changedItem: onCheckMaterial(lstModels, lstMaterials, changedItem, window))
    
    # Populate lists
    # Make the lists pretty
    AddCandidateAnimsetsToList(lstModels, lstMaterials)
    ListStripeItems(GetListItems(lstModels)) 
    ListStripeItems(GetListItems(lstMaterials)) 
    
    # Save button
    btnSave = QtGui.QPushButton()
    btnSave.setText("Done")
    btnSave.clicked.connect(lambda: window.close()) # Don't need to do anything fancy.
    grid.addWidget(btnSave, 7, 0, 1, 3)
    
    # Display the window
    window.setLayout(grid)
    window.setModal(True)
    window.exec_()

if len(MatTog_GameModFolderNames) < 1:
    ParseGameInfo()
    
CreateUI()
