#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
IKFK Switch Tool V1.0
Seamless switching between IK and FK animation for Maya rigs
Compatible with Maya 2018+
"""

import maya.cmds as cmds
import maya.mel as mel


class IKFKSwitchUI(object):
    """IKFK Switch UI Class"""
    
    WINDOW_NAME = "IKFK_Switch_UI"
    WINDOW_TITLE = "IKFK V1.0"
    WINDOW_WIDTH = 300
    WINDOW_HEIGHT = 578
    
    def __init__(self):
        """Initialize UI"""
        self.target_fields = {}
        self.target_values = {
            'target_01': '',  # FK Joint Root
            'target_02': '',  # FK Joint Mid
            'target_03': '',  # FK Joint End
            'target_04': '',  # FK Ctrl Root
            'target_05': '',  # FK Ctrl Mid
            'target_06': '',  # FK Ctrl End
            'target_07': '',  # IK Ctrl Root
            'target_08': '',  # IK Ctrl Pole
            'target_09': '',  # Switch Ctrl
            'target_10': ''   # Switch Attr
        }
        
    def create_ui(self):
        """Create the main UI window"""
        # Delete existing window if it exists
        if cmds.window(self.WINDOW_NAME, exists=True):
            cmds.deleteUI(self.WINDOW_NAME)
        
        # Create window
        cmds.window(
            self.WINDOW_NAME,
            title=self.WINDOW_TITLE,
            widthHeight=(self.WINDOW_WIDTH, self.WINDOW_HEIGHT),
            sizeable=True
        )
        
        # Main layout
        main_layout = cmds.columnLayout(rowSpacing=1, adjustableColumn=True)
        
        # Edit frame
        edit_frame = cmds.frameLayout(
            label="Edit",
            collapsable=True,
            collapse=True,
            collapseCommand=lambda: cmds.window(self.WINDOW_NAME, edit=True, height=self.WINDOW_HEIGHT)
        )
        cmds.columnLayout(rowSpacing=2, adjustableColumn=True)
        cmds.button(label="<<< ADV Build >>>", command=lambda x: self.adv_build())
        cmds.button(label="<<< Empty >>>", command=lambda x: self.empty_fields())
        cmds.setParent('..')
        cmds.setParent('..')
        
        # FK Controls frame
        fk_frame = cmds.frameLayout(label="Load FK Ctrl")
        cmds.columnLayout(rowSpacing=2, adjustableColumn=True)
        
        self._create_text_field_row('target_01', '< FK Joint Root')
        self._create_text_field_row('target_02', '< FK Joint Mid')
        self._create_text_field_row('target_03', '< FK Joint End')
        cmds.separator(style='in', height=5)
        self._create_text_field_row('target_04', '< FK Ctrl Root')
        self._create_text_field_row('target_05', '< FK Ctrl Mid')
        self._create_text_field_row('target_06', '< FK Ctrl End')
        
        cmds.setParent('..')
        cmds.setParent('..')
        
        # IK Controls frame
        ik_frame = cmds.frameLayout(label="Load IK Ctrl")
        cmds.columnLayout(rowSpacing=2, adjustableColumn=True)
        
        self._create_text_field_row('target_07', '< IK Ctrl Root')
        self._create_text_field_row('target_08', '< IK Ctrl Pole')
        
        cmds.setParent('..')
        cmds.setParent('..')
        
        # Switch Control frame
        switch_frame = cmds.frameLayout(label="Load Switch Ctrl")
        cmds.columnLayout(rowSpacing=2, adjustableColumn=True)
        
        self._create_text_field_row('target_09', '< Switch Ctrl')
        
        # Special row for attribute selection
        cmds.rowLayout(numberOfColumns=2, adjustableColumn=1, columnWidth2=(45, 60))
        self.target_fields['target_10'] = cmds.textField(text='')
        cmds.button(width=120, label='< Switch Attr', command=lambda x: self.load_attr())
        cmds.setParent('..')
        
        cmds.setParent('..')
        cmds.setParent('..')
        
        # Execute button
        cmds.separator(style='in', height=5)
        cmds.button(height=32, label="<<< Build Seamless Switching >>>", command=lambda x: self.execute())
        
        # Show window
        cmds.window(self.WINDOW_NAME, edit=True, height=self.WINDOW_HEIGHT)
        cmds.showWindow(self.WINDOW_NAME)
    
    def _create_text_field_row(self, target_name, button_label):
        """Create a text field row with button"""
        cmds.rowLayout(numberOfColumns=2, adjustableColumn=1, columnWidth2=(45, 60))
        self.target_fields[target_name] = cmds.textField(text=self.target_values[target_name])
        cmds.button(
            width=120,
            label=button_label,
            command=lambda x: self.set_text_field(target_name)
        )
        cmds.setParent('..')
    
    def set_text_field(self, target_name):
        """Set text field from selection"""
        selection = cmds.ls(selection=True)
        if selection:
            obj_name = selection[0]
            cmds.textField(self.target_fields[target_name], edit=True, text=obj_name)
            self.target_values[target_name] = obj_name
    
    def load_attr(self):
        """Load attribute from channel box selection"""
        try:
            # Get selected attributes from channel box
            attrs = cmds.channelBox('mainChannelBox', query=True, selectedMainAttributes=True)
            if not attrs:
                cmds.error("Select an attribute...")
                return
            
            attr_name = attrs[0]
            cmds.textField(self.target_fields['target_10'], edit=True, text=attr_name)
            self.target_values['target_10'] = attr_name
        except Exception as e:
            cmds.warning("Failed to load attribute: {}".format(e))
    
    def execute(self):
        """Execute the seamless switching setup"""
        # Get values from text fields
        fk_joint_root = cmds.textField(self.target_fields['target_01'], query=True, text=True)
        fk_joint_mid = cmds.textField(self.target_fields['target_02'], query=True, text=True)
        fk_joint_end = cmds.textField(self.target_fields['target_03'], query=True, text=True)
        fk_ctrl_root = cmds.textField(self.target_fields['target_04'], query=True, text=True)
        fk_ctrl_mid = cmds.textField(self.target_fields['target_05'], query=True, text=True)
        fk_ctrl_end = cmds.textField(self.target_fields['target_06'], query=True, text=True)
        ik_ctrl_root = cmds.textField(self.target_fields['target_07'], query=True, text=True)
        ik_ctrl_pole = cmds.textField(self.target_fields['target_08'], query=True, text=True)
        switch_ctrl = cmds.textField(self.target_fields['target_09'], query=True, text=True)
        switch_attr = cmds.textField(self.target_fields['target_10'], query=True, text=True)
        
        # Call the seamless switching function
        # Note: IK joints use FK joints for reference (same as MEL version)
        seamless_switching(
            fk_joint_root, fk_joint_mid, fk_joint_end,
            fk_ctrl_root, fk_ctrl_mid, fk_ctrl_end,
            fk_joint_root, fk_joint_mid, fk_joint_end,  # IK joints = FK joints
            ik_ctrl_root, ik_ctrl_pole,
            switch_ctrl, switch_attr
        )
    
    def adv_build(self):
        """Build ADV rig presets"""
        # Right arm
        seamless_switching(
            "FKXShoulder_R", "FKXElbow_R", "FKXWrist_R",
            "FKShoulder_R", "FKElbow_R", "FKWrist_R",
            "IKXShoulder_R", "IKXElbow_R", "IKXWrist_R",
            "IKArm_R", "PoleArm_R",
            "FKIKArm_R", "FKIKBlend"
        )
        
        # Left arm
        seamless_switching(
            "FKXShoulder_L", "FKXElbow_L", "FKXWrist_L",
            "FKShoulder_L", "FKElbow_L", "FKWrist_L",
            "IKXShoulder_L", "IKXElbow_L", "IKXWrist_L",
            "IKArm_L", "PoleArm_L",
            "FKIKArm_L", "FKIKBlend"
        )
        
        # Right leg
        seamless_switching(
            "FKXHip_R", "FKXKnee_R", "FKXAnkle_R",
            "FKHip_R", "FKKnee_R", "FKAnkle_R",
            "IKXHip_R", "IKXKnee_R", "IKXAnkle_R",
            "IKLeg_R", "PoleLeg_R",
            "FKIKLeg_R", "FKIKBlend"
        )
        
        # Left leg
        seamless_switching(
            "FKXHip_L", "FKXKnee_L", "FKXAnkle_L",
            "FKHip_L", "FKKnee_L", "FKAnkle_L",
            "IKXHip_L", "IKXKnee_L", "IKXAnkle_L",
            "IKLeg_L", "PoleLeg_L",
            "FKIKLeg_L", "FKIKBlend"
        )
    
    def empty_fields(self):
        """Clear all text fields"""
        for target_name in self.target_values.keys():
            self.target_values[target_name] = ''
            if target_name in self.target_fields:
                cmds.textField(self.target_fields[target_name], edit=True, text='')


def seamless_switching(
    fk_joint_root, fk_joint_mid, fk_joint_end,
    fk_ctrl_root, fk_ctrl_mid, fk_ctrl_end,
    ik_joint_root, ik_joint_mid, ik_joint_end,
    ik_ctrl_root, ik_ctrl_pole,
    switch_ctrl, switch_attr
):
    """
    Setup seamless IKFK switching
    
    Args:
        fk_joint_root: FK joint root
        fk_joint_mid: FK joint mid
        fk_joint_end: FK joint end
        fk_ctrl_root: FK control root
        fk_ctrl_mid: FK control mid
        fk_ctrl_end: FK control end
        ik_joint_root: IK joint root
        ik_joint_mid: IK joint mid
        ik_joint_end: IK joint end
        ik_ctrl_root: IK control root
        ik_ctrl_pole: IK pole vector control
        switch_ctrl: Switch control
        switch_attr: Switch attribute name
    """
    # Create locator for storing data
    locator_name = "{}_Switch_Locator".format(switch_ctrl)
    
    if not cmds.objExists(locator_name):
        locator_shape = cmds.createNode('locator', name=locator_name)
        locator_transform = cmds.listRelatives(locator_shape, parent=True)[0]
        cmds.parent(locator_shape, switch_ctrl, shape=True, add=True)
        cmds.delete(locator_transform)
        cmds.setAttr("{}.visibility".format(locator_name), 0)
    
    # Add IKFK_Seamless attribute if it doesn't exist
    if not cmds.objExists("{}.IKFK_Seamless".format(switch_ctrl)):
        cmds.addAttr(switch_ctrl, longName="IKFK_Seamless", attributeType="enum", enumName="IK:FK:", keyable=False)
        cmds.setAttr("{}.IKFK_Seamless".format(switch_ctrl), channelBox=True)
    
    # Add Location attribute to all objects
    objs = [
        fk_joint_root, fk_joint_mid, fk_joint_end,
        fk_ctrl_root, fk_ctrl_mid, fk_ctrl_end,
        ik_joint_root, ik_joint_mid, ik_joint_end,
        ik_ctrl_root, ik_ctrl_pole, switch_ctrl
    ]
    
    for obj in objs:
        # Check if object name is not empty and exists
        if obj and cmds.objExists(obj):
            if not cmds.objExists("{}.Location".format(obj)):
                cmds.addAttr(obj, longName="Location", dataType="string", keyable=True)
        else:
            cmds.warning("Object does not exist or is empty: {}".format(obj))
    
    # Add attributes to locator and connect
    attrs_data = {
        "IKFK_Seamless_Switching": locator_name,
        "FK_Joint_Root": fk_joint_root,
        "FK_Joint_Mid": fk_joint_mid,
        "FK_Joint_End": fk_joint_end,
        "FK_Ctrl_Root": fk_ctrl_root,
        "FK_Ctrl_Mid": fk_ctrl_mid,
        "FK_Ctrl_End": fk_ctrl_end,
        "IK_Joint_Root": ik_joint_root,
        "IK_Joint_Mid": ik_joint_mid,
        "IK_Joint_End": ik_joint_end,
        "IK_Ctrl_Root": ik_ctrl_root,
        "IK_Ctrl_Pole": ik_ctrl_pole,
        "Switch_Ctrl": switch_ctrl,
        "Switch_Attr": switch_attr
    }
    
    for attr_name, attr_value in attrs_data.items():
        if not cmds.objExists("{}.{}".format(locator_name, attr_name)):
            cmds.addAttr(locator_name, longName=attr_name, dataType="string", keyable=True)
        
        cmds.setAttr("{}.{}".format(locator_name, attr_name), attr_value, type="string")
        
        # Connect to Location attribute (except for special attributes)
        if attr_name not in ["IKFK_Seamless_Switching", "Switch_Attr"]:
            target_obj = attr_value
            if cmds.objExists("{}.Location".format(target_obj)):
                if not cmds.isConnected("{}.{}".format(locator_name, attr_name), "{}.Location".format(target_obj)):
                    cmds.connectAttr("{}.{}".format(locator_name, attr_name), "{}.Location".format(target_obj), force=True)
    
    # Create switching script
    create_switching_script()
    
    # Create script job for this control
    create_script_job(switch_ctrl, locator_name)
    
    print("IKFK Seamless Switching setup completed for {}".format(switch_ctrl))


def create_switching_script():
    """Create the global switching script"""
    script_node_name = "SG_IKFK_Switching_Script"
    
    # MEL script for switching logic
    mel_script = '''
global proc sg_switching (string $Switch)
{
    int $state_tmp;
    if (!`optionVar -ex $Switch`){
        $state_tmp = 0;
    }
    $state_tmp = `optionVar -q $Switch`;
    string $tmp[] = `listConnections ($Switch+".Switch_Ctrl")`;
    string $Switch_Ctrl = $tmp[0];
    int $state=`getAttr ($Switch_Ctrl+".IKFK_Seamless")`;
    if($state_tmp == $state)
    {
        return;
    }
    
    string $sel[] = `ls -sl`;
    string $tmp01[] = `listConnections ($Switch+".FK_Joint_Root")`;
    string $tmp02[] = `listConnections ($Switch+".FK_Joint_Mid")`;
    string $tmp03[] = `listConnections ($Switch+".FK_Joint_End")`;
    string $tmp04[] = `listConnections ($Switch+".FK_Ctrl_Root")`;
    string $tmp05[] = `listConnections ($Switch+".FK_Ctrl_Mid")`;
    string $tmp06[] = `listConnections ($Switch+".FK_Ctrl_End")`;
    string $tmp07[] = `listConnections ($Switch+".IK_Joint_Root")`;
    string $tmp08[] = `listConnections ($Switch+".IK_Joint_Mid")`;
    string $tmp09[] = `listConnections ($Switch+".IK_Joint_End")`;
    string $tmp10[] = `listConnections ($Switch+".IK_Ctrl_Root")`;
    string $tmp11[] = `listConnections ($Switch+".IK_Ctrl_Pole")`;
    
    string $FK_Joint_Root = $tmp01[0];
    string $FK_Joint_Mid = $tmp02[0];
    string $FK_Joint_End = $tmp03[0];
    string $FK_Ctrl_Root = $tmp04[0];
    string $FK_Ctrl_Mid = $tmp05[0];
    string $FK_Ctrl_End = $tmp06[0];
    string $IK_Joint_Root = $tmp07[0];
    string $IK_Joint_Mid = $tmp08[0];
    string $IK_Joint_End = $tmp09[0];
    string $IK_Ctrl_Root = $tmp10[0];
    string $IK_Ctrl_Pole = $tmp11[0];
    
    string $Switch_Attr = `getAttr ($Switch+".Switch_Attr")`;
    
    int $min = `addAttr -q -min ($Switch_Ctrl+"."+$Switch_Attr)`;
    int $max = `addAttr -q -max ($Switch_Ctrl+"."+$Switch_Attr)`;
    int $time=`currentTime -q`;
    setKeyframe $FK_Ctrl_Root;
    setKeyframe $FK_Ctrl_Mid;
    setKeyframe $FK_Ctrl_End;
    setKeyframe $IK_Ctrl_Root;
    setKeyframe $IK_Ctrl_Pole;
    if($state==0){
        string $tempGroup_A = `group -em`;
        string $tempGroup_B = `group -em`;
        delete `parentConstraint -weight 1 $IK_Joint_End $tempGroup_A`;
        delete `parentConstraint -weight 1 $IK_Joint_End $tempGroup_B`;
        delete `orientConstraint -offset 0 0 0 -weight 1 $IK_Ctrl_Root $tempGroup_B`;
        parent $tempGroup_B $tempGroup_A;
        delete `parentConstraint -weight 1 $FK_Joint_End $tempGroup_A`;
        string $con_A[] = `parentConstraint -weight 1 $tempGroup_B $IK_Ctrl_Root`;
        setKeyframe $IK_Ctrl_Root;
        delete $con_A;
        string $con_B[] = `pointConstraint -offset 0 0 0 -weight 1 $FK_Joint_Mid $IK_Ctrl_Pole`;
        setKeyframe $IK_Ctrl_Pole;
        delete $con_B;
        setAttr ($Switch_Ctrl+"."+$Switch_Attr) $min;
        setKeyframe -t ($time-1) -at $Switch_Attr $Switch_Ctrl;
        setAttr ($Switch_Ctrl+"."+$Switch_Attr) $max;
        setKeyframe -t $time -at $Switch_Attr $Switch_Ctrl;
        delete $tempGroup_A;
    }
    if($state==1){
        string $con_A[] = `parentConstraint -weight 1 $IK_Joint_Root $FK_Ctrl_Root`;
        setKeyframe $FK_Ctrl_Root;
        delete $con_A;
        string $con_B[] = `parentConstraint -weight 1 $IK_Joint_Mid $FK_Ctrl_Mid`;
        setKeyframe $FK_Ctrl_Mid;
        delete $con_B;
        string $con_C[] = `parentConstraint -weight 1 $IK_Joint_End $FK_Ctrl_End`;
        setKeyframe $FK_Ctrl_End;
        delete $con_C;
        setAttr ($Switch_Ctrl+"."+$Switch_Attr) $max;
        setKeyframe -t ($time-1) -at $Switch_Attr $Switch_Ctrl;
        setAttr ($Switch_Ctrl+"."+$Switch_Attr) $min;
        setKeyframe -t $time -at $Switch_Attr $Switch_Ctrl;
    }
    optionVar -iv $Switch $state;
    select -r $sel;
}
'''
    
    if not cmds.objExists(script_node_name):
        cmds.scriptNode(beforeScript=mel_script, name=script_node_name)
    else:
        cmds.scriptNode(script_node_name, edit=True, beforeScript=mel_script)
    
    cmds.setAttr("{}.scriptType".format(script_node_name), 1)
    cmds.scriptNode(script_node_name, executeBefore=True)


def create_script_job(switch_ctrl, locator_name):
    """Create script job for automatic switching"""
    script_node_name = "{}_Switching_Script".format(switch_ctrl)
    
    # MEL script for script job creation
    mel_script = '''
string $Locator_all[]=`ls -typ "locator"`;
for($Locator in $Locator_all){{
    if(`objExists ($Locator+".IKFK_Seamless_Switching")`){{
        string $tmp=`getAttr ($Locator+".IKFK_Seamless_Switching")`;
        if($tmp=="{0}"){{
            string $Switch[]=`listConnections ($Locator+".Switch_Ctrl")`;
            scriptJob -ac ($Switch[0]+".IKFK_Seamless") ("sg_switching "+$Locator) -kws;
        }}
    }}
}}
'''.format(locator_name)
    
    if not cmds.objExists(script_node_name):
        cmds.scriptNode(beforeScript=mel_script, name=script_node_name)
    else:
        cmds.scriptNode(script_node_name, edit=True, beforeScript=mel_script)
    
    cmds.setAttr("{}.scriptType".format(script_node_name), 1)
    cmds.scriptNode(script_node_name, executeBefore=True)


def show():
    """Main entry point to show the UI"""
    ui = IKFKSwitchUI()
    ui.create_ui()


# For backwards compatibility
def start():
    """Alternative entry point"""
    show()


if __name__ == "__main__":
    show()