import bpy
import math

from . import bones_skel_struct
from . import bones_fingers_struct
from . import bones_toes_struct
from . import bones_web_struct
from . import bones_ik_struct



# Modify Bone Function
def modifybone(bone, head_empt, tail_empt, armature, matrixworldinv):
    bpy.context.view_layer.objects.active = armature
    bpy.ops.object.mode_set(mode='EDIT')

    edit_bone = armature.data.edit_bones.get(bone)
    empty_head = bpy.data.objects.get(head_empt)
    empty_tail = bpy.data.objects.get(tail_empt)
    allowhead = False
    allowtail = False
    alloweditbone = False

    # print(f"modifying bone: {edit_bone}")
    
    # Fetch world coordinates of the empties
    if empty_head:
        allowhead = True
        
    if empty_tail:  
        allowtail = True
    
    if edit_bone:
        alloweditbone = True
        
    if allowhead and allowtail and alloweditbone:
        head_world_pos = empty_head.matrix_world.translation
        tail_world_pos = empty_tail.matrix_world.translation

        # Multiply the inverse matrix by the world positions using the '@' operator
        edit_bone.head = matrixworldinv @ head_world_pos
        edit_bone.tail = matrixworldinv @ tail_world_pos



# Fix StretchTo Function
def fixbonestretching(bone, armature):
    bpy.context.view_layer.objects.active = armature
    edit_bone = armature.pose.bones.get(bone)
    for constaint in edit_bone.constraints:
        if constaint.type == 'STRETCH_TO':
            constaint.rest_length = 0.0


# Setup IK Bones
def setupIKbones(bone, helper_name, armature, matrixworldinv):
    bpy.context.view_layer.objects.active = armature
    bpy.ops.object.mode_set(mode='EDIT')

    edit_bone = armature.data.edit_bones.get(bone)
    helper_empt = bpy.data.objects.get(helper_name)

    if edit_bone:
        target_length = 0.025
        direction = (edit_bone.tail - edit_bone.head).normalized()
        
        empty_pos = helper_empt.matrix_world.translation
        edit_bone.head = matrixworldinv @ empty_pos
        edit_bone.tail = edit_bone.head + direction * target_length



def startAutoRig():
    mabulig_coll_name = "Mabulig-Rig"
    og_amt_obj = bpy.data.objects.get("OG_Armature")   

    # =========== Duplicate OG armature for future editing ===========
    target_coll = bpy.data.collections.get(mabulig_coll_name)
    armaturecheck = bpy.data.objects.get("Armature")


    if armaturecheck:
        print("Armature already exist")
        new_amt_obj = armaturecheck

    else:
        # Duplicate the armature object
        new_amt_obj = og_amt_obj.copy()

        # Duplicate the underlying armature data (bones, constraints, etc.)
        new_amt_obj.data = og_amt_obj.data.copy()

        new_amt_obj.name = "Armature"
        new_amt_obj.data.name = "Armature"

        target_coll.objects.link(new_amt_obj)

        
        print("Armature created")

    new_amt_obj.hide_viewport = False
        
        
    # Calculate the armature's inverse world matrix  
    matrix_world_inv = new_amt_obj.matrix_world.inverted()



    # =========== Initial State Checks ===========

    # Check if 'helpers' collection is hidden
    if bpy.data.collections['helpers'].hide_viewport:
        ishelperscollhidden = True
    else:
        ishelperscollhidden = False


    # Check if 'belly_mesh' is hidden
    if bpy.data.objects['belly_mesh'].hide_viewport:
        isbellymeshhidden = True
    else:
        isbellymeshhidden = False

    # Check if 'gular_mesh' is hidden
    if bpy.data.objects['gular_mesh'].hide_viewport:
        isgularmeshhidden = True
    else:
        isgularmeshhidden = False

    # Check if 'nostril_mesh' is hidden
    if bpy.data.objects['nostril_mesh'].hide_viewport:
        isnostrilmeshhidden = True
    else:
        isnostrilmeshhidden = False

    # Check if 'fingers' collection is hidden
    if bpy.data.collections['fingers'].hide_viewport:
        isfingerscollhidden = True
    else:
        isfingerscollhidden = False

    # Check if 'toes' collection is hidden
    if bpy.data.collections['toes'].hide_viewport:
        istoescollhidden = True
    else:
        istoescollhidden = False



    # =========== Skeleton Bones Setup ===========

    bone_data_list = bones_skel_struct.bones_data

    # Force 'helpers' collection to unhide
    bpy.data.collections['helpers'].hide_viewport = False

    for item in bone_data_list:
        b_name = item["bone_name"]
        h_helper = item["head_helper"]
        t_helper = item["tail_helper"]
        
        modifybone(b_name, h_helper, t_helper, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function



    # =========== Belly Bones Setup ===========

    belly_l_collection = bpy.data.collections.get("belly_l")
    belly_r_collection = bpy.data.collections.get("belly_r")

    # Force belly collections to unhide
    belly_l_collection.hide_viewport = False
    belly_r_collection.hide_viewport = False


    # For Left Belly
    for obj in belly_l_collection.objects:
        modifybone(obj.name, "belly_base_l", obj.name, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function

    # For Right Belly
    for obj in belly_r_collection.objects:
        modifybone(obj.name, "belly_base_r", obj.name, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function

    # Return belly collections to hidden
    belly_l_collection.hide_viewport = True
    belly_r_collection.hide_viewport = True




    # =========== Gullar Bones Setup ===========

    gular_collection = bpy.data.collections.get("gular")

    # Force gular collections to unhide
    gular_collection.hide_viewport = False

    for obj in gular_collection.objects:
        modifybone(obj.name, "gular_base", obj.name, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function

    # Return gular collections to hidden
    gular_collection.hide_viewport = True



    # =========== Nostril Bones Setup ===========

    nostril_l_collection = bpy.data.collections.get("nostril_l")
    nostril_r_collection = bpy.data.collections.get("nostril_r")

    # Force nostril collections to unhide
    nostril_l_collection.hide_viewport = False
    nostril_r_collection.hide_viewport = False

    # For Left Nostril
    for obj in nostril_l_collection.objects:
        modifybone(obj.name, "nostril_base_l", obj.name, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function


    # For Right Nostril
    for obj in nostril_r_collection.objects:
        modifybone(obj.name, "nostril_base_r", obj.name, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function

    # Return belly collections to hidden
    nostril_l_collection.hide_viewport = True
    nostril_r_collection.hide_viewport = True



    # =========== Metacarpal and Finger Bones Setup ===========

    bones_fingers_data_list = bones_fingers_struct.fingers_data

    # Force 'helpers' collection to unhide
    bpy.data.collections['fingers'].hide_viewport = False

    for item in bones_fingers_data_list:
        b_name = item["bone_name"]
        h_helper = item["head_helper"]
        t_helper = item["tail_helper"]
        
        modifybone(b_name, h_helper, t_helper, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function



    # =========== Metatarsus and Toe Bones Setup ===========

    bones_toes_data_list = bones_toes_struct.toes_data

    # Force 'toes' collection to unhide
    bpy.data.collections['toes'].hide_viewport = False

    for item in bones_toes_data_list:
        b_name = item["bone_name"]
        h_helper = item["head_helper"]
        t_helper = item["tail_helper"]
        
        modifybone(b_name, h_helper, t_helper, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function



    # =========== Webbing Bones ===========

    bones_web_data_list = bones_web_struct.webbing_data

    for item in bones_web_data_list:
        b_name = item["bone_name"]
        h_helper = item["head_helper"]
        t_helper = item["tail_helper"]
        
        modifybone(b_name, h_helper, t_helper, new_amt_obj, matrix_world_inv)  # Call Modify Bone Function
        fixbonestretching(b_name, new_amt_obj) # Call Fix StretchTo Function



    # =========== IK Bones ===========

    bones_ik_data_list = bones_ik_struct.bones_ik_data

    for item in bones_ik_data_list:
        b_name = item["bone_name"]
        pos_helper = item["pos_helper"]

        setupIKbones(b_name, pos_helper, new_amt_obj, matrix_world_inv) # Call Setup IK Bone Function



    # Returns 'helpers' collection to restore hide state prior to bone modification
    if ishelperscollhidden:
        bpy.data.collections['helpers'].hide_viewport = True
    else:
        bpy.data.collections['helpers'].hide_viewport = False

    # Returns belly collections and mesh to restore hide state prior to bone modification
    if isbellymeshhidden:
        bpy.data.objects['belly_mesh'].hide_viewport = True
    else:
        bpy.data.objects['belly_mesh'].hide_viewport = False

    # Returns gular collection and mesdh to restore hide state prior to bone modification
    if isgularmeshhidden:
        bpy.data.objects['gular_mesh'].hide_viewport = True
    else:
        bpy.data.objects['gular_mesh'].hide_viewport = False

    # Returns nostril collection and mesh to restore hide state prior to bone modification
    if isnostrilmeshhidden:
        bpy.data.objects['nostril_mesh'].hide_viewport = True
    else:
        bpy.data.objects['nostril_mesh'].hide_viewport = False

    # Returns 'fingers' collection to restore hide state prior to bone modification
    if isfingerscollhidden:
        bpy.data.collections['fingers'].hide_viewport = True
    else:
        bpy.data.collections['fingers'].hide_viewport = False

    # Returns 'toes' collection to restore hide state prior to bone modification
    if istoescollhidden:
        bpy.data.collections['toes'].hide_viewport = True
    else:
        bpy.data.collections['toes'].hide_viewport = False




    # Return to Object Mode
    bpy.ops.object.mode_set(mode='OBJECT')