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

"""
Tool - Mesh Utilities
Provides utilities for working with mesh objects in Maya
"""

import os
import sys
import json
import maya.cmds as cmds
import maya.mel as mel

# Import configuration
import config

class MeshUtils:
    """Utilities for working with mesh objects in Maya"""
    
    def __init__(self):
        """Initialize mesh utilities"""
        pass
    
    def validate_metahuman_topology(self, mesh):
        """
        Validate if a mesh has MetaHuman topology
        
        Args:
            mesh (str): Name of the mesh
            
        Returns:
            bool: True if the mesh has MetaHuman topology, False otherwise
        """
        try:
            # Check if mesh exists
            if not cmds.objExists(mesh):
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            # Get mesh shape
            shapes = cmds.listRelatives(mesh, shapes=True, type="mesh")
            if not shapes:
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            mesh_shape = shapes[0]
            
            # Get vertex count
            vertex_count = cmds.polyEvaluate(mesh, vertex=True)
            
            # MetaHuman topology has a specific vertex count
            # For MetaHuman 4.0, the vertex count is typically around 6700-6800
            # This is a very basic check and should be enhanced with more detailed topology checks
            if vertex_count < 6000 or vertex_count > 7000:
                print(f"Mesh {mesh} has vertex count ({vertex_count}) that does not meet MetaHuman topology requirements (6000-7000)")
                return False
            
            # TODO: Add more detailed topology checks
            # - Check for specific edge loops
            # - Check for specific vertex positions
            # - Check for specific face structure
            
            return True
        except Exception as e:
            cmds.warning(f"Error validating MetaHuman topology: {str(e)}")
            return False
    
    def bind_mesh(self, mesh, skeleton, method="Smooth Bind", max_influences=4, 
                 maintain_max_influences=True, normalize_weights=True, 
                 post_normalize=True, use_metahuman_weights=True):
        """
        Bind a mesh to a skeleton
        
        Args:
            mesh (str): Name of the mesh
            skeleton (str): Name of the skeleton root
            method (str, optional): Binding method
            max_influences (int, optional): Maximum number of joint influences per vertex
            maintain_max_influences (bool, optional): Maintain maximum influences
            normalize_weights (bool, optional): Normalize weights
            post_normalize (bool, optional): Post normalize weights
            use_metahuman_weights (bool, optional): Use MetaHuman weights if available
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if mesh and skeleton exist
            if not cmds.objExists(mesh):
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            if not cmds.objExists(skeleton):
                cmds.warning(f"Skeleton not found: {skeleton}")
                return False
            
            # Get all joints in the skeleton
            joints = self._get_all_joints(skeleton)
            
            if not joints:
                cmds.warning(f"Skeleton {skeleton} has no joints")
                return False
            
            # Check if mesh is already bound
            skin_clusters = self._get_skin_clusters(mesh)
            if skin_clusters:
                # Ask if we should unbind first
                result = cmds.confirmDialog(
                    title="Mesh already bound",
                    message=f"Mesh {mesh} is already bound to skeleton. Do you want to unbind first?",
                    button=["Yes", "No"],
                    defaultButton="Yes",
                    cancelButton="No",
                    dismissString="No"
                )
                
                if result == "Yes":
                    self.unbind_mesh(mesh)
                else:
                    cmds.warning(f"Mesh {mesh} is already bound to skeleton")
                    return False
            
            # Bind mesh to skeleton
            skin_cluster = None
            
            if method == "Smooth Bind":
                skin_cluster = cmds.skinCluster(
                    joints,
                    mesh,
                    name=f"{mesh}_skinCluster",
                    toSelectedBones=True,
                    bindMethod=0,  # Closest distance
                    skinMethod=0,  # Linear
                    normalizeWeights=2 if normalize_weights else 0,  # 0=None, 1=Interactive, 2=Post
                    maximumInfluences=max_influences,
                    obeyMaxInfluences=maintain_max_influences,
                    dropoffRate=4.0
                )[0]
            elif method == "Rigid Bind":
                skin_cluster = cmds.skinCluster(
                    joints,
                    mesh,
                    name=f"{mesh}_skinCluster",
                    toSelectedBones=True,
                    bindMethod=1,  # Closest joint
                    skinMethod=0,  # Linear
                    normalizeWeights=2 if normalize_weights else 0,
                    maximumInfluences=1,
                    obeyMaxInfluences=True,
                    dropoffRate=4.0
                )[0]
            elif method == "Heat Map Bind":
                # Heat Map binding requires Maya 2020 or later
                skin_cluster = cmds.skinCluster(
                    joints,
                    mesh,
                    name=f"{mesh}_skinCluster",
                    toSelectedBones=True,
                    bindMethod=0,  # Closest distance
                    skinMethod=1,  # Dual quaternion
                    normalizeWeights=2 if normalize_weights else 0,
                    maximumInfluences=max_influences,
                    obeyMaxInfluences=maintain_max_influences,
                    dropoffRate=4.0,
                    heatmapFalloff=0.68
                )[0]
            elif method == "Geodesic Voxel Bind":
                # Geodesic Voxel binding requires Maya 2018 or later
                skin_cluster = cmds.skinCluster(
                    joints,
                    mesh,
                    name=f"{mesh}_skinCluster",
                    toSelectedBones=True,
                    bindMethod=2,  # Geodesic Voxel
                    skinMethod=0,  # Linear
                    normalizeWeights=2 if normalize_weights else 0,
                    maximumInfluences=max_influences,
                    obeyMaxInfluences=maintain_max_influences,
                    dropoffRate=4.0,
                    volumeType=1,  # 0=Volume, 1=Surface
                    volumeBinding=0.5
                )[0]
            
            if not skin_cluster:
                cmds.warning(f"Failed to create skin cluster")
                return False
            
            # Post-normalize weights if needed
            if post_normalize and normalize_weights:
                cmds.skinPercent(skin_cluster, mesh, normalize=True)
            
            # Load MetaHuman weights if available and requested
            if use_metahuman_weights:
                weights_file = os.path.join(config.METAHUMAN_PATH, "weights", f"{mesh}_weights.xml")
                if os.path.exists(weights_file):
                    self.import_weights(mesh, weights_file)
            
            print(f"Successfully bound {mesh} to {skeleton}")
            return True
        except Exception as e:
            cmds.warning(f"Error binding mesh: {str(e)}")
            return False
    
    def unbind_mesh(self, mesh):
        """
        Unbind a mesh from its skeleton
        
        Args:
            mesh (str): Name of the mesh
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if mesh exists
            if not cmds.objExists(mesh):
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            # Get skin clusters
            skin_clusters = self._get_skin_clusters(mesh)
            
            if not skin_clusters:
                print(f"Mesh {mesh} is not bound to skeleton")
                return True
            
            # Unbind mesh
            for skin_cluster in skin_clusters:
                cmds.skinCluster(skin_cluster, edit=True, unbind=True)
            
            print(f"Successfully unbound {mesh}")
            return True
        except Exception as e:
            cmds.warning(f"Error unbinding mesh: {str(e)}")
            return False
    
    def export_weights(self, mesh, file_path):
        """
        Export skin weights to a file
        
        Args:
            mesh (str): Name of the mesh
            file_path (str): Path to save the weights file
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if mesh exists
            if not cmds.objExists(mesh):
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            # Get skin clusters
            skin_clusters = self._get_skin_clusters(mesh)
            
            if not skin_clusters:
                cmds.warning(f"Mesh {mesh} is not bound to skeleton")
                return False
            
            # Determine file format
            if file_path.lower().endswith(".xml"):
                # Export as XML
                cmds.deformerWeights(
                    file_path,
                    export=True,
                    deformer=skin_clusters[0],
                    format="XML"
                )
            else:
                # Export as JSON
                weights_data = self._get_skin_weights(mesh, skin_clusters[0])
                
                with open(file_path, "w") as f:
                    json.dump(weights_data, f, indent=2)
            
            print(f"Successfully exported weights for {mesh} to {file_path}")
            return True
        except Exception as e:
            cmds.warning(f"Error exporting weights: {str(e)}")
            return False
    
    def import_weights(self, mesh, file_path):
        """
        Import skin weights from a file
        
        Args:
            mesh (str): Name of the mesh
            file_path (str): Path to the weights file
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if mesh exists
            if not cmds.objExists(mesh):
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            # Check if file exists
            if not os.path.exists(file_path):
                cmds.warning(f"Weights file not found: {file_path}")
                return False
            
            # Get skin clusters
            skin_clusters = self._get_skin_clusters(mesh)
            
            if not skin_clusters:
                cmds.warning(f"Mesh {mesh} is not bound to skeleton")
                return False
            
            # Determine file format
            if file_path.lower().endswith(".xml"):
                # Import from XML
                cmds.deformerWeights(
                    file_path,
                    im=True,
                    deformer=skin_clusters[0],
                    format="XML"
                )
            else:
                # Import from JSON
                with open(file_path, "r") as f:
                    weights_data = json.load(f)
                
                self._set_skin_weights(mesh, skin_clusters[0], weights_data)
            
            print(f"Successfully imported weights to {mesh}")
            return True
        except Exception as e:
            cmds.warning(f"Error importing weights: {str(e)}")
            return False
    
    def mirror_weights(self, mesh):
        """
        Mirror skin weights across the X axis
        
        Args:
            mesh (str): Name of the mesh
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if mesh exists
            if not cmds.objExists(mesh):
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            # Get skin clusters
            skin_clusters = self._get_skin_clusters(mesh)
            
            if not skin_clusters:
                cmds.warning(f"Mesh {mesh} is not bound to skeleton")
                return False
            
            # Mirror weights
            cmds.copySkinWeights(
                sourceSkin=skin_clusters[0],
                destinationSkin=skin_clusters[0],
                mirrorMode="YZ",
                surfaceAssociation="closestPoint",
                influenceAssociation=["label", "oneToOne"]
            )
            
            print(f"Successfully mirrored weights for {mesh}")
            return True
        except Exception as e:
            cmds.warning(f"Error mirroring weights: {str(e)}")
            return False
    
    def prune_weights(self, mesh, threshold=0.01):
        """
        Prune small skin weights
        
        Args:
            mesh (str): Name of the mesh
            threshold (float, optional): Threshold for pruning
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if mesh exists
            if not cmds.objExists(mesh):
                cmds.warning(f"Mesh model not found: {mesh}")
                return False
            
            # Get skin clusters
            skin_clusters = self._get_skin_clusters(mesh)
            
            if not skin_clusters:
                cmds.warning(f"Mesh {mesh} is not bound to skeleton")
                return False
            
            # Prune weights
            cmds.skinPercent(
                skin_clusters[0],
                mesh,
                pruneWeights=threshold
            )
            
            print(f"Successfully pruned weights for {mesh}")
            return True
        except Exception as e:
            cmds.warning(f"Error pruning weights: {str(e)}")
            return False
    
    def transfer_weights(self, source_mesh, target_mesh):
        """
        Transfer skin weights from one mesh to another
        
        Args:
            source_mesh (str): Name of the source mesh
            target_mesh (str): Name of the target mesh
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if meshes exist
            if not cmds.objExists(source_mesh):
                cmds.warning(f"Source mesh model not found: {source_mesh}")
                return False
            
            if not cmds.objExists(target_mesh):
                cmds.warning(f"Target mesh model not found: {target_mesh}")
                return False
            
            # Get source skin cluster
            source_skin_clusters = self._get_skin_clusters(source_mesh)
            
            if not source_skin_clusters:
                cmds.warning(f"Source mesh {source_mesh} is not bound to skeleton")
                return False
            
            # Get target skin cluster
            target_skin_clusters = self._get_skin_clusters(target_mesh)
            
            if not target_skin_clusters:
                # Get joints from source skin cluster
                joints = cmds.skinCluster(source_skin_clusters[0], query=True, influence=True)
                
                # Bind target mesh to the same joints
                target_skin_clusters = [cmds.skinCluster(
                    joints,
                    target_mesh,
                    name=f"{target_mesh}_skinCluster",
                    toSelectedBones=True,
                    bindMethod=0,
                    skinMethod=0,
                    normalizeWeights=2,
                    maximumInfluences=4,
                    obeyMaxInfluences=True,
                    dropoffRate=4.0
                )[0]]
            
            # Transfer weights
            cmds.copySkinWeights(
                sourceSkin=source_skin_clusters[0],
                destinationSkin=target_skin_clusters[0],
                noMirror=True,
                surfaceAssociation="closestPoint",
                influenceAssociation=["label", "name"]
            )
            
            print(f"Successfully transferred weights from {source_mesh} to {target_mesh}")
            return True
        except Exception as e:
            cmds.warning(f"Error transferring weights: {str(e)}")
            return False
    
    def _get_all_joints(self, root_joint):
        """
        Get all joints in a hierarchy
        
        Args:
            root_joint (str): Name of the root joint
            
        Returns:
            list: List of joint names
        """
        try:
            # Check if root joint exists
            if not cmds.objExists(root_joint):
                return []
            
            # Get all joints in the hierarchy
            joints = cmds.listRelatives(root_joint, allDescendents=True, type="joint") or []
            joints.append(root_joint)
            
            return joints
        except Exception as e:
            cmds.warning(f"Error getting joints: {str(e)}")
            return []
    
    def _get_skin_clusters(self, mesh):
        """
        Get skin clusters for a mesh
        
        Args:
            mesh (str): Name of the mesh
            
        Returns:
            list: List of skin cluster names
        """
        try:
            # Check if mesh exists
            if not cmds.objExists(mesh):
                return []
            
            # Get mesh shape
            shapes = cmds.listRelatives(mesh, shapes=True, type="mesh")
            if not shapes:
                return []
            
            mesh_shape = shapes[0]
            
            # Get skin clusters
            skin_clusters = []
            
            # Get all deformers
            deformers = cmds.listHistory(mesh_shape, pruneDagObjects=True)
            
            if deformers:
                # Filter for skin clusters
                skin_clusters = cmds.ls(deformers, type="skinCluster")
            
            return skin_clusters
        except Exception as e:
            cmds.warning(f"Error getting skin clusters: {str(e)}")
            return []
    
    def _get_skin_weights(self, mesh, skin_cluster):
        """
        Get skin weights for a mesh
        
        Args:
            mesh (str): Name of the mesh
            skin_cluster (str): Name of the skin cluster
            
        Returns:
            dict: Dictionary of skin weights
        """
        try:
            # Check if mesh and skin cluster exist
            if not cmds.objExists(mesh) or not cmds.objExists(skin_cluster):
                return {}
            
            # Get influences
            influences = cmds.skinCluster(skin_cluster, query=True, influence=True)
            
            if not influences:
                return {}
            
            # Get vertex count
            vertex_count = cmds.polyEvaluate(mesh, vertex=True)
            
            # Get weights for each vertex
            weights_data = {
                "influences": influences,
                "weights": {}
            }
            
            for i in range(vertex_count):
                vertex_weights = {}
                
                # Get weights for this vertex
                for influence in influences:
                    weight = cmds.skinPercent(
                        skin_cluster,
                        f"{mesh}.vtx[{i}]",
                        query=True,
                        transform=influence
                    )
                    
                    if weight > 0.0001:  # Only store non-zero weights
                        vertex_weights[influence] = weight
                
                weights_data["weights"][i] = vertex_weights
            
            return weights_data
        except Exception as e:
            cmds.warning(f"Error getting skin weights: {str(e)}")
            return {}
    
    def _set_skin_weights(self, mesh, skin_cluster, weights_data):
        """
        Set skin weights for a mesh
        
        Args:
            mesh (str): Name of the mesh
            skin_cluster (str): Name of the skin cluster
            weights_data (dict): Dictionary of skin weights
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if mesh and skin cluster exist
            if not cmds.objExists(mesh) or not cmds.objExists(skin_cluster):
                return False
            
            # Get influences from skin cluster
            current_influences = cmds.skinCluster(skin_cluster, query=True, influence=True)
            
            # Get influences from weights data
            data_influences = weights_data.get("influences", [])
            
            # Map influences if needed
            influence_map = {}
            
            for data_influence in data_influences:
                # Check if influence exists
                if cmds.objExists(data_influence) and data_influence in current_influences:
                    influence_map[data_influence] = data_influence
                else:
                    # Try to find a matching influence by name
                    for current_influence in current_influences:
                        if current_influence.endswith(data_influence.split("|")[-1]):
                            influence_map[data_influence] = current_influence
                            break
            
            # Set weights for each vertex
            weights = weights_data.get("weights", {})
            
            for vertex_index, vertex_weights in weights.items():
                # Convert to transformation value pairs
                transform_value_pairs = []
                
                for influence, weight in vertex_weights.items():
                    # Map influence if needed
                    mapped_influence = influence_map.get(influence)
                    
                    if mapped_influence:
                        transform_value_pairs.extend([mapped_influence, weight])
                
                if transform_value_pairs:
                    # Set weights for this vertex
                    cmds.skinPercent(
                        skin_cluster,
                        f"{mesh}.vtx[{vertex_index}]",
                        transformValue=transform_value_pairs
                    )
            
            # Normalize weights
            cmds.skinPercent(skin_cluster, mesh, normalize=True)
            
            return True
        except Exception as e:
            cmds.warning(f"Error setting skin weights: {str(e)}")
            return False