import bpy

# --- Control Calibration Variable ---
calibration = 40

# --- Configurable Axis & Direction Mapping ---
DIRECTION_MAP = {
    'x_pos': ('LOC_X', 1),
    'x_neg': ('LOC_X', -1),
    'y_pos': ('LOC_Y', 1),
    'y_neg': ('LOC_Y', -1),
    'z_pos': ('LOC_Z', 1),
    'z_neg': ('LOC_Z', -1),
    'rx_pos': ('ROT_X', 1),
    'rx_neg': ('ROT_X', -1),
    'ry_pos': ('ROT_Y', 1),
    'ry_neg': ('ROT_Y', -1),
    'rz_pos': ('ROT_Z', 1),
    'rz_neg': ('ROT_Z', -1)
}

# --- Start Execution ---
obj = bpy.context.object

if obj is None or obj.type != 'MESH':
    print("No mesh object selected.")
else:
    arm_mod = next((mod for mod in obj.modifiers if mod.type == 'ARMATURE'), None)
    if not arm_mod or not arm_mod.object or not arm_mod.object.type == 'ARMATURE':
        print("Mesh is not bound to an armature.")
    elif not obj.data.shape_keys or not obj.data.shape_keys.key_blocks:
        print("Mesh has no shape keys.")
    else:
        armature = arm_mod.object
        sk_data = obj.data.shape_keys

        for key in sk_data.key_blocks:
            if key.lock_shape:
                continue

            sk_name = key.name
            if not sk_name.startswith("SK-"):
                continue

            base = sk_name[3:]

            # Extract side if present
            side = None
            if base.endswith('.L'):
                side = '.L'
                base = base[:-2]
            elif base.endswith('.R'):
                side = '.R'
                base = base[:-2]

            parts = base.split('_')

            if len(parts) not in {4, 5}:
                continue

            area = parts[0]
            local_area = parts[1]

            if len(parts) == 5:
                lip_pos = parts[2]
                axis_dir = f"{parts[3]}_{parts[4]}"
            else:
                lip_pos = None
                axis_dir = f"{parts[2]}_{parts[3]}"

            if axis_dir not in DIRECTION_MAP:
                continue

            transform_type, sign = DIRECTION_MAP[axis_dir]

            # Flip sign for right side X-axis shape keys
            if side == '.R' and axis_dir.startswith('x_'):
                sign *= -1


            # Build bone name
            bone_base = f"CTL-{area}_{local_area}"
            if lip_pos:
                bone_base += f"_{lip_pos}"
            if side:
                bone_base += side

            bone_name = bone_base

            if bone_name not in armature.pose.bones:
                continue
            # Check if a driver already exists on the shape key's value
            existing_driver = None
            if sk_data.animation_data and sk_data.animation_data.drivers:
                for fc in sk_data.animation_data.drivers:
                    if fc.data_path == f'key_blocks["{key.name}"].value':
                        existing_driver = fc
                        break

            if existing_driver:
                print(f"Shape key '{key.name}' already has a driver. Skipping.")
                continue
            
            
            # Add driver
            fcurve = key.driver_add("value")
            driver = fcurve.driver
            driver.type = 'SCRIPTED'
            while driver.variables:
                driver.variables.remove(driver.variables[0])

            # Variable 1: Bone Transform
            var_bone = driver.variables.new()
            var_bone.name = 'BONE'
            var_bone.type = 'TRANSFORMS'
            targ_bone = var_bone.targets[0]
            targ_bone.id = armature
            targ_bone.bone_target = bone_name
            targ_bone.transform_type = transform_type
            targ_bone.transform_space = 'LOCAL_SPACE'

            # Expression
            if transform_type.startswith("ROT"):
                driver.expression = f"{'-' if sign == -1 else ''}BONE"
            else:
                # Variable 2: Calibration Property
                driver.expression = f"BONE * {'-' if sign == -1 else ''}" + str(calibration)