// ===========================================================================
// 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.
// ===========================================================================
//
//  Procedure Name:
//	    AEswitchControlTemplate
//
//  Description Name;
//	    Creates the custom attribute controls for the switch node.
//
//  Input Value:
//	    nodeName        (currently unused)
//      inputAttr
//      shapeAttr
//      shadeAttr
//
//  Output Value:
//	    None
// 

proc string switchPlug( string $switchNode,
                        string $inputAttr,
                        string $inputIndex,
                        string $compoundAttr )
{
    return ($switchNode+"."+$inputAttr+"["+$inputIndex+"]."+$compoundAttr);
}

proc addLast( string $itemArray[], string $item )
{
    $itemArray[size($itemArray)] = $item;
}

proc downStreamNetwork( string $node, string $result[], string $traversed[] )
{
	// Check if node has already been traversed
	for ($item in $traversed) {
		if ($item == $node) {
			return;
		}
	}
	addLast( $traversed, $node );
	
	if (`objectType $node` == "shadingEngine") {
		// No need to check if the node is already in the result,
		// because we've already checked if it's been traversed
        addLast($result, $node);
    }
    else {
        string $connections[] = `listConnections -connections true
            -source false -destination true $node`;

        // For each destination (output) connection on node.
        for ($i = 0; $i < size($connections); $i += 2) {
            string $srcPlug = $connections[$i];
            string $dstNode = $connections[$i + 1];

            // Is it not a message connection ?
            if (plugAttr($srcPlug) != "message") {
                // We follow the first non-message output connection.
				downStreamNetwork($dstNode, $result, $traversed);
            }
        }
    }
}

//
//  Procedure Name:
//      AEswitchRemoveCB
//
//  Description:
//      Procedure to remove a row from the switch table.
//
global proc AEswitchRemoveCB( string $shapePlug,
                              string $shadePlug)
{
    string $shapeCons[] = `listConnections -plugs true $shapePlug`;

    for ($item in $shapeCons) {
        disconnectAttr $item $shapePlug;
    }

    string $shadeCons[] = `listConnections -plugs true $shadePlug`;

    for ($item in $shadeCons) {
        disconnectAttr $item $shadePlug;
    }
}

//
//  Procedure Name:
//      AEswitchMenuCB
//
//  Description:
//      Procedure to build a popup menu specific to the spread sheet
//      row the mouse is above. This menu will be displayed when the
//      user presses the RMB.
//
global proc AEswitchMenuCB( string $parent,
                            string $menuName,
                            string $inputAttr,
                            string $shapeAttr,
                            string $shadeAttr )
{
    setParent $parent;

    setParent -m $menuName;

    popupMenu -e -deleteAllItems $menuName;

    string $switchNode = `switchTable -q -switchNode switchTable`;

    int $row = `switchTable -q -underPointerRow switchTable`;

    if ($row == -1) return;

    // Create shape and shade plugs for the element under the pointer.
    //

    string $shapePlug = switchPlug($switchNode, $inputAttr, $row, $shapeAttr);
    string $shadePlug = switchPlug($switchNode, $inputAttr, $row, $shadeAttr);

    // Create menu item to show the AE of the nodes connected to shape plug
    //

    string $shapeCons[] = `listConnections $shapePlug`;

    for ($item in $shapeCons) {
        menuItem -l ("Edit "+$item+" ...") -c ("showEditor "+$item);
    }

    // Create menu item to show the AE of the nodes connected to shade plug
    //

    string $shadeCons[] = `listConnections -plugs true $shadePlug`;

    for ($item in $shadeCons) {
        menuItem -l ("Edit "+$item+" ...") -c ("showEditor "+$item);
    }

    // Create a menu item to map a new shader to the shade plug
    //

    menuItem -l "Map ..."
        -c ("createRenderNode -allWithTexturesUp "+
            ("\"defaultNavigation -connectToExisting -source %node "+
             "-destination "+$shadePlug+"\" ")+
            "\"\"");

    // Create a menu item to disconnect the shape and shade plugs.
    //

    menuItem -l "Remove item"
        -c ("AEswitchRemoveCB "+$shapePlug+" "+$shadePlug);
}

//
//  Procedure Name:
//      AEswitchAddButtonCB
//
//  Description:
//      Procedure to connect all surfaces downstream of the
//      switch. This be invoked from the add button in the AE.
//
global proc AEswitchAddButtonCB( string $switchNode,
                                 string $inputAttr,
                                 string $shapeAttr,
                                 string $shadeAttr )
{
    int $nextAvailable = 0;

    string $currentSurfaces[];

    // Create the current list of surfaces. We filter switchConnections
    // because it will contain connections to both the shape and shade.
    //

    string $switchConnections[] = `listConnections -connections true
        -plugs true ($switchNode+"."+$inputAttr)`;

    for ($i = 0; $i < size($switchConnections); $i += 2) {
        // Is this a shape connection ?
        if (match("\\."+$shapeAttr+"$", $switchConnections[$i]) != "") {
            addLast($currentSurfaces, $switchConnections[$i+1]);
        }
    }

    // Search downstream to find all shading groups
    //

    string $shadingGroups[];
	downStreamNetwork( $switchNode, $shadingGroups, {} );

    if (size($shadingGroups) == 0) {
		string $msg = (uiRes("m_AEswitchControlTemplate.kNotConnected"));
        error -showLineNumber false `format -stringArg $switchNode $msg`;
    }

    // Foreach shading group add the surfaces not already connected
    //

    int $wereSurfacesAdded = false;

    for ($group in $shadingGroups) {
        string $surfaces[] =
            `listConnections -plugs true ($group+".dagSetMembers")`;

        for ($surface in $surfaces) {
            // Ensure the surface is not already connected
            //

            int $isFound = false;

            for ($currentSurface in $currentSurfaces) {
                if ($surface == $currentSurface) {
                    $isFound = true;
                    break;
                }
            }

            if (!$isFound) {
                $wereSurfacesAdded = true;

                // Find the next available connection slot
                //

                string $shapePlug, $shadePlug;

                while (true) {
                    $shapePlug = switchPlug( $switchNode, $inputAttr,
                                             $nextAvailable, $shapeAttr );

                    $shadePlug = switchPlug( $switchNode, $inputAttr,
                                             $nextAvailable, $shadeAttr );

                    string $shape[] = `listConnections $shapePlug`;

                    string $shade[] = `listConnections $shadePlug`;

                    // Is there no connections to the shape and shade ?
                    if (size($shape) == 0 && size($shade) == 0)
                        break;

                    $nextAvailable += 1;
                }

                connectAttr $surface $shapePlug;

                // The same surface cannot be in two shading groups,
                // but it doesn't hurt to ensure we don't add it
                // again.

                addLast($currentSurfaces, $surface);
            }
        }
    }

    if (!$wereSurfacesAdded) {
		string $msg = (uiRes("m_AEswitchControlTemplate.kNoMoreSurface"));
        warning -showLineNumber false `format -stringArg $switchNode $msg`;
    }
}

//
//  Procedure Name:
//      AEswitchRemoveButtonCB
//
//  Description:
//      Procedure to break connections to the shape and shade of the
//      selected switch table row. This will be invoked from the
//      remove button in the AE.
//
global proc AEswitchRemoveButtonCB( string $switchNode,
                                    string $inputAttr,
                                    string $shapeAttr,
                                    string $shadeAttr )
{
    int $row = `switchTable -q -selectedRow switchTable`;

    if ($row == -1) return;

    string $shapePlug = switchPlug($switchNode, $inputAttr, $row, $shapeAttr);

    string $shadePlug = switchPlug($switchNode, $inputAttr, $row, $shadeAttr);

    AEswitchRemoveCB($shapePlug, $shadePlug);
}

//
//  Procedure Name:
//      AEswitchMapButtonCB
//
//  Description:
//      Procedure to create a new shader and connect it to the
//      selected switch table row. This will be invoked from the map
//      button in the AE.
//
global proc AEswitchMapButtonCB( string $switchNode,
                                 string $inputAttr,
                                 string $shapeAttr,
                                 string $shadeAttr )
{
    int $row = `switchTable -q -selectedRow switchTable`;

    if ($row == -1) return;

    string $shadePlug = switchPlug($switchNode, $inputAttr, $row, $shadeAttr);

    createRenderNode -allWithTexturesUp
        ("defaultNavigation -connectToExisting -source %node -destination "+
         $shadePlug) "";
}

//
//  Procedure Name:
//      AEswitchControlNew
//
//  Description:
//      Procedure to create the switch table for a switch node.
//
global proc AEswitchControlNew( string $inputAttr,
                                string $shapeAttr,
                                string $shadeAttr,
                                string $messagePlug )
{
	global int $gTextColumnWidthIndex;
	global int $gAESingleWidgetWidthIndex;

    separator;

	int	$switchButtonSize = ((3 * $gAESingleWidgetWidthIndex) / 2 );

    rowLayout -numberOfColumns 3 
		-cw 1 $gTextColumnWidthIndex 
		-cw 2 $switchButtonSize 
		-cw 3 $switchButtonSize

		-cal 1 "center" -cal 2 "center" -cal 3 "center"
		-cat 1 "right" 1 -cat 2 "both" 1 -cat 3 "both" 1;

		button -label (uiRes("m_AEswitchControlTemplate.kAddSurfaces")) -w $switchButtonSize addButton;
        button -label (uiRes("m_AEswitchControlTemplate.kRemoveItem")) removeButton;
		button -label (uiRes("m_AEswitchControlTemplate.kMapItem")) mapButton;
        setParent ..;

    separator -style "none";

	formLayout switchForm;
        switchTable switchTable;
        popupMenu -button 3 switchMenu;
        setParent ..;

    formLayout -e
        -height 150
        -af switchTable top 0
        -af switchTable left ($gTextColumnWidthIndex - $switchButtonSize)
        -af switchTable bottom 0
        -af switchTable right 0
        switchForm;

    separator;

    AEswitchControlReplace($inputAttr, $shapeAttr, $shadeAttr, $messagePlug);
}

//
//  Procedure Name:
//      AEswitchControlNew
//
//  Description:
//      Procedure to initialize the switch table for a node.
//
global proc AEswitchControlReplace( string $inputAttr,
                                    string $shapeAttr,
                                    string $shadeAttr,
                                    string $messagePlug )
{
	string $buffer[];
	tokenize($messagePlug, ".", $buffer);
    string $nodeName = $buffer[0];

	// Here, $shapeAttr is the "inShape" child of multi
	// .input[0].  We may or may not have a specific
	// instance here; don't use attributeName here, since
	// we're not actually interested in the UI name of a
	// specific plug like "input[0].inShape".  We just want
	// the indexless name of the attribute as the column
	// title in the AttrEd spreadsheet ("In Shape").
	//
	string $shapeAttrUI = `attributeQuery -nn -node $nodeName $shapeAttr`;
	string $shadeAttrUI = `attributeQuery -nn -node $nodeName $shadeAttr`;

	switchTable -e
        -label1 $shapeAttrUI -label2 $shadeAttrUI
        -switchNode $nodeName
        switchTable;

    string $parent = `setParent -q`;
    string $menuName = $parent + "|switchForm|switchTable|switchMenu";
    popupMenu -e -postMenuCommand
        ("AEswitchMenuCB "
         +$parent+" "+$menuName+" "+$inputAttr+" "+$shapeAttr+" "+$shadeAttr)
        switchMenu;

    button -e -command
        ("AEswitchAddButtonCB "+$nodeName+" "+ $inputAttr+" "+
         $shapeAttr+" "+$shadeAttr) addButton;

    button -e -command
        ("AEswitchRemoveButtonCB "+$nodeName+" "+$inputAttr+" "+
         $shapeAttr+" "+ $shadeAttr) removeButton;

    button -e -command
        ("AEswitchMapButtonCB "+$nodeName+" "+$inputAttr+" "+
         $shapeAttr+" "+ $shadeAttr) mapButton;
}

global proc AEswitchControlTemplate( string $nodeName,
                                     string $inputAttr,
                                     string $shapeAttr,
                                     string $shadeAttr)
{
    editorTemplate -callCustom
        ("AEswitchControlNew "+$inputAttr+" "+$shapeAttr+" "+$shadeAttr)
        ("AEswitchControlReplace "+$inputAttr+" "+$shapeAttr+" "+$shadeAttr)
        "message";
}