import traceback
import bpy
from typing import Dict, List, Set, Optional, Tuple, Any
from bpy.types import Operator, Context, Object, PoseBone, EditBone, Bone, Constraint
from ...core.common import get_active_armature, transfer_vertex_weights, get_all_meshes
from ...core.logging_setup import logger
from ...core.translations import t
from ...core.dictionaries import rigify_unity_names, rigify_basic_unity_names, rigify_unnecessary_bones
from ...core.armature_validation import validate_armature
import traceback

class AvatarToolkit_OT_ConvertRigifyToUnity(Operator):
    """Convert Rigify armature to Unity-compatible format"""
    bl_idname = "avatar_toolkit.convert_rigify_to_unity"
    bl_label = t("Tools.convert_rigify_to_unity")
    bl_description = t("Tools.convert_rigify_to_unity_desc")
    bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context: Context) -> bool:
        armature = get_active_armature(context)
        if not armature:
            return False
        return ("DEF-spine" in armature.data.bones or 
                "spine" in armature.data.bones and "metarig" in armature.name.lower())

    def execute(self, context: Context) -> Set[str]:
        try:
            logger.info("Starting Rigify to Unity conversion")
            armature = get_active_armature(context)
            if not armature:
                logger.error("No armature found")
                self.report({'ERROR'}, t("Tools.no_armature"))
                return {'CANCELLED'}

            logger.debug(f"Converting armature: {armature.name}")
            armature.name = "Armature"
            armature.data.name = "Armature"
            logger.debug("Renamed armature to 'Armature'")
            
            if "DEF-spine" in armature.data.bones:
                logger.info("Processing DEF bones")
                self.move_def_bones(armature)
                self.rename_bones_for_unity(armature)
            else:
                logger.info("Processing basic bones")
                self.cleanup_extra_bones(armature)
                self.rename_basic_bones_for_unity(armature)
            
            logger.debug("Cleaning up bone collections")
            self.cleanup_bone_collections(armature)
            
            if context.scene.avatar_toolkit.merge_twist_bones:
                logger.info("Merging twist bones")
                self.handle_twist_bones(armature)
                
            logger.info("Successfully converted Rigify armature to Unity format")
            self.report({'INFO'}, t("Tools.rigify_converted"))
            return {'FINISHED'}
            
        except Exception:
            logger.error(f"Failed to convert Rigify: {traceback.format_exc()}", exc_info=True)
            self.report({'ERROR'}, traceback.format_exc())
            return {'CANCELLED'}

    def cleanup_extra_bones(self, armature: Object) -> None:
        """Remove unnecessary bones and merge neck bones"""
        logger.debug("Starting cleanup of extra bones")
        
        # Set armature as active object before mode switch
        bpy.context.view_layer.objects.active = armature
        
        # Get all meshes for weight transfer
        meshes = get_all_meshes(bpy.context)
        
        bpy.ops.object.mode_set(mode='EDIT')
        
        bones_to_remove: List[str] = []
        for bone in armature.data.edit_bones:
            bone_name_lower = bone.name.lower()
            if any(bone_name_lower.startswith(pattern) or bone_name_lower == pattern 
                   for pattern in rigify_unnecessary_bones):
                bones_to_remove.append(bone.name)
        
        # Check for neck bones that need merging
        merge_neck_bones = 'spine.004' in armature.data.edit_bones and 'spine.005' in armature.data.edit_bones
        
        bpy.ops.object.mode_set(mode='OBJECT')
        
        # Transfer weights from bones being removed
        for bone_name in bones_to_remove:
            if bone_name in armature.data.bones:
                logger.debug(f"Transferring weights from bone: {bone_name}")
                for mesh in meshes:
                    if bone_name in mesh.vertex_groups:
                        # Remove the vertex group since we don't need the weights
                        mesh.vertex_groups.remove(mesh.vertex_groups[bone_name])
        
        # Transfer weights for neck bone merging
        if merge_neck_bones:
            logger.debug("Transferring weights from spine.005 to spine.004")
            for mesh in meshes:
                if 'spine.005' in mesh.vertex_groups:
                    transfer_vertex_weights(mesh, 'spine.005', 'spine.004')
        
        bpy.ops.object.mode_set(mode='EDIT')
        
        # Remove unnecessary bones
        for bone_name in bones_to_remove:
            if bone_name in armature.data.edit_bones:
                logger.debug(f"Removing bone: {bone_name}")
                armature.data.edit_bones.remove(armature.data.edit_bones[bone_name])
   
        # Merge neck bones
        if merge_neck_bones:
            logger.debug("Merging neck bones")
            neck_start = armature.data.edit_bones['spine.004']
            neck_end = armature.data.edit_bones['spine.005']
            neck_start.tail = neck_end.tail
            armature.data.edit_bones.remove(neck_end)
            neck_start.name = "Neck"
            
        # Rename head bone
        if 'spine.006' in armature.data.edit_bones:
            logger.debug("Renaming head bone")
            head_bone = armature.data.edit_bones['spine.006']
            head_bone.name = "Head"

    def move_def_bones(self, armature: Object) -> None:
        """Move DEF bones to their correct positions"""
        logger.debug("Moving DEF bones to correct positions")
        
        # Set armature as active object
        bpy.context.view_layer.objects.active = armature
        remap: Dict[str, str] = self.get_org_remap(armature)
        remap.update(self.get_special_remap())

        remove_bones_in_chain: List[str] = [
            'DEF-upper_arm.L.001', 'DEF-forearm.L.001',
            'DEF-upper_arm.R.001', 'DEF-forearm.R.001',
            'DEF-thigh.L.001', 'DEF-shin.L.001',
            'DEF-thigh.R.001', 'DEF-shin.R.001'
        ]

        transform_copies: List[str] = self.get_transform_copies(armature)

        logger.debug("Setting up transform copies")
        bpy.ops.object.mode_set(mode='POSE')
        for bone_name in transform_copies:
            bone = armature.pose.bones[bone_name]
            org_name = 'ORG-' + self.get_proto_name(bone_name)
            if org_name in armature.pose.bones:
                constraint = bone.constraints.new('COPY_TRANSFORMS')
                constraint.target = armature
                constraint.subtarget = org_name
                constr_count = len(bone.constraints)
                if constr_count > 1:
                    bone.constraints.move(constr_count-1, 0)

        logger.debug("Remapping bone parents")
        bpy.ops.object.mode_set(mode='EDIT')
        for remap_key in remap:
            if remap_key in armature.data.edit_bones and remap[remap_key] in armature.data.edit_bones:
                armature.data.edit_bones[remap_key].parent = armature.data.edit_bones[remap[remap_key]]

        logger.debug("Processing bone chain removal")
        bpy.ops.object.mode_set(mode='OBJECT')
        for bone_name in remove_bones_in_chain:
            if bone_name in armature.data.bones:
                armature.data.bones[bone_name].use_deform = False

        # Get all meshes for weight transfer
        meshes = get_all_meshes(bpy.context)
        
        bpy.ops.object.mode_set(mode='OBJECT')
        for bone_name in remove_bones_in_chain:
            if bone_name in armature.data.bones:
                parent_name = armature.data.bones[bone_name].parent.name if armature.data.bones[bone_name].parent else None
                if parent_name:
                    logger.debug(f"Transferring weights from {bone_name} to {parent_name}")
                    for mesh in meshes:
                        if bone_name in mesh.vertex_groups and parent_name in mesh.vertex_groups:
                            transfer_vertex_weights(mesh, bone_name, parent_name)
                        elif bone_name in mesh.vertex_groups:
                            # Remove weights if no parent to merge to
                            mesh.vertex_groups.remove(mesh.vertex_groups[bone_name])

        bpy.ops.object.mode_set(mode='EDIT')
        for bone_name in remove_bones_in_chain:
            if bone_name in armature.data.bones:
                remove_bone = armature.data.edit_bones[bone_name]
                parent_bone = remove_bone.parent
                parent_bone.tail = remove_bone.tail
                retarget_bones = list(remove_bone.children)
                for bone in retarget_bones:
                    bone.parent = parent_bone
                armature.data.edit_bones.remove(remove_bone)

    def rename_bones_for_unity(self, armature: Object) -> None:
        """Rename bones to Unity-compatible names"""
        logger.debug("Renaming bones to Unity format")
        for old_name, new_name in rigify_unity_names.items():
            bone = armature.pose.bones.get(old_name)
            if bone:
                logger.debug(f"Renaming bone: {old_name} -> {new_name}")
                bone.name = new_name

    def rename_basic_bones_for_unity(self, armature: Object) -> None:
        """Rename basic metarig bones to Unity-compatible names"""
        logger.debug("Renaming basic metarig bones")
        for old_name, new_name in rigify_basic_unity_names.items():
            bone = armature.pose.bones.get(old_name)
            if bone:
                logger.debug(f"Renaming basic bone: {old_name} -> {new_name}")
                bone.name = new_name

    def cleanup_bone_collections(self, armature: Object) -> None:
        """Remove all bone collections since they're not needed for Unity"""
        logger.debug("Cleaning up bone collections")
        if hasattr(armature.data, 'collections') and armature.data.collections:
            while len(armature.data.collections) > 0:
                collection = armature.data.collections[0]
                armature.data.collections.remove(collection)

            while len(armature.data.collections) > 1:
                collection = armature.data.collections[1]
                armature.data.collections.remove(collection)

    def handle_twist_bones(self, armature: Object) -> None:
        """Handle twist bones during conversion"""
        logger.debug("Processing twist bones")
        twist_bones: List[Tuple[str, str]] = [
            ("DEF-upper_arm_twist.L", "DEF-upper_arm.L"),
            ("DEF-upper_arm_twist.R", "DEF-upper_arm.R"),
            ("DEF-forearm_twist.L", "DEF-forearm.L"),
            ("DEF-forearm_twist.R", "DEF-forearm.R"),
            ("DEF-thigh_twist.L", "DEF-thigh.L"),
            ("DEF-thigh_twist.R", "DEF-thigh.R")
        ]

        # Get all meshes for weight transfer
        meshes = get_all_meshes(bpy.context)
        
        bpy.ops.object.mode_set(mode='OBJECT')
        for twist_bone, parent_bone in twist_bones:
            if twist_bone in armature.data.bones and parent_bone in armature.data.bones:
                logger.debug(f"Transferring weights from {twist_bone} to {parent_bone}")
                for mesh in meshes:
                    if twist_bone in mesh.vertex_groups:
                        transfer_vertex_weights(mesh, twist_bone, parent_bone)

        bpy.ops.object.mode_set(mode='EDIT')
        for twist_bone, parent_bone in twist_bones:
            if twist_bone in armature.data.edit_bones and parent_bone in armature.data.edit_bones:
                logger.debug(f"Merging twist bone: {twist_bone} into {parent_bone}")
                twist = armature.data.edit_bones[twist_bone]
                parent = armature.data.edit_bones[parent_bone]
                parent.tail = twist.tail
                for child in twist.children:
                    child.parent = parent
                armature.data.edit_bones.remove(twist)

        bpy.ops.object.mode_set(mode='OBJECT')

    def get_org_remap(self, armature: Object) -> Dict[str, str]:
        """Get original bone remapping"""
        logger.debug("Getting original bone remapping")
        remap: Dict[str, str] = {}
        for bone in armature.data.bones:
            if self.is_def_bone(bone.name):
                name = self.get_proto_name(bone.name)
                parent = bone.parent
                while parent:
                    parent_name = self.get_proto_name(parent.name)
                    if parent_name != name:
                        if ('DEF-' + parent_name) in armature.data.bones:
                            remap[bone.name] = 'DEF-' + parent_name
                            break
                    parent = parent.parent
        return remap

    def get_special_remap(self) -> Dict[str, str]:
        """Get special bone remapping cases"""
        logger.debug("Getting special bone remapping")
        return {
            'DEF-thigh.L': 'DEF-pelvis.L',
            'DEF-thigh.R': 'DEF-pelvis.R',
            'DEF-upper_arm.L': 'DEF-shoulder.L',
            'DEF-upper_arm.R': 'DEF-shoulder.R',
        }

    def get_transform_copies(self, armature: Object) -> List[str]:
        """Get bones that need transform copies"""
        logger.debug("Getting transform copy bones")
        result: List[str] = []
        for bone in armature.pose.bones:
            if self.is_def_bone(bone.name) and not self.has_transform_copies(bone):
                result.append(bone.name)
        return result

    def has_transform_copies(self, bone: PoseBone) -> bool:
        """Check if bone has transform copy constraints"""
        return any(constraint.type == 'COPY_TRANSFORMS' for constraint in bone.constraints)

    def is_def_bone(self, bone_name: str) -> bool:
        """Check if bone is a DEF bone"""
        return bone_name.startswith('DEF-')

    def is_org_bone(self, bone_name: str) -> bool:
        """Check if bone is an ORG bone"""
        return bone_name.startswith('ORG-')

    def get_proto_name(self, bone_name: str) -> str:
        """Get the prototype name of a bone"""
        if self.is_def_bone(bone_name) or self.is_org_bone(bone_name):
            return bone_name[4:]
        return bone_name
