import unreal
import re

class IKrig():

	def __init__(self) -> None:
		engineVersion = unreal.SystemLibrary.get_engine_version()
		self.utilityBase = unreal.GlobalEditorUtilityBase.get_default_object()
		
		self.get_select_controll_rig()

	def get_select_controll_rig(self):

		selectedAssets = self.utilityBase.get_selected_assets()

		#要选中 controlrig 资产
		if len(selectedAssets)<1:
			raise Exception('Nothing selected, you must select a ControlRig')

		self.IKRigDef = selectedAssets[0]
		print("{0}".format(self.IKRigDef.get_class().get_name()))

		if self.IKRigDef.get_class().get_name() != 'IKRigDefinition':
			raise Exception('Selected object is not a IKRig, you must select a IKRig')

	def remove_history(self):
		self.ikr_controller = unreal.IKRigController.get_controller(self.IKRigDef)
		ikrig_mesh = self.ikr_controller.get_skeletal_mesh()
		
		#获取分配给IK Rig的骨骼网格体。
		print(ikrig_mesh.get_name())

		## 删除旧的解算器
		num_solvers = self.ikr_controller.get_num_solvers()
		for i in range(num_solvers - 1 , -1 , -1):
			print("remove solver - {0}".format( self.ikr_controller.get_solver_at_index(i).get_name()))
			self.ikr_controller.remove_solver(i)

		## 删除旧的 Goal 控制器
		goal = self.ikr_controller.get_all_goals()
		for i in goal :
			print("remove => {0} -> {1}".format(i.goal_name,i.bone_name))
			self.ikr_controller.remove_goal(i.goal_name)

		## 删除旧的 chain 组件
		chain = self.ikr_controller.get_retarget_chains()
		for i in chain:
			print("remove chain => {0}".format(i.chain_name))
			self.ikr_controller.remove_retarget_chain(i.chain_name)

	def add_full_body_solver(self,solver_root , retarget_root):
		fbik_index = self.ikr_controller.add_solver(unreal.IKRigFBIKSolver)
		self.ikr_controller.set_root_bone(solver_root, 0)
		self.ikr_controller.set_retarget_root(retarget_root)

	def add_retarget_chain(self,chain_name,start_name,end_name,goal_name):
		#root_bone = ikr_controller.get_root_bone(0)
		#print(root_bone)
		#添加chain 组件
		self.ikr_controller.add_retarget_chain(chain_name,start_name,end_name,goal_name)
	
	def add_goal(self,goal_name,bone_name):
		#创建 IK goal 
		#self.ikr_controller.add_new_goal("L-Hand_goal","Bip001-L-Hand")
		self.ikr_controller.add_new_goal(goal_name,bone_name)
		
		#查询 goal 是否被添加到解算器
		is_exist = self.ikr_controller.is_goal_connected_to_solver(goal_name,0)
		
		#把 goal 指认到解算器
		if not is_exist:
			self.ikr_controller.connect_goal_to_solver(goal_name, 0)

		#创建 chain 并指认 Ik goal
		#self.ikr_controller.add_retarget_chain("L-Hand","Bip001-L-Clavicle","Bip001-L-Hand","L-Hand_goal")
		
	def add_bone_setting(self,bone_name):

		#添加 解算器的 谋根骨骼设置
		self.ikr_controller.add_bone_setting(bone_name, 0)

		# for i in dir(ikr_controller):
		# 	print(i)

		# library_controller = self.selectedAsset.get_controller(library)
		# self.hierarchy =self.selectedAsset.hierarchy

		# #创建的控制器
		# self.hierarchy_controller =self. hierarchy.get_controller()

if __name__ == "__main__":

	#https://dev.epicgames.com/documentation/zh-cn/unreal-engine/ik-rig-solvers-in-unreal-engine?application_version=5.3
	test_rig = IKrig()
	test_rig.remove_history()
	#test_rig.add_full_body_solver("Bip001","Root_Bip001")
	test_rig.add_full_body_solver('CC_Base_Hip','RL_BoneRoot')
	test_rig.add_retarget_chain('pelvis','CC_Base_Hip','CC_Base_Pelvis','')
	test_rig.add_retarget_chain('Spine','CC_Base_Waist','CC_Base_Spine02','')
	test_rig.add_retarget_chain('Neck','CC_Base_NeckTwist01','CC_Base_NeckTwist02','')
	test_rig.add_retarget_chain('Head','CC_Base_Head','CC_Base_Head','')
	test_rig.add_goal('LlegIK','CC_Base_L_Foot')
	test_rig.add_retarget_chain('Lleg','CC_Base_L_Thigh','CC_Base_L_Foot','LlegIK')
	test_rig.add_goal('RlegIK','CC_Base_R_Foot')
	test_rig.add_retarget_chain('Rleg','CC_Base_R_Thigh','CC_Base_R_Foot','RlegIK')
	test_rig.add_goal('LarmIK','CC_Base_L_Hand')
	test_rig.add_retarget_chain('Lram','CC_Base_L_Clavicle','CC_Base_L_Hand','LarmIK')
	test_rig.add_retarget_chain('Lfinger1','CC_Base_L_Thumb1','CC_Base_L_Thumb3','')
	test_rig.add_retarget_chain('Lfinger2','CC_Base_L_Index1','CC_Base_L_Index3','')
	test_rig.add_goal('RarmIK','CC_Base_R_Hand')
	test_rig.add_retarget_chain('Rram','CC_Base_R_Clavicle','CC_Base_R_Hand','RarmIK')
	test_rig.add_retarget_chain('Rfinger1','CC_Base_R_Thumb1','CC_Base_R_Thumb3','')
	test_rig.add_retarget_chain('Rfinger2','CC_Base_R_Index1','CC_Base_R_Index3','')