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

"""
Rigging function module
绑定系统功能模块 - 提供绑定系统相关的功能函数
"""

#===================================== IMPORT MODULES =====================================
import maya.cmds as cmds
import pymel.core as pm
import importlib
import sys
import os
import random
from PySide2 import QtWidgets, QtCore, QtGui
from functools import partial

#========================================== GLOBALS ========================================
# 存储当前选中的关节和控制器信息
selected_joint = None
selected_controller = None

# 存储关节和控制器的原始属性，用于撤销操作
original_joint_properties = {}
original_controller_properties = {}

#========================================== FUNCTIONS ========================================

#------------------------------------ UI UTILITIES ------------------------------------
def handle_ui_event(event_type, event_data):
    """
    处理UI事件
    
    Args:
        event_type: 事件类型
        event_data: 事件数据
    """
    if event_type == "button_click":
        # 处理按钮点击事件
        print(f"Button clicked: {event_data}")
    elif event_type == "slider_change":
        # 处理滑条改变事件
        print(f"Slider changed: {event_data}")
    else:
        print(f"Unknown event type: {event_type}")

def connect_ui_signals(ui_widget, event_type, event_handler):
    """
    连接UI信号
    
    Args:
        ui_widget: UI控件
        event_type: 事件类型
        event_handler: 事件处理函数
    """
    if event_type == "button_click":
        # 连接按钮点击信号
        ui_widget.clicked.connect(event_handler)
    elif event_type == "slider_change":
        # 连接滑条改变信号
        ui_widget.valueChanged.connect(event_handler)
    else:
        print(f"Unknown event type: {event_type}")

#------------------------------------ JOINT OPERATIONS ------------------------------------
def add_joint():
    """
    添加关节
    
    在场景中创建新的关节并添加到骨骼层级结构中
    """
    # 在场景中创建新的关节
    try:
        # 创建新关节
        new_joint = cmds.joint(name="joint1", position=[0, 0, 0])
        print(f"成功创建关节: {new_joint}")
        return new_joint
    except Exception as e:
        print(f"创建关节失败: {str(e)}")
        return None

def remove_joint():
    """
    移除选中的关节
    
    从场景中删除选中的关节
    """
    # 获取当前选中的关节
    selected = cmds.ls(selection=True, type="joint")
    if not selected:
        print("没有选中关节")
        return False
    
    try:
        # 删除选中的关节
        for joint in selected:
            cmds.delete(joint)
        print(f"成功删除关节: {selected}")
        return True
    except Exception as e:
        print(f"删除关节失败: {str(e)}")
        return False

def duplicate_joint():
    """
    复制选中的关节
    
    复制场景中选中的关节及其属性
    """
    # 获取当前选中的关节
    selected = cmds.ls(selection=True, type="joint")
    if not selected:
        print("没有选中关节")
        return None
    
    try:
        # 复制选中的关节
        duplicated = cmds.duplicate(selected, returnRootsOnly=True)
        print(f"成功复制关节: {duplicated}")
        return duplicated
    except Exception as e:
        print(f"复制关节失败: {str(e)}")
        return None

def update_joint_properties(joint_name, position=None, rotation=None, scale=None):
    """
    更新关节属性
    
    Args:
        joint_name: 关节名称
        position: 位置坐标 [x, y, z]
        rotation: 旋转角度 [rx, ry, rz]
        scale: 缩放比例 [sx, sy, sz]
    """
    if not cmds.objExists(joint_name):
        print(f"关节不存在: {joint_name}")
        return False
    
    try:
        # 更新位置
        if position:
            cmds.setAttr(f"{joint_name}.translateX", position[0])
            cmds.setAttr(f"{joint_name}.translateY", position[1])
            cmds.setAttr(f"{joint_name}.translateZ", position[2])
        
        # 更新旋转
        if rotation:
            cmds.setAttr(f"{joint_name}.rotateX", rotation[0])
            cmds.setAttr(f"{joint_name}.rotateY", rotation[1])
            cmds.setAttr(f"{joint_name}.rotateZ", rotation[2])
        
        # 更新缩放
        if scale:
            cmds.setAttr(f"{joint_name}.scaleX", scale[0])
            cmds.setAttr(f"{joint_name}.scaleY", scale[1])
            cmds.setAttr(f"{joint_name}.scaleZ", scale[2])
        
        print(f"成功更新关节属性: {joint_name}")
        return True
    except Exception as e:
        print(f"更新关节属性失败: {str(e)}")
        return False

def reset_joint_properties(joint_name):
    """
    重置关节属性
    
    Args:
        joint_name: 关节名称
    """
    if not cmds.objExists(joint_name):
        print(f"关节不存在: {joint_name}")
        return False
    
    try:
        # 重置位置
        cmds.setAttr(f"{joint_name}.translateX", 0)
        cmds.setAttr(f"{joint_name}.translateY", 0)
        cmds.setAttr(f"{joint_name}.translateZ", 0)
        
        # 重置旋转
        cmds.setAttr(f"{joint_name}.rotateX", 0)
        cmds.setAttr(f"{joint_name}.rotateY", 0)
        cmds.setAttr(f"{joint_name}.rotateZ", 0)
        
        # 重置缩放
        cmds.setAttr(f"{joint_name}.scaleX", 1)
        cmds.setAttr(f"{joint_name}.scaleY", 1)
        cmds.setAttr(f"{joint_name}.scaleZ", 1)
        
        print(f"成功重置关节属性: {joint_name}")
        return True
    except Exception as e:
        print(f"重置关节属性失败: {str(e)}")
        return False

#------------------------------------ CONTROLLER OPERATIONS ------------------------------------
def add_controller():
    """
    添加控制器
    
    在场景中创建新的控制器
    """
    try:
        # 创建控制器曲线
        circle = cmds.circle(name="controller1", normal=[0, 1, 0], radius=1)[0]
        print(f"成功创建控制器: {circle}")
        return circle
    except Exception as e:
        print(f"创建控制器失败: {str(e)}")
        return None

def remove_controller():
    """
    移除选中的控制器
    
    从场景中删除选中的控制器
    """
    # 获取当前选中的控制器
    selected = cmds.ls(selection=True, type="transform")
    if not selected:
        print("没有选中控制器")
        return False
    
    try:
        # 删除选中的控制器
        for ctrl in selected:
            cmds.delete(ctrl)
        print(f"成功删除控制器: {selected}")
        return True
    except Exception as e:
        print(f"删除控制器失败: {str(e)}")
        return False

def duplicate_controller():
    """
    复制选中的控制器
    
    复制场景中选中的控制器及其属性
    """
    # 获取当前选中的控制器
    selected = cmds.ls(selection=True, type="transform")
    if not selected:
        print("没有选中控制器")
        return None
    
    try:
        # 复制选中的控制器
        duplicated = cmds.duplicate(selected, returnRootsOnly=True)
        print(f"成功复制控制器: {duplicated}")
        return duplicated
    except Exception as e:
        print(f"复制控制器失败: {str(e)}")
        return None

def update_controller_properties(controller_name, position=None, rotation=None, scale=None, color=None):
    """
    更新控制器属性
    
    Args:
        controller_name: 控制器名称
        position: 位置坐标 [x, y, z]
        rotation: 旋转角度 [rx, ry, rz]
        scale: 缩放比例 [sx, sy, sz]
        color: 控制器颜色索引
    """
    if not cmds.objExists(controller_name):
        print(f"控制器不存在: {controller_name}")
        return False
    
    try:
        # 更新位置
        if position:
            cmds.setAttr(f"{controller_name}.translateX", position[0])
            cmds.setAttr(f"{controller_name}.translateY", position[1])
            cmds.setAttr(f"{controller_name}.translateZ", position[2])
        
        # 更新旋转
        if rotation:
            cmds.setAttr(f"{controller_name}.rotateX", rotation[0])
            cmds.setAttr(f"{controller_name}.rotateY", rotation[1])
            cmds.setAttr(f"{controller_name}.rotateZ", rotation[2])
        
        # 更新缩放
        if scale:
            cmds.setAttr(f"{controller_name}.scaleX", scale[0])
            cmds.setAttr(f"{controller_name}.scaleY", scale[1])
            cmds.setAttr(f"{controller_name}.scaleZ", scale[2])
        
        # 更新颜色
        if color is not None:
            shapes = cmds.listRelatives(controller_name, shapes=True)
            if shapes:
                for shape in shapes:
                    cmds.setAttr(f"{shape}.overrideEnabled", 1)
                    cmds.setAttr(f"{shape}.overrideColor", color)
        
        print(f"成功更新控制器属性: {controller_name}")
        return True
    except Exception as e:
        print(f"更新控制器属性失败: {str(e)}")
        return False

def reset_controller_properties(controller_name):
    """
    重置控制器属性
    
    Args:
        controller_name: 控制器名称
    """
    if not cmds.objExists(controller_name):
        print(f"控制器不存在: {controller_name}")
        return False
    
    try:
        # 重置位置
        cmds.setAttr(f"{controller_name}.translateX", 0)
        cmds.setAttr(f"{controller_name}.translateY", 0)
        cmds.setAttr(f"{controller_name}.translateZ", 0)
        
        # 重置旋转
        cmds.setAttr(f"{controller_name}.rotateX", 0)
        cmds.setAttr(f"{controller_name}.rotateY", 0)
        cmds.setAttr(f"{controller_name}.rotateZ", 0)
        
        # 重置缩放
        cmds.setAttr(f"{controller_name}.scaleX", 1)
        cmds.setAttr(f"{controller_name}.scaleY", 1)
        cmds.setAttr(f"{controller_name}.scaleZ", 1)
        
        print(f"成功重置控制器属性: {controller_name}")
        return True
    except Exception as e:
        print(f"重置控制器属性失败: {str(e)}")
        return False

#------------------------------------ DNA OPERATIONS ------------------------------------
def import_dna():
    """
    导入DNA文件
    
    从文件中导入DNA数据
    """
    try:
        # 打开文件对话框选择DNA文件
        file_path = cmds.fileDialog2(fileFilter="DNA Files (*.dna);;All Files (*.*)", dialogStyle=2, fileMode=1)
        if not file_path:
            return None
            
        file_path = file_path[0]  # 获取选中的文件路径
        
        # 这里应该调用DNA导入API
        # 暂时使用打印信息代替
        print(f"导入DNA文件: {file_path}")
        return file_path
    except Exception as e:
        print(f"导入DNA文件失败: {str(e)}")
        return None

def export_dna():
    """
    导出DNA文件
    
    将当前绑定数据导出为DNA文件
    """
    try:
        # 打开文件对话框选择保存路径
        file_path = cmds.fileDialog2(fileFilter="DNA Files (*.dna);;All Files (*.*)", dialogStyle=2, fileMode=0)
        if not file_path:
            return None
            
        file_path = file_path[0]  # 获取选中的文件路径
        
        # 确保文件扩展名为.dna
        if not file_path.lower().endswith(".dna"):
            file_path += ".dna"
        
        # 这里应该调用DNA导出API
        # 暂时使用打印信息代替
        print(f"导出DNA文件: {file_path}")
        return file_path
    except Exception as e:
        print(f"导出DNA文件失败: {str(e)}")
        return None

def calibrate_dna():
    """
    校准DNA数据
    
    校准当前绑定数据与DNA标准
    """
    try:
        # 这里应该调用DNA校准API
        # 暂时使用打印信息代替
        print("校准DNA数据")
        return True
    except Exception as e:
        print(f"校准DNA数据失败: {str(e)}")
        return False

#------------------------------------ UTILITY FUNCTIONS ------------------------------------
def update_ui_list(list_widget, items):
    """
    更新UI列表控件
    
    Args:
        list_widget: 列表控件
        items: 要添加的项目列表
    """
    if not list_widget:
        return
    
    # 清空列表
    list_widget.clear()
    
    # 添加项目
    for item in items:
        list_widget.addItem(item)

def update_group_count(list_widget, group_name):
    """
    更新组标题中的计数
    
    Args:
        list_widget: 列表控件
        group_name: 组名称
    """
    if not list_widget:
        return
    
    # 获取列表项目数量
    count = list_widget.count()
    
    # 更新组标题
    group_box = list_widget.parent()
    if isinstance(group_box, QtWidgets.QGroupBox):
        group_box.setTitle(f"{group_name} ({count})")

#------------------------------------ UI EVENT HANDLERS ------------------------------------
def handle_joint_name_changed(ui_widget):
    """
    处理关节名称变化事件
    
    Args:
        ui_widget: UI控件，应该是一个文本输入框
    """
    global selected_joint
    # 获取当前选中的关节和新名称
    if selected_joint and ui_widget:
        new_name = ui_widget.text()
        if new_name and new_name != selected_joint:
            try:
                # 重命名关节
                cmds.rename(selected_joint, new_name)
                selected_joint = new_name
                return True
            except Exception as e:
                print(f"重命名关节失败: {str(e)}")
                return False
    return False

def handle_joint_position_changed(axis, value=None, ui_widgets=None):
    """
    处理关节位置变化事件
    
    Args:
        axis: 轴向 (0=X, 1=Y, 2=Z)
        value: 新的位置值，如果为None则从ui_widgets中获取
        ui_widgets: UI控件列表 [x_spin, y_spin, z_spin]
    """
    global selected_joint
    if not selected_joint:
        return False
        
    # 获取当前位置
    pos = [
        cmds.getAttr(f"{selected_joint}.translateX"),
        cmds.getAttr(f"{selected_joint}.translateY"),
        cmds.getAttr(f"{selected_joint}.translateZ")
    ]
    
    # 更新指定轴向的位置
    if value is not None:
        pos[axis] = value
    elif ui_widgets and len(ui_widgets) >= 3:
        if axis == 0 and ui_widgets[0]:
            pos[0] = ui_widgets[0].value()
        elif axis == 1 and ui_widgets[1]:
            pos[1] = ui_widgets[1].value()
        elif axis == 2 and ui_widgets[2]:
            pos[2] = ui_widgets[2].value()
    else:
        return False
            
    # 更新关节属性
    return update_joint_properties(selected_joint, position=pos)

def handle_joint_rotation_changed(axis, value=None, ui_widgets=None):
    """
    处理关节旋转变化事件
    
    Args:
        axis: 轴向 (0=X, 1=Y, 2=Z)
        value: 新的旋转值，如果为None则从ui_widgets中获取
        ui_widgets: UI控件列表 [x_spin, y_spin, z_spin]
    """
    global selected_joint
    if not selected_joint:
        return False
        
    # 获取当前旋转
    rot = [
        cmds.getAttr(f"{selected_joint}.rotateX"),
        cmds.getAttr(f"{selected_joint}.rotateY"),
        cmds.getAttr(f"{selected_joint}.rotateZ")
    ]
    
    # 更新指定轴向的旋转
    if value is not None:
        rot[axis] = value
    elif ui_widgets and len(ui_widgets) >= 3:
        if axis == 0 and ui_widgets[0]:
            rot[0] = ui_widgets[0].value()
        elif axis == 1 and ui_widgets[1]:
            rot[1] = ui_widgets[1].value()
        elif axis == 2 and ui_widgets[2]:
            rot[2] = ui_widgets[2].value()
    else:
        return False
            
    # 更新关节属性
    return update_joint_properties(selected_joint, rotation=rot)

def handle_joint_scale_changed(axis, value=None, ui_widgets=None):
    """
    处理关节缩放变化事件
    
    Args:
        axis: 轴向 (0=X, 1=Y, 2=Z)
        value: 新的缩放值，如果为None则从ui_widgets中获取
        ui_widgets: UI控件列表 [x_spin, y_spin, z_spin]
    """
    global selected_joint
    if not selected_joint:
        return False
        
    # 获取当前缩放
    scale = [
        cmds.getAttr(f"{selected_joint}.scaleX"),
        cmds.getAttr(f"{selected_joint}.scaleY"),
        cmds.getAttr(f"{selected_joint}.scaleZ")
    ]
    
    # 更新指定轴向的缩放
    if value is not None:
        scale[axis] = value
    elif ui_widgets and len(ui_widgets) >= 3:
        if axis == 0 and ui_widgets[0]:
            scale[0] = ui_widgets[0].value()
        elif axis == 1 and ui_widgets[1]:
            scale[1] = ui_widgets[1].value()
        elif axis == 2 and ui_widgets[2]:
            scale[2] = ui_widgets[2].value()
    else:
        return False
            
    # 更新关节属性
    return update_joint_properties(selected_joint, scale=scale)

def apply_joint_properties_from_ui(pos_widgets=None, rot_widgets=None, scale_widgets=None):
    """
    从UI控件中应用关节属性
    
    Args:
        pos_widgets: 位置控件列表 [x_spin, y_spin, z_spin]
        rot_widgets: 旋转控件列表 [x_spin, y_spin, z_spin]
        scale_widgets: 缩放控件列表 [x_spin, y_spin, z_spin]
    """
    global selected_joint
    if not selected_joint:
        return False
        
    # 获取UI中的属性值
    pos = None
    rot = None
    scale = None
    
    if pos_widgets and len(pos_widgets) >= 3:
        pos = [
            pos_widgets[0].value() if pos_widgets[0] else 0,
            pos_widgets[1].value() if pos_widgets[1] else 0,
            pos_widgets[2].value() if pos_widgets[2] else 0
        ]
    
    if rot_widgets and len(rot_widgets) >= 3:
        rot = [
            rot_widgets[0].value() if rot_widgets[0] else 0,
            rot_widgets[1].value() if rot_widgets[1] else 0,
            rot_widgets[2].value() if rot_widgets[2] else 0
        ]
    
    if scale_widgets and len(scale_widgets) >= 3:
        scale = [
            scale_widgets[0].value() if scale_widgets[0] else 1,
            scale_widgets[1].value() if scale_widgets[1] else 1,
            scale_widgets[2].value() if scale_widgets[2] else 1
        ]
    
    # 更新关节属性
    return update_joint_properties(selected_joint, position=pos, rotation=rot, scale=scale)

def reset_joint_properties_ui(pos_widgets=None, rot_widgets=None, scale_widgets=None):
    """
    重置关节属性UI
    
    Args:
        pos_widgets: 位置控件列表 [x_spin, y_spin, z_spin]
        rot_widgets: 旋转控件列表 [x_spin, y_spin, z_spin]
        scale_widgets: 缩放控件列表 [x_spin, y_spin, z_spin]
    """
    global selected_joint
    if not selected_joint:
        return False
    
    # 重置关节属性
    reset_joint_properties(selected_joint)
    
    # 更新UI显示
    if pos_widgets and len(pos_widgets) >= 3:
        if pos_widgets[0]:
            pos_widgets[0].setValue(0)
        if pos_widgets[1]:
            pos_widgets[1].setValue(0)
        if pos_widgets[2]:
            pos_widgets[2].setValue(0)
    
    if rot_widgets and len(rot_widgets) >= 3:
        if rot_widgets[0]:
            rot_widgets[0].setValue(0)
        if rot_widgets[1]:
            rot_widgets[1].setValue(0)
        if rot_widgets[2]:
            rot_widgets[2].setValue(0)
    
    if scale_widgets and len(scale_widgets) >= 3:
        if scale_widgets[0]:
            scale_widgets[0].setValue(1)
        if scale_widgets[1]:
            scale_widgets[1].setValue(1)
        if scale_widgets[2]:
            scale_widgets[2].setValue(1)
    
    return True

def update_joint_ui(joint_name, name_widget=None, pos_widgets=None, rot_widgets=None, scale_widgets=None):
    """
    更新关节UI
    
    Args:
        joint_name: 关节名称
        name_widget: 名称控件
        pos_widgets: 位置控件列表 [x_spin, y_spin, z_spin]
        rot_widgets: 旋转控件列表 [x_spin, y_spin, z_spin]
        scale_widgets: 缩放控件列表 [x_spin, y_spin, z_spin]
    """
    global selected_joint
    if not cmds.objExists(joint_name):
        return False
    
    # 更新选中的关节
    selected_joint = joint_name
    
    # 更新名称控件
    if name_widget:
        name_widget.setText(joint_name)
    
    # 获取关节属性
    pos = [
        cmds.getAttr(f"{joint_name}.translateX"),
        cmds.getAttr(f"{joint_name}.translateY"),
        cmds.getAttr(f"{joint_name}.translateZ")
    ]
    
    rot = [
        cmds.getAttr(f"{joint_name}.rotateX"),
        cmds.getAttr(f"{joint_name}.rotateY"),
        cmds.getAttr(f"{joint_name}.rotateZ")
    ]
    
    scale = [
        cmds.getAttr(f"{joint_name}.scaleX"),
        cmds.getAttr(f"{joint_name}.scaleY"),
        cmds.getAttr(f"{joint_name}.scaleZ")
    ]
    
    # 更新UI控件
    if pos_widgets and len(pos_widgets) >= 3:
        if pos_widgets[0]:
            pos_widgets[0].setValue(pos[0])
        if pos_widgets[1]:
            pos_widgets[1].setValue(pos[1])
        if pos_widgets[2]:
            pos_widgets[2].setValue(pos[2])
            
    if rot_widgets and len(rot_widgets) >= 3:
        if rot_widgets[0]:
            rot_widgets[0].setValue(rot[0])
        if rot_widgets[1]:
            rot_widgets[1].setValue(rot[1])
        if rot_widgets[2]:
            rot_widgets[2].setValue(rot[2])
            
    if scale_widgets and len(scale_widgets) >= 3:
        if scale_widgets[0]:
            scale_widgets[0].setValue(scale[0])
        if scale_widgets[1]:
            scale_widgets[1].setValue(scale[1])
        if scale_widgets[2]:
            scale_widgets[2].setValue(scale[2])
    
    return True

def on_selection_changed():
    """
    选择变化事件处理
    当Maya中的选择变化时更新UI
    """
    global selected_joint, selected_controller
    
    # 获取当前选中的对象
    selected = cmds.ls(selection=True)
    
    if not selected:
        # 清除选中状态
        selected_joint = None
        selected_controller = None
        return
    
    # 检查选中的对象类型
    for obj in selected:
        # 检查是否是关节
        if cmds.objectType(obj) == "joint":
            selected_joint = obj
            # 可以在这里添加代码来更新UI
            print(f"选中关节: {selected_joint}")
            break
        
        # 检查是否是控制器（通常是transform节点）
        elif cmds.objectType(obj) == "transform":
            # 这里可以添加额外的检查来确定是否是控制器
            # 例如检查是否有特定的属性或命名规则
            selected_controller = obj
            # 可以在这里添加代码来更新UI
            print(f"选中控制器: {selected_controller}")
            break