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

"""
Tool - Maya Utilities
Provides utilities for working with Maya
"""

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

# Import configuration
import config

class MayaUtils:
    """Utilities for working with Maya"""
    
    def __init__(self):
        """Initialize Maya utilities"""
        # Ensure required plugins are loaded
        self._load_required_plugins()
    
    def _load_required_plugins(self):
        """Load required Maya plugins"""
        required_plugins = ["fbxmaya", "matrixNodes"]
        
        for plugin in required_plugins:
            if not cmds.pluginInfo(plugin, query=True, loaded=True):
                try:
                    cmds.loadPlugin(plugin)
                    print(f"Loaded plugin: {plugin}")
                except Exception as e:
                    cmds.warning(f"Error loading plugin: {plugin}, Error: {str(e)}")
    
    def export_fbx(self, file_path):
        """
        Export the current scene or selection as an FBX file
        
        Args:
            file_path (str): Path to save the FBX file
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Ensure FBX plugin is loaded
            if not cmds.pluginInfo("fbxmaya", query=True, loaded=True):
                cmds.loadPlugin("fbxmaya")
            
            # Get selection
            selection = cmds.ls(selection=True)
            
            # Set FBX export options
            mel.eval("FBXResetExport")
            
            # Set export options
            mel.eval("FBXExportFileVersion -v FBX201900")
            mel.eval("FBXExportUpAxis y")
            mel.eval("FBXExportShapes -v true")
            mel.eval("FBXExportScaleFactor 1.0")
            mel.eval("FBXExportInAscii -v false")
            mel.eval("FBXExportConstraints -v false")
            mel.eval("FBXExportLights -v false")
            mel.eval("FBXExportSkins -v true")
            mel.eval("FBXExportSmoothingGroups -v true")
            mel.eval("FBXExportSmoothMesh -v true")
            mel.eval("FBXExportEmbeddedTextures -v false")
            mel.eval("FBXExportCameras -v false")
            mel.eval("FBXExportBakeComplexAnimation -v true")
            mel.eval("FBXExportBakeComplexStart -v 0")
            mel.eval("FBXExportBakeComplexEnd -v 100")
            mel.eval("FBXExportBakeComplexStep -v 1")
            mel.eval("FBXExportBakeResampleAnimation -v true")
            mel.eval("FBXExportAnimationOnly -v false")
            mel.eval("FBXExportQuaternion -v euler")
            mel.eval("FBXExportAxisConversionMethod \"none\"")
            mel.eval("FBXExportApplyConstantKeyReducer -v false")
            mel.eval("FBXExportUseSceneName -v false")
            mel.eval("FBXExportGenerateLog -v false")
            
            # Export FBX
            if selection:
                # Export selected objects
                result = mel.eval(f'FBXExport -f "{file_path}" -s')
            else:
                # Export entire scene
                result = mel.eval(f'FBXExport -f "{file_path}"')
            
            if result:
                print(f"Successfully exported FBX file: {file_path}")
                return True
            else:
                cmds.warning(f"Failed to export FBX file: {file_path}")
                return False
        except Exception as e:
            cmds.warning(f"Error exporting FBX file: {str(e)}")
            return False
    
    def import_fbx(self, file_path):
        """
        Import an FBX file into the current scene
        
        Args:
            file_path (str): Path to the FBX file
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Ensure FBX plugin is loaded
            if not cmds.pluginInfo("fbxmaya", query=True, loaded=True):
                cmds.loadPlugin("fbxmaya")
            
            # Set FBX import options
            mel.eval("FBXResetImport")
            
            # Set import options
            mel.eval("FBXImportMode -v add")
            mel.eval("FBXImportMergeAnimationLayers -v false")
            mel.eval("FBXImportProtectDrivenKeys -v true")
            mel.eval("FBXImportConvertDeformingNullsToJoint -v false")
            mel.eval("FBXImportMergeBackNullPivots -v false")
            mel.eval("FBXImportSetLockedAttribute -v true")
            mel.eval("FBXImportConstraints -v false")
            mel.eval("FBXImportCameras -v false")
            mel.eval("FBXImportLights -v false")
            mel.eval("FBXImportSkins -v true")
            mel.eval("FBXImportShapes -v true")
            mel.eval("FBXImportUnlockNormals -v false")
            mel.eval("FBXImportForcedFileAxis \"y\"")
            
            # Import FBX
            result = mel.eval(f'FBXImport -f "{file_path}"')
            
            if result:
                print(f"Successfully imported FBX file: {file_path}")
                return True
            else:
                cmds.warning(f"Failed to import FBX file: {file_path}")
                return False
        except Exception as e:
            cmds.warning(f"Error importing FBX file: {str(e)}")
            return False
    
    def save_scene(self, file_path):
        """
        Save the current scene to a Maya file
        
        Args:
            file_path (str): Path to save the Maya file
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Determine file type
            if file_path.lower().endswith(".ma"):
                file_type = "mayaAscii"
            else:
                file_type = "mayaBinary"
            
            # Save scene
            cmds.file(rename=file_path)
            cmds.file(save=True, type=file_type)
            
            print(f"Successfully saved scene: {file_path}")
            return True
        except Exception as e:
            cmds.warning(f"Error saving scene: {str(e)}")
            return False
    
    def load_scene(self, file_path):
        """
        Load a Maya scene file
        
        Args:
            file_path (str): Path to the Maya file
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Check if current scene has changes
            if cmds.file(query=True, modified=True):
                # Ask user to save changes
                result = cmds.confirmDialog(
                    title="Save Changes",
                    message="Current scene has unsaved changes. Save?",
                    button=["Save", "Don't Save", "Cancel"],
                    defaultButton="Save",
                    cancelButton="Cancel",
                    dismissString="Cancel"
                )
                
                if result == "Save":
                    cmds.file(save=True)
                elif result == "Don't Save":
                    return False
            
            # Load scene
            cmds.file(file_path, open=True, force=True)
            
            print(f"Successfully loaded scene: {file_path}")
            return True
        except Exception as e:
            cmds.warning(f"Error loading scene: {str(e)}")
            return False
    
    def create_reference(self, file_path, namespace=None):
        """
        Create a reference to a Maya file
        
        Args:
            file_path (str): Path to the Maya file
            namespace (str, optional): Namespace for the reference
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Create reference
            if namespace:
                cmds.file(file_path, reference=True, namespace=namespace)
            else:
                cmds.file(file_path, reference=True)
            
            print(f"Successfully created reference: {file_path}")
            return True
        except Exception as e:
            cmds.warning(f"Error creating reference: {str(e)}")
            return False
    
    def remove_reference(self, reference_node):
        """
        Remove a reference
        
        Args:
            reference_node (str): Name of the reference node
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Remove reference
            cmds.file(referenceNode=reference_node, removeReference=True)
            
            print(f"Successfully removed reference: {reference_node}")
            return True
        except Exception as e:
            cmds.warning(f"Error removing reference: {str(e)}")
            return False
    
    def get_scene_info(self):
        """
        Get information about the current scene
        
        Returns:
            dict: Dictionary containing scene information
        """
        try:
            info = {}
            
            # Get scene path
            scene_path = cmds.file(query=True, sceneName=True)
            info["scene_path"] = scene_path
            
            # Get scene name
            scene_name = os.path.basename(scene_path) if scene_path else "untitled"
            info["scene_name"] = scene_name
            
            # Get scene statistics
            info["mesh_count"] = len(cmds.ls(type="mesh"))
            info["joint_count"] = len(cmds.ls(type="joint"))
            info["blendshape_count"] = len(cmds.ls(type="blendShape"))
            info["skincluster_count"] = len(cmds.ls(type="skinCluster"))
            
            # Get selection
            info["selection"] = cmds.ls(selection=True)
            
            return info
        except Exception as e:
            cmds.warning(f"Error getting scene information: {str(e)}")
            return {}
    
    def create_locator(self, name, position=(0, 0, 0)):
        """
        Create a locator
        
        Args:
            name (str): Name for the locator
            position (tuple, optional): Position for the locator
            
        Returns:
            str: Name of the created locator
        """
        try:
            # Create locator
            locator = cmds.spaceLocator(name=name)[0]
            
            # Set position
            cmds.setAttr(f"{locator}.translate", *position)
            
            return locator
        except Exception as e:
            cmds.warning(f"Error creating locator: {str(e)}")
            return None
    
    def create_group(self, name, children=None):
        """
        Create a group
        
        Args:
            name (str): Name for the group
            children (list, optional): List of objects to parent to the group
            
        Returns:
            str: Name of the created group
        """
        try:
            # Create group
            group = cmds.group(empty=True, name=name)
            
            # Parent children
            if children:
                cmds.parent(children, group)
            
            return group
        except Exception as e:
            cmds.warning(f"Error creating group: {str(e)}")
            return None
    
    def create_control(self, name, shape="circle", radius=1.0, color=None):
        """
        Create a control curve
        
        Args:
            name (str): Name for the control
            shape (str, optional): Shape of the control
            radius (float, optional): Radius of the control
            color (int, optional): Color index for the control
            
        Returns:
            str: Name of the created control
        """
        try:
            # Create control based on shape
            control = None
            
            if shape == "circle":
                control = cmds.circle(name=name, radius=radius, normal=(0, 1, 0), constructionHistory=False)[0]
            elif shape == "square":
                points = [
                    (-radius, 0, -radius),
                    (radius, 0, -radius),
                    (radius, 0, radius),
                    (-radius, 0, radius),
                    (-radius, 0, -radius)
                ]
                control = cmds.curve(name=name, degree=1, point=points)
            elif shape == "cube":
                points = [
                    (-radius, -radius, -radius),
                    (radius, -radius, -radius),
                    (radius, -radius, radius),
                    (-radius, -radius, radius),
                    (-radius, -radius, -radius),
                    (-radius, radius, -radius),
                    (radius, radius, -radius),
                    (radius, -radius, -radius),
                    (radius, radius, -radius),
                    (radius, radius, radius),
                    (radius, -radius, radius),
                    (radius, radius, radius),
                    (-radius, radius, radius),
                    (-radius, -radius, radius),
                    (-radius, radius, radius),
                    (-radius, radius, -radius)
                ]
                control = cmds.curve(name=name, degree=1, point=points)
            elif shape == "sphere":
                # Create circles in different planes
                circle1 = cmds.circle(radius=radius, normal=(1, 0, 0), constructionHistory=False)[0]
                circle2 = cmds.circle(radius=radius, normal=(0, 1, 0), constructionHistory=False)[0]
                circle3 = cmds.circle(radius=radius, normal=(0, 0, 1), constructionHistory=False)[0]
                
                # Combine circles
                control = cmds.group(empty=True, name=name)
                cmds.parent(
                    cmds.listRelatives(circle1, shapes=True),
                    cmds.listRelatives(circle2, shapes=True),
                    cmds.listRelatives(circle3, shapes=True),
                    control,
                    shape=True,
                    relative=True
                )
                
                # Delete original circles
                cmds.delete(circle1, circle2, circle3)
            else:
                # Default to circle
                control = cmds.circle(name=name, radius=radius, normal=(0, 1, 0), constructionHistory=False)[0]
            
            # Set color
            if color is not None:
                shapes = cmds.listRelatives(control, shapes=True)
                for shape in shapes:
                    cmds.setAttr(f"{shape}.overrideEnabled", 1)
                    cmds.setAttr(f"{shape}.overrideColor", color)
            
            return control
        except Exception as e:
            cmds.warning(f"Error creating control: {str(e)}")
            return None
    
    def create_joint(self, name, position=(0, 0, 0), parent=None):
        """
        Create a joint
        
        Args:
            name (str): Name for the joint
            position (tuple, optional): Position for the joint
            parent (str, optional): Parent for the joint
            
        Returns:
            str: Name of the created joint
        """
        try:
            # Select parent if provided
            if parent:
                cmds.select(parent)
            else:
                cmds.select(clear=True)
            
            # Create joint
            joint = cmds.joint(name=name, position=position)
            
            return joint
        except Exception as e:
            cmds.warning(f"Error creating joint: {str(e)}")
            return None
    
    def set_keyframe(self, obj, attribute, value, time):
        """
        Set a keyframe
        
        Args:
            obj (str): Name of the object
            attribute (str): Name of the attribute
            value (float): Value to set
            time (float): Time to set the keyframe
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Set attribute value
            cmds.setAttr(f"{obj}.{attribute}", value)
            
            # Set keyframe
            cmds.setKeyframe(obj, attribute=attribute, time=time)
            
            return True
        except Exception as e:
            cmds.warning(f"Error setting keyframe: {str(e)}")
            return False
    
    def bake_animation(self, objects, start_frame, end_frame, step=1):
        """
        Bake animation
        
        Args:
            objects (list): List of objects to bake
            start_frame (float): Start frame
            end_frame (float): End frame
            step (float, optional): Frame step
            
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Bake animation
            cmds.bakeResults(
                objects,
                simulation=True,
                time=(start_frame, end_frame),
                sampleBy=step,
                disableImplicitControl=True,
                preserveOutsideKeys=True,
                sparseAnimCurveBake=False,
                removeBakedAttributeFromLayer=False,
                bakeOnOverrideLayer=False,
                minimizeRotation=True,
                controlPoints=False
            )
            
            return True
        except Exception as e:
            cmds.warning(f"Error baking animation: {str(e)}")
            return False
    
    def clean_scene(self):
        """
        Clean the scene by removing unused nodes
        
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Delete unused nodes
            mel.eval("MLdeleteUnused")
            
            # Delete turtle nodes if they exist
            if cmds.pluginInfo("Turtle", query=True, loaded=True):
                turtle_nodes = cmds.ls("*turtle*")
                if turtle_nodes:
                    cmds.delete(turtle_nodes)
            
            # Delete unknown nodes
            unknown_nodes = cmds.ls(type="unknown")
            if unknown_nodes:
                cmds.delete(unknown_nodes)
            
            # Delete unused animation curves
            unused_anim_curves = []
            for curve in cmds.ls(type=["animCurveTA", "animCurveTL", "animCurveTU"]):
                if not cmds.connectionInfo(f"{curve}.output", destinationFromSource=True):
                    unused_anim_curves.append(curve)
            
            if unused_anim_curves:
                cmds.delete(unused_anim_curves)
            
            # Delete empty groups
            empty_groups = []
            for transform in cmds.ls(type="transform"):
                if not cmds.listRelatives(transform, children=True) and not cmds.listRelatives(transform, shapes=True):
                    empty_groups.append(transform)
            
            if empty_groups:
                cmds.delete(empty_groups)
            
            # Optimize scene
            mel.eval("OptimizeScene")
            
            return True
        except Exception as e:
            cmds.warning(f"Error cleaning scene: {str(e)}")
            return False