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

#===================================== 1. Module Imports =====================================
import maya.cmds as cmds
import maya.mel as mel
import sys
import os
import json
from scripts.config import data

try:
    from PySide2 import QtCore, QtGui, QtWidgets
    from shiboken2 import wrapInstance
    print("从PySide2加载Qt和shiboken2")
except ImportError:
    try:
        from PySide6 import QtCore, QtGui, QtWidgets
        from shiboken6 import wrapInstance
        print("从PySide6加载Qt和shiboken6")
    except ImportError:
        try:
            from PySide import QtCore, QtGui, QtWidgets
            from shiboken import wrapInstance
            print("从PySide加载Qt和shiboken")
        except ImportError as e:
            print(f"Qt加载失败: {str(e)}")
            QtCore = QtGui = QtWidgets = None
            wrapInstance = None

#===================================== 2. DNA Definition Class =====================================
class DNADefinition:
    """DNA定义类"""
    def __init__(self):
        self.file_path = ""
        self.content = {}
        
    def load(self, file_path):
        """加载DNA文件"""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"DNA文件不存在: {file_path}")
            
        try:
            with open(file_path, "r") as f:
                self.content = json.load(f)
            self.file_path = file_path
            return True
        except Exception as e:
            cmds.warning(f"加载DNA文件失败: {str(e)}")
            return False
            
    def save(self, file_path=None):
        """保存DNA文件"""
        save_path = file_path or self.file_path
        if not save_path:
            cmds.warning("未指定保存路径")
            return False
            
        try:
            with open(save_path, "w") as f:
                json.dump(self.content, f, indent=4)
            return True
        except Exception as e:
            cmds.warning(f"保存DNA文件失败: {str(e)}")
            return False
            
    def validate(self):
        """验证DNA内容"""
        required_keys = ["joints", "blendshapes", "description"]
        for key in required_keys:
            if key not in self.content:
                return False
        return True

# 全局DNA定义实例
dna_definition = DNADefinition()

def browse_dna_file():
    """浏览DNA文件"""
    file_path, _ = cmds.fileDialog2(
        fileFilter="DNA Files (*.dna)",
        dialogStyle=2,
        fileMode=1
    )
    
    if not file_path:
        return
        
    # 加载DNA文件
    if dna_definition.load(file_path[0]):
        # 更新UI预览
        update_dna_preview()
        # 更新骨骼列表
        update_joint_list()
        # 更新BlendShape列表
        update_blendshape_list()

def update_dna_preview():
    """更新DNA预览"""
    # 获取UI实例
    from scripts.MetaFusion import app
    if not app:
        return
        
    define_tab = app.window.findChild(QtWidgets.QWidget, "DefineTab")
    if not define_tab:
        return
        
    # 更新预览内容
    preview_text = json.dumps(dna_definition.content, indent=4)
    define_tab.dna_preview.setText(preview_text)

def update_joint_list():
    """更新骨骼列表"""
    from scripts.MetaFusion import app
    if not app:
        return
        
    define_tab = app.window.findChild(QtWidgets.QWidget, "DefineTab")
    if not define_tab:
        return
        
    # 清空列表
    define_tab.joint_list.clear()
    
    # 添加骨骼项
    joints = dna_definition.content.get("joints", [])
    for joint in joints:
        item = QtWidgets.QTreeWidgetItem([
            joint["name"],
            str(joint["position"]),
            str(joint["rotation"]),
            str(joint["scale"])
        ])
        define_tab.joint_list.addTopLevelItem(item)

def update_blendshape_list():
    """更新BlendShape列表"""
    from scripts.MetaFusion import app
    if not app:
        return
        
    define_tab = app.window.findChild(QtWidgets.QWidget, "DefineTab")
    if not define_tab:
        return
        
    # 清空列表
    define_tab.bs_list.clear()
    
    # 添加BlendShape项
    blendshapes = dna_definition.content.get("blendshapes", [])
    for bs in blendshapes:
        item = QtWidgets.QTreeWidgetItem([
            bs["name"],
            bs["target"],
            str(bs["weight"])
        ])
        define_tab.bs_list.addTopLevelItem(item)

# 骨骼功能
def add_joint():
    """添加骨骼"""
    # 获取选中的骨骼
    sel = cmds.ls(sl=True, type="joint")
    if not sel:
        cmds.warning("请先选择要添加的骨骼")
        return
        
    # 获取骨骼信息
    joint = sel[0]
    pos = cmds.xform(joint, q=True, ws=True, t=True)
    rot = cmds.xform(joint, q=True, ws=True, ro=True)
    scl = cmds.xform(joint, q=True, ws=True, s=True)
    
    # 添加到DNA定义
    joint_data = {
        "name": joint,
        "position": pos,
        "rotation": rot,
        "scale": scl
    }
    
    if "joints" not in dna_definition.content:
        dna_definition.content["joints"] = []
    dna_definition.content["joints"].append(joint_data)
    
    # 更新UI
    update_joint_list()

def delete_joint():
    """删除骨骼"""
    from scripts.MetaFusion import app
    if not app:
        return
        
    define_tab = app.window.findChild(QtWidgets.QWidget, "DefineTab")
    if not define_tab:
        return
        
    # 获取选中项
    items = define_tab.joint_list.selectedItems()
    if not items:
        cmds.warning("请先选择要删除的骨骼")
        return
        
    # 从DNA定义中删除
    for item in items:
        joint_name = item.text(0)
        joints = dna_definition.content.get("joints", [])
        dna_definition.content["joints"] = [
            j for j in joints if j["name"] != joint_name
        ]
    
    # 更新UI
    update_joint_list()

def modify_joint():
    """修改骨骼"""
    from scripts.MetaFusion import app
    if not app:
        return
        
    define_tab = app.window.findChild(QtWidgets.QWidget, "DefineTab")
    if not define_tab:
        return
        
    # 获取选中项
    items = define_tab.joint_list.selectedItems()
    if not items:
        cmds.warning("请先选择要修改的骨骼")
        return
        
    # 获取场景中选中的骨骼
    sel = cmds.ls(sl=True, type="joint")
    if not sel:
        cmds.warning("请先选择要用于更新的骨骼")
        return
        
    # 更新骨骼信息
    joint = sel[0]
    pos = cmds.xform(joint, q=True, ws=True, t=True)
    rot = cmds.xform(joint, q=True, ws=True, ro=True)
    scl = cmds.xform(joint, q=True, ws=True, s=True)
    
    # 更新DNA定义
    joint_name = items[0].text(0)
    joints = dna_definition.content.get("joints", [])
    for j in joints:
        if j["name"] == joint_name:
            j["position"] = pos
            j["rotation"] = rot
            j["scale"] = scl
            break
    
    # 更新UI
    update_joint_list()

# BlendShape功能
def add_blendshape():
    """添加BlendShape"""
    # 获取选中的BlendShape
    sel = cmds.ls(sl=True, type="blendShape")
    if not sel:
        cmds.warning("请先选择要添加的BlendShape")
        return
        
    # 获取BlendShape信息
    bs = sel[0]
    target = cmds.blendShape(bs, q=True, g=True)[0]
    weight = cmds.getAttr(f"{bs}.weight[0]")
    
    # 添加到DNA定义
    bs_data = {
        "name": bs,
        "target": target,
        "weight": weight
    }
    
    if "blendshapes" not in dna_definition.content:
        dna_definition.content["blendshapes"] = []
    dna_definition.content["blendshapes"].append(bs_data)
    
    # 更新UI
    update_blendshape_list()

def delete_blendshape():
    """删除BlendShape"""
    from scripts.MetaFusion import app
    if not app:
        return
        
    define_tab = app.window.findChild(QtWidgets.QWidget, "DefineTab")
    if not define_tab:
        return
        
    # 获取选中项
    items = define_tab.bs_list.selectedItems()
    if not items:
        cmds.warning("请先选择要删除的BlendShape")
        return
        
    # 从DNA定义中删除
    for item in items:
        bs_name = item.text(0)
        blendshapes = dna_definition.content.get("blendshapes", [])
        dna_definition.content["blendshapes"] = [
            bs for bs in blendshapes if bs["name"] != bs_name
        ]
    
    # 更新UI
    update_blendshape_list()

def modify_blendshape():
    """修改BlendShape"""
    from scripts.MetaFusion import app
    if not app:
        return
        
    define_tab = app.window.findChild(QtWidgets.QWidget, "DefineTab")
    if not define_tab:
        return
        
    # 获取选中项
    items = define_tab.bs_list.selectedItems()
    if not items:
        cmds.warning("请先选择要修改的BlendShape")
        return
        
    # 获取场景中选中的BlendShape
    sel = cmds.ls(sl=True, type="blendShape")
    if not sel:
        cmds.warning("请先选择要用于更新的BlendShape")
        return
        
    # 更新BlendShape信息
    bs = sel[0]
    target = cmds.blendShape(bs, q=True, g=True)[0]
    weight = cmds.getAttr(f"{bs}.weight[0]")
    
    # 更新DNA定义
    bs_name = items[0].text(0)
    blendshapes = dna_definition.content.get("blendshapes", [])
    for b in blendshapes:
        if b["name"] == bs_name:
            b["target"] = target
            b["weight"] = weight
            break
    
    # 更新UI
    update_blendshape_list()

# DNA相关功能
def load_dna_preview(dna_file):
    """加载DNA预览
    Args:
        dna_file: DNA文件路径
    Returns:
        bool: 是否加载成功
    """
    try:
        # 加载DNA文件
        if not dna_definition.load(dna_file):
            return False
            
        # 更新UI预览
        update_dna_preview()
        return True
        
    except Exception as e:
        cmds.warning(f"加载DNA预览失败: {str(e)}")
        return False

def validate_dna(dna_file):
    """验证DNA文件
    Args:
        dna_file: DNA文件路径
    Returns:
        bool: 是否验证通过
    """
    try:
        # 加载DNA文件
        if not dna_definition.load(dna_file):
            return False
            
        # 验证必要字段
        if not dna_definition.validate():
            cmds.warning("DNA文件缺少必要字段")
            return False
            
        # 验证骨骼数据
        joints = dna_definition.content.get("joints", [])
        for joint in joints:
            required = ["name", "position", "rotation", "scale"]
            if not all(key in joint for key in required):
                cmds.warning(f"骨骼数据不完整: {joint.get('name', 'unknown')}")
                return False
                
        # 验证BlendShape数据
        blendshapes = dna_definition.content.get("blendshapes", [])
        for bs in blendshapes:
            required = ["name", "target", "weight"]
            if not all(key in bs for key in required):
                cmds.warning(f"BlendShape数据不完整: {bs.get('name', 'unknown')}")
                return False
                
        return True
        
    except Exception as e:
        cmds.warning(f"验证DNA文件失败: {str(e)}")
        return False

def export_dna_definition(dna_file):
    """导出DNA定义
    Args:
        dna_file: 导出文件路径
    Returns:
        bool: 是否导出成功
    """
    try:
        # 收集场景中的骨骼数据
        joints = []
        for joint in cmds.ls(type="joint"):
            pos = cmds.xform(joint, q=True, ws=True, t=True)
            rot = cmds.xform(joint, q=True, ws=True, ro=True)
            scl = cmds.xform(joint, q=True, ws=True, s=True)
            
            joints.append({
                "name": joint,
                "position": pos,
                "rotation": rot,
                "scale": scl
            })
            
        # 收集场景中的BlendShape数据
        blendshapes = []
        for bs in cmds.ls(type="blendShape"):
            target = cmds.blendShape(bs, q=True, g=True)[0]
            weight = cmds.getAttr(f"{bs}.weight[0]")
            
            blendshapes.append({
                "name": bs,
                "target": target,
                "weight": weight
            })
            
        # 更新DNA内容
        dna_definition.content.update({
            "joints": joints,
            "blendshapes": blendshapes,
            "description": {
                "name": cmds.file(q=True, sn=True, shn=True),
                "version": data.TOOL_VERSION,
                "author": data.TOOL_AUTHOR,
                "date": cmds.date(format="YYYY-MM-DD HH:mm:ss")
            }
        })
        
        # 保存文件
        return dna_definition.save(dna_file)
        
    except Exception as e:
        cmds.warning(f"导出DNA定义失败: {str(e)}")
        return False

def import_dna_definition(dna_file):
    """导入DNA定义
    Args:
        dna_file: DNA文件路径
    Returns:
        bool: 是否导入成功
    """
    try:
        # 验证DNA文件
        if not validate_dna(dna_file):
            return False
            
        # 导入骨骼
        joints = dna_definition.content.get("joints", [])
        for joint_data in joints:
            # 检查骨骼是否存在
            joint_name = joint_data["name"]
            if not cmds.objExists(joint_name):
                # 创建骨骼
                joint = cmds.joint(name=joint_name)
            else:
                joint = joint_name
                
            # 设置变换
            cmds.xform(joint, 
                ws=True,
                t=joint_data["position"],
                ro=joint_data["rotation"],
                s=joint_data["scale"]
            )
            
        # 导入BlendShape
        blendshapes = dna_definition.content.get("blendshapes", [])
        for bs_data in blendshapes:
            # 检查目标是否存在
            if not cmds.objExists(bs_data["target"]):
                cmds.warning(f"BlendShape目标不存在: {bs_data['target']}")
                continue
                
            # 创建或获取BlendShape
            bs_name = bs_data["name"]
            if not cmds.objExists(bs_name):
                bs = cmds.blendShape(
                    bs_data["target"],
                    name=bs_name,
                    frontOfChain=True
                )[0]
            else:
                bs = bs_name
                
            # 设置权重
            cmds.setAttr(f"{bs}.weight[0]", bs_data["weight"])
            
        return True
        
    except Exception as e:
        cmds.warning(f"导入DNA定义失败: {str(e)}")
        return False