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

import maya.cmds as cmds
from scripts.utils.Core import (
    GetMeshes,
    Descriptor,
    ReadJson,
    GetNeutralJointTranslations,
    ProgressBar,
    RefreshNeutralJointTranslation
)

def set_head_neutral_joint_translation():
    """
    设置头部模型的关节中性位置
    """
    # 获取所需的网格模型
    head_mesh = GetMeshes(m=0)
    teeth_mesh = GetMeshes(m=1)
    eye_left_mesh = GetMeshes(m=3)
    eye_right_mesh = GetMeshes(m=4)
    
    # 检查必要的模型是否都存在
    if not all(cmds.objExists(mesh) for mesh in [head_mesh, teeth_mesh, eye_left_mesh, eye_right_mesh]):
        return
    
    # 获取关节信息
    json_file = Descriptor(ti="jointsInfo")
    object_list = ReadJson(f=json_file, k="head", t="object")
    
    namespace = "DHIhead:"
    
    # 如果有关节信息，开始处理
    if object_list:
        count = len(object_list)
        
        # 初始化进度条
        ProgressBar(sp=True)
        
        # 创建临时定位器
        locator = "MetaHumanLocatorTemp"
        if not cmds.objExists(locator):
            cmds.spaceLocator(n=locator)
        
        # 设置进度条
        ProgressBar(max_value=count)
        ProgressBar(t="Set Head Joint Translation...")
        
        # 获取中性关节位置
        pos = GetNeutralJointTranslations(h=True)
        
        # 检查数据一致性
        if count != len(pos) // 3:
            raise RuntimeError("count error......")
        
        # 处理每个关节
        for i in range(count):
            # 更新进度条
            ProgressBar(apr=1)
            
            # 获取关节信息
            joints = ReadJson(d=object_list[i], k="joint", t="string")
            
            # 处理命名空间
            joint = f"{namespace}{joints[0]}" if cmds.namespace(exists=namespace) else joints[0]
            
            # 检查是否需要定位
            enable = ReadJson(d=object_list[i], k="locate", t="bool")[0]
            
            if enable:
                # 设置定位器位置
                x = pos[i*3]
                y = pos[i*3+1]
                z = pos[i*3+2]
                cmds.setAttr(f"{locator}.t", x, y, z, type="float3")
                
                # 创建并删除点约束
                constraint = cmds.pointConstraint(locator, joint, weight=1)
                cmds.delete(constraint)
        
        # 清理和刷新
        cmds.delete(locator)
        RefreshNeutralJointTranslation()
        ProgressBar(ep=True)
        cmds.refresh()

# 如果直接运行此脚本
if __name__ == '__main__':
    set_head_neutral_joint_translation()