import unreal
import re
# UE5.7.4 版本
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("IKRigDef -> {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())
		print("ikr_controller -> {0}".format(self.ikr_controller.get_name()))

		## 删除旧的解算器
		num_solvers = self.ikr_controller.get_num_solvers()
		for i in range(num_solvers - 1 , -1 , -1):
			_solver_controller = self.ikr_controller.get_solver_controller(i)
			_solver = _solver_controller.get_solver_settings()
			print("solver_controller -> {0}".format(_solver_controller))
			print("solver get_name- {0}".format( _solver_controller.get_class().get_name()))
			print("solver struct - {0}".format( _solver.static_struct()))
			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 debug_print_ikrig_solver_names(self):
		print("===== Unreal IKRig/Solver/FBIK candidates =====")
		matched = []
		for attr_name in dir(unreal):
			if ("IKRig" in attr_name) or ("Solver" in attr_name) or ("FBIK" in attr_name):
				matched.append(attr_name)

		if len(matched) == 0:
			print("No related class names found in unreal module.")
		else:
			for name in sorted(matched):
				print(name)
		print("===== End =====")

	def add_full_body_solver(self,solver_root , retarget_root):
		# UE5.7 文档要求传 UStruct 完整包路径
		fbik_index = self.ikr_controller.add_solver("/Script/IKRig.IKRigFullBodyIKSolver")
		print("fbik_index -> {0}".format(fbik_index))
		if fbik_index >= 0 :
			self.ikr_controller.set_solver_enabled(fbik_index, False)
		else:
			raise RuntimeError(
				"Failed to add FullBodyIK solver. solver_type={0}".format("/Script/IKRig.IKRigFullBodyIKSolver")
			)
		
		if hasattr(self.ikr_controller, "set_start_bone"):
			self.ikr_controller.set_start_bone(solver_root, fbik_index)
		else:
			self.ikr_controller.set_root_bone(solver_root, fbik_index)
		self.ikr_controller.set_retarget_root(retarget_root)
	
	def set_retarget_root(self,retarget_root):
		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('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','')