// ===========================================================================
// Copyright 2018 Autodesk, Inc. All rights reserved.
//
// Use of this software is subject to the terms of the Autodesk license
// agreement provided at the time of installation or download, or which
// otherwise accompanies this software in either electronic or hard copy form.
// ===========================================================================
// This script fixes several render layer outAdjustment problems caused by a bug which
// will result damaged scenes. The bug is caused by uncontrolled reference operations when 
// render layer outAdjustment presence.
// 
// If you get errors like this:
// Error:, Connection not made: 'XXXShape.instObjGroups[0].objectGroups[0]' -> 'XXXSG.dagSetMembers[-1]'.,  Destination node will not allow the connection., // 
// This script will help.
//
// Usage:
// 1. Make sure all reference are loaded
// 2. Source this script and run fixRenderLayerOutAdjustmentErrors

global proc fixMissingAdjustments() {
    print "------------------------------------------------\n";
    print (uiRes("m_fixRenderLayerOutAdjustmentErrors.kFixTitle"));
    print "------------------------------------------------\n";
    string $layers[] = `ls -type renderLayer`;
    $currentLayer = `editRenderLayerGlobals -q -currentRenderLayer`;
    for ($l in $layers) {
        if (`referenceQuery -isNodeReferenced $l` && gmatch($l, "*defaultRenderLayer"))
            continue;
        print `format -stringArg $l (uiRes("m_fixRenderLayerOutAdjustmentErrors.kLayer"))`;
        string $conns[] = `listConnections -d 0 -c 1 -p 1 ($l + ".outAdjustments")`;

        for ($i=0; $i<size($conns); $i+=2) {
            $adjPlug = $conns[$i];
            $scnPlug = $conns[$i+1];
            $adjValue = substitute("outPlug", $conns[$i], "outValue");
            string $dsgPlugs[] = `connectionInfo -dfs $adjValue`;
            if(!size($dsgPlugs)) {
                print `format -stringArg $scnPlug (uiRes("m_fixRenderLayerOutAdjustmentErrors.kOutAdjMissing"))`;
                $scnPlugParent = "";
                string $scnParentDstPlug[] = {};
                $scnPlug = `connectionInfo -ges $scnPlug`;
                string $scnDstPlug[] = `connectionInfo -dfs $scnPlug`;
                if(`gmatch $scnPlug "*objectGroups*"`) {
                    $scnPlugParent = `match ".*instObjGroups\\[[0-9]+\\]" $scnPlug`;
                    $scnParentDstPlug = `connectionInfo -dfs $scnPlugParent`;
                }
                $SG = "";
                $defaultSG = "initialShadingGroup";
                $parentSG = "";
                $parentSGPlug = "";
                // Try to find a suitable ShadingGroup which can be used as the outAdjustment
                for($p = 0; $p < size($scnDstPlug); $p++) {
                    $scnDstNode = plugNode($scnDstPlug[$p]);
                    if(`nodeType $scnDstNode` == "shadingEngine") {
                        $SG = $scnDstNode;
                        break;
                    }
                    else if($scnDstNode == "defaultRenderLayer" && $l != "defaultRenderLayer") {
                        $defaultAdjValue = substitute("outPlug", $scnDstPlug[$p], "outValue");
                        string $defaultDsgPlugs[] = `connectionInfo -dfs $defaultAdjValue`;
                        if(size($defaultDsgPlugs)) {
                            $defaultSG = plugNode($defaultDsgPlugs[0]);
                        }
                    }
                }
                if($SG == "") {
                    for($p =0; $p < size($scnParentDstPlug); $p++) {
                        $scnParentDstNode = plugNode($scnParentDstPlug[$p]);
                        if(`nodeType $scnParentDstNode` == "shadingEngine") {
                            $parentSG = $scnParentDstNode;
                            $parentSGPlug = $scnParentDstPlug[$p];
                            break;
                        }
                    }
                }
                if($l == $currentLayer && $SG == "") {
                    if($parentSG == "")
                        connectAttr -na $scnPlug ($defaultSG + ".dagSetMembers");
                    else {
                        disconnectAttr $scnPlugParent $parentSGPlug;
                        connectAttr -na $scnPlug ($parentSG + ".dagSetMembers");
                    }
                }
                if($SG == "")
                    if($parentSG == "")
                        $SG = $defaultSG;
                    else
                        $SG = $parentSG;
                print `format -stringArg $SG -stringArg $scnPlug (uiRes("m_fixRenderLayerOutAdjustmentErrors.kUseOutAdj"))`;
                connectAttr -na $adjValue ($SG + ".dagSetMembers");
            }
        }
    }
}

global proc cleanDuplicateAdjustments() {
    print "------------------------------------------------\n";
    print (uiRes("m_fixRenderLayerOutAdjustmentErrors.kCleanDupOutAdj"));
    print "------------------------------------------------\n";
    string $layers[] = `ls -type renderLayer`;
    for ($l in $layers) {
        if (`referenceQuery -isNodeReferenced $l` && gmatch($l, "*defaultRenderLayer"))
            continue;
        print `format -stringArg $l (uiRes("m_fixRenderLayerOutAdjustmentErrors.kLayer2"))`;
        string $conns[] = `listConnections -d 0 -c 1 -p 1 ($l + ".outAdjustments")`;

        for ($i=0; $i<size($conns); $i+=2) {
            for ($j=$i+2; $j<size($conns); $j+=2) {
                if ($conns[$j+1] == $conns[$i+1]) {
                    print `format -stringArg $conns[$i] -stringArg $conns[$i+1] (uiRes("m_fixRenderLayerOutAdjustmentErrors.kBreakDupPlug"))`;
                    disconnectAttr $conns[$i+1] $conns[$i];
                    $plugName = substitute("outPlug", $conns[$i], "outValue");
                    string $other[] = `connectionInfo -dfs $plugName`;
                    print `format -stringArg $plugName -stringArg $other[0] (uiRes("m_fixRenderLayerOutAdjustmentErrors.kBreakDupVal"))`;
                    disconnectAttr $plugName $other[0];
                    break;
                }
            }
        }
    }
}

global proc fixMismatchAdjustments() {
    print "--------------------------------------------------\n";
    print (uiRes("m_fixRenderLayerOutAdjustmentErrors.kFixMismatch"));
    print "--------------------------------------------------\n";
    $currentLayer = `editRenderLayerGlobals -q -currentRenderLayer`;
    string $layers[] = `ls -type renderLayer`;
    stringArrayInsertAtIndex(0, $layers, $currentLayer);
    stringArrayRemoveDuplicates($layers);
    for ($l in $layers) {
        if (`referenceQuery -isNodeReferenced $l` && gmatch($l, "*defaultRenderLayer"))
            continue;
        string $SGOverrides[] = `listConnections -type "shadingEngine" -s 1 -d 0 ($l+".shadingGroupOverride")`;
        editRenderLayerGlobals -currentRenderLayer $l;
        print `format -stringArg $l (uiRes("m_fixRenderLayerOutAdjustmentErrors.kLayer3"))`;
        string $conns[] = `listConnections -d 0 -c 1 -p 1 ($l + ".outAdjustments")`;
        for ($i=0; $i<size($conns); $i+=2) {
            $adjPlug = $conns[$i];
            $scnPlug = $conns[$i+1];
            $adjValue = substitute("outPlug", $conns[$i], "outValue");
            string $dsgPlugs[] = `connectionInfo -dfs $adjValue`;
            $SG = plugNode($dsgPlugs[0]);
            if(size($SGOverrides))
                $SG = $SGOverrides[0];
            $nodeType = `nodeType $SG`;
            if($nodeType != "shadingEngine")
                continue;
            $scnPlugParent = "";
            string $scnParentDstPlug[] = {};
            $scnPlug = `connectionInfo -ges $scnPlug`;
            string $scnDstPlug[] = `connectionInfo -dfs $scnPlug`;
            if(`gmatch $scnPlug "*objectGroups*"`) {
                $scnPlugParent = `match ".*instObjGroups\\[[0-9]+\\]" $scnPlug`;
                $scnParentDstPlug = `connectionInfo -dfs $scnPlugParent`;
            }
            $done = false;
            for($p = 0; $p < size($scnDstPlug); $p++) {
                $scnDstNode = plugNode($scnDstPlug[$p]);
                if(`nodeType $scnDstNode` == "shadingEngine") {
                    if($SG != $scnDstNode) {
                        print `format -stringArg $adjValue -stringArg $dsgPlugs[0] (uiRes("m_fixRenderLayerOutAdjustmentErrors.kBreakBadAdj"))`;
                        disconnectAttr $adjValue $dsgPlugs[0];
                        print `format -stringArg $adjValue -stringArg $scnDstNode (uiRes("m_fixRenderLayerOutAdjustmentErrors.kMakeGoodAdj"))`;
                        connectAttr -na $adjValue ($scnDstNode + ".dagSetMembers");
                    }
                    $done = true;
                    break;
                }
            }
            if(!$done) {
                for($p =0; $p < size($scnParentDstPlug); $p++) {
                    $scnParentDstNode = plugNode($scnParentDstPlug[$p]);
                    if(`nodeType $scnParentDstNode` == "shadingEngine") {
                        if($SG != $scnParentDstNode) {
                            print `format -stringArg $adjValue -stringArg $dsgPlugs[0] (uiRes("m_fixRenderLayerOutAdjustmentErrors.kBreakBadAdj2"))`;
                            disconnectAttr $adjValue $dsgPlugs[0];
                            print `format -stringArg $adjValue -stringArg $scnParentDstNode (uiRes("m_fixRenderLayerOutAdjustmentErrors.kMakeGoodAdj2"))`;
                            connectAttr -na $adjValue ($scnParentDstNode + ".dagSetMembers");
                        }
                        disconnectAttr $scnPlugParent $scnParentDstPlug[$p];
                        connectAttr -na $scnPlug ($scnParentDstNode + ".dagSetMembers");
                    }
                }
            }
        }
    }
    editRenderLayerGlobals -currentRenderLayer $currentLayer;
}

global proc fixRenderLayerOutAdjustmentErrors()
{
    print "------------------------------------------------\n";
    print (uiRes("m_fixRenderLayerOutAdjustmentErrors.kCleanup"));
    print "------------------------------------------------\n";
    string $referenceNodes[] = `ls -type reference`;
    for($i = 0; $i < size($referenceNodes); $i++) {
        if(gmatch($referenceNodes[$i], "*sharedReferenceNode"))
            continue;
        print `format -stringArg $referenceNodes[$i] (uiRes("m_fixRenderLayerOutAdjustmentErrors.kCleaning"))`;
        file -cleanReference $referenceNodes[$i];
    }
    fixMissingAdjustments;
    cleanDuplicateAdjustments;
    fixMismatchAdjustments;
    print (uiRes("m_fixRenderLayerOutAdjustmentErrors.kDone"));
}