import maya.cmds as cmds


def DrawUI():
    windowID = "rsAO_Pass"
    if cmds.window(windowID, exists=True):
        cmds.deleteUI(windowID)
    cmds.window(windowID, mnb=0, mxb=0, t="rsAO Pass")
    cmds.window(windowID, e=1, wh=(180,20), rtf=1, s=0)
    cmds.columnLayout()
    cmds.text("Select objects and hit Apply!\nUndo needs no selection." ,w=170, h=50)
    cmds.rowColumnLayout(nc=3, cw= [(1,80),(2,20) ,(3,80)])
    cmds.button(l="Apply", w=70, c=applyAO)
    cmds.separator(h=10, style="none")
    cmds.button(l="Undo", w=70, c=undoAO)
    cmds.separator(h=10, style="none")
    cmds.showWindow()

def applyAO(*Args):
    selList = cmds.ls(sl=1)
    if cmds.objExists("set_AO"):
        for sel in selList:
            cmds.sets(sel, add="set_AO")
    else:
        cmds.sets(selList, n="set_AO")
    sgList = []


    def rsMaterialAO(mtl, mtlAO):
        cmds.duplicate(mtl, n=mtlAO, rr=1, ic=1)
        # break connections
        for attr in [".diffuse_weight", ".refl_weight", ".ms_amount", ".coat_weight", ".emission_weight"]:
            if cmds.listConnections(mtlAO + attr):
                cmds.disconnectAttr((cmds.listConnections(mtlAO + attr, p=1)[0]), mtlAO + attr)
                cmds.setAttr(mtlAO + attr, 0)
            else:
                cmds.setAttr(mtlAO + attr, 0)
        cmds.connectAttr("rsAO_texture.outColor", mtlAO + ".emission_color", f=1)
        cmds.setAttr(mtlAO + ".emission_weight", 1)
        cmds.setAttr(mtlAO + ".refl_fresnel_mode", 3)

    def rsArch(mtl, mtlAO):
        cmds.duplicate(mtl, n=mtlAO, rr=1, ic=1)
        for attr in [".diffuse_weight", ".reflectivity", ".refl_base", ".incandescent_scale"]:
            if cmds.listConnections(mtlAO + attr):
                cmds.disconnectAttr((cmds.listConnections(mtlAO + attr, p=1)[0]), mtlAO + attr)
                cmds.setAttr(mtlAO + attr, 0)
            else:
                cmds.setAttr(mtlAO + attr, 0)
        cmds.connectAttr("rsAO_texture.outColor", mtlAO + ".additional_color", f=1)
        cmds.setAttr(mtlAO + ".incandescent_scale", 1)

    def rsCarPaint(mtl, mtlAO):
        cmds.shadingNode("RedshiftMaterial", asShader=1, n=mtlAO)
        cmds.setAttr(mtlAO + ".diffuse_weight", 0)
        cmds.setAttr(mtlAO + ".refl_weight", 0)
        if cmds.listConnections(mtl + ".bump_input"):
            normalMap = cmds.listConnections(mtl + ".bump_input", p=1)[0]
            print "normalMap found: " + str(normalMap)
            cmds.connectAttr(normalMap, mtlAO + ".bump_input")

        cmds.setAttr(mtlAO + ".emission_weight", 1)

    def rsBlend(mtl, mtlAO, sg, sgConnection):
        cmds.duplicate(mtl, n=mtlAO, rr=1, ic=1)
        mtlBase = cmds.listConnections(mtlAO + ".baseColor")[0]
        mtlType = cmds.objectType(mtlBase)
        if ("_AO" in mtlBase) or (cmds.objExists(mtlBase + "_AO" )):
            print mtlBase + " is already an AO-Material! Skipping."
        else:
            print mtlBase + " has not been created yet."
            createAO(mtlBase, mtlType, sg, sgConnection, sgIgnore=1)
        #createAO(mtlBase, mtlType, sg, sgConnection)
        mtlBaseAO = mtlBase + "_AO"
        cmds.connectAttr(mtlBaseAO + ".outColor", mtlAO + ".baseColor", f=1)
        for count in range(1,6):
            if cmds.listConnections(mtl + ".layerColor" + str(count)) and (cmds.getAttr(mtl + ".additiveMode" + str(count)) ==0):
                print str(mtl) + ".layerColor" + str(count) + " has a material and is not additive"
                mtlLayer = cmds.listConnections(mtl + ".layerColor" + str(count))[0]
                mltLayerAO = mtlLayer + "_AO"
                mtlType = cmds.objectType(mtlLayer)
                if ("_AO" in mtlLayer) or (cmds.objExists(mtlLayer + "_AO")):
                    print mtlLayer + " is already an AO-Material! Skipping."
                else:
                    print mtlLayer + " has not been created yet."
                    createAO(mtlLayer, mtlType, sg, sgConnection, sgIgnore=1)
                #createAO(mtlLayer, mtlType, sg, sgConnection, sgIgnore=1)
                cmds.connectAttr(mltLayerAO + ".outColor", mtlAO + ".layerColor" + str(count), f=1)
            elif cmds.listConnections(mtl + ".layerColor" + str(count)) and (cmds.getAttr(mtl + ".additiveMode" + str(count)) ==1):
                print str(mtl) + ".layerColor" + str(count) + " has a material but is additive. Breaking connection and ignoring!"
                cmds.disconnectAttr(cmds.listConnections(mtl + ".layerColor" + str(count),p=1)[0], mtlAO + ".layerColor" + str(count))

    def rsSprite(mtl, mtlAO, sg, sgConnection):
        cmds.duplicate(mtl, n=mtlAO, rr=1, ic=1)
        mtlBase = cmds.listConnections(mtlAO + ".input")[0]
        mtlType = cmds.objectType(mtlBase)
        #createAO(mtlBase, mtlType, sg, sgConnection, sgIgnore=1)
        if ("_AO" in mtlBase) or (cmds.objExists(mtlBase + "_AO" )):
            print mtlBase + " is already an AO-Material! Skipping."
        else:
            print mtlBase + " has not been created yet."
            createAO(mtlBase, mtlType, sg, sgConnection, sgIgnore=1)
        mtlBaseAO = mtlBase + "_AO"
        cmds.connectAttr(mtlBaseAO + ".outColor", mtlAO + ".input", f=1)

    def createAO(mtl, mtlType, sg, sgConnection, sgIgnore=0):
        print "_"*20
        print "starting 'createAO'"
        print "mtl is " + mtl
        print "mtlType is " + mtlType
        print "sg is " + sg
        print "sgConnection is " + sgConnection
        print "_"*20
        mtlAO = mtl + "_AO"
        if cmds.objExists("rsAO_texture") != True:
            cmds.shadingNode("RedshiftAmbientOcclusion", au=1, n="rsAO_texture")
            cmds.setAttr("rsAO_texture.numSamples", 256)
            cmds.setAttr("rsAO_texture.spread", 0.75)
            cmds.setAttr("rsAO_texture.fallOff", 20)
            cmds.setAttr("rsAO_texture.maxDistance", 40)

        # do shader specific stuff
        if mtlType == "RedshiftMaterial":
            print "mtl is rsMtl"
            rsMaterialAO(mtl, mtlAO)
        elif mtlType == "RedshiftArchitectural":
            rsArch(mtl, mtlAO)
            print "mtl is rsArch"
        elif mtlType == "RedshiftCarPaint":
            print "mtl is rsCarPaint"
            rsCarPaint(mtl, mtlAO)
        elif mtlType == "RedshiftMaterialBlender":
            rsBlend(mtl, mtlAO, sg, sgConnection)
            print "mtl is rsBlend"
        elif mtlType == "RedshiftSkin":
            print "mtl is rsSkin"
        elif mtlType == "RedshiftSprite":
            rsSprite(mtl, mtlAO, sg, sgConnection)
            print "mtl is rsSprite"
        elif mtlType == "RedshiftSubSurfaceScatter":
            print "mtl is rsSSS"
        elif mtlType == "RedshiftMatteShadowCatcher":
            print "mtl is rsMatte"
        if sgIgnore != True:
            cmds.connectAttr(mtlAO + ".outColor", sg + sgConnection, f=1)

    #start of scripts
    #get all materials without duplicates
    for sel in selList:
        selShape = cmds.listRelatives(sel, typ="shape")[0]
        selSGTemp = cmds.listConnections(selShape, t="shadingEngine")
        for selSG in selSGTemp:
            if selSG not in sgList:
                sgList.append(selSG)

    for sg in sgList:
        try:
            mtl = cmds.listConnections(sg + ".rsSurfaceShader")[0]
            sgConnection = ".rsSurfaceShader"
        except:
            mtl = cmds.listConnections(sg + ".surfaceShader")[0]
            sgConnection = ".surfaceShader"
        mtlType = cmds.objectType(mtl)
        if ("_AO" in mtl) or (cmds.objExists(mtl + "_AO")):
            print mtl + "_AO" + " is already an AO-Material! Skipping."
            try:
                cmds.connectAttr(mtl + "_AO.outColor", sg + sgConnection, f=1)
            except:
                pass
        else:
            print mtl + "_AO" + " has not been created yet."
            createAO(mtl, mtlType, sg, sgConnection)

    cmds.select("rsAO_texture")

def undoAO(*Args):
    selList = cmds.sets("set_AO", q=1)
    sgList = []

    for sel in selList:
        selShape = cmds.listRelatives(sel, typ="shape")[0]
        selSGTemp = cmds.listConnections(selShape, t="shadingEngine")
        for selSG in selSGTemp:
            if selSG not in sgList:
                sgList.append(selSG)


    for sg in sgList:
        try:
            mtlAO = cmds.listConnections(sg + ".rsSurfaceShader")[0]
            sgConnection = ".rsSurfaceShader"
        except:
            mtlAO = cmds.listConnections(sg + ".surfaceShader")[0]
            sgConnection = ".surfaceShader"
        mtl = mtlAO.replace("_AO", "")
        cmds.connectAttr(mtl + ".outColor", sg + sgConnection, f=1)
        cmds.delete(mtlAO)
    cmds.delete("rsAO_texture")
    cmds.delete("set_AO")

DrawUI()