# SPDX-FileCopyrightText: 2019-2022 Blender Foundation
#
# SPDX-License-Identifier: GPL-2.0-or-later

import bpy
import importlib
import importlib.util
import re

from itertools import count
from collections import defaultdict
from typing import TYPE_CHECKING, Any, Optional
from bpy.types import bpy_struct, Constraint, Object, PoseBone, Bone, Armature

from bpy.types import bpy_prop_array, bpy_prop_collection  # noqa
from idprop.types import IDPropertyArray
from mathutils import Vector

from .misc import ArmatureObject, wrap_list_to_lines, IdPropSequence, find_index, flatten_children

if TYPE_CHECKING:
    from ..base_rig import BaseRig
    from .. import RigifyColorSet

RIG_DIR = "rigs"  # Name of the directory where rig types are kept
METARIG_DIR = "metarigs"  # Name of the directory where metarigs are kept
TEMPLATE_DIR = "ui_templates"  # Name of the directory where ui templates are kept

# noinspection SpellCheckingInspection
outdated_types = {"pitchipoy.limbs.super_limb": "limbs.super_limb",
                  "pitchipoy.limbs.super_arm": "limbs.super_limb",
                  "pitchipoy.limbs.super_leg": "limbs.super_limb",
                  "pitchipoy.limbs.super_front_paw": "limbs.super_limb",
                  "pitchipoy.limbs.super_rear_paw": "limbs.super_limb",
                  "pitchipoy.limbs.super_finger": "limbs.super_finger",
                  "pitchipoy.super_torso_turbo": "spines.super_spine",
                  "pitchipoy.simple_tentacle": "limbs.simple_tentacle",
                  "pitchipoy.super_face": "faces.super_face",
                  "pitchipoy.super_palm": "limbs.super_palm",
                  "pitchipoy.super_copy": "basic.super_copy",
                  "pitchipoy.tentacle": "",
                  "palm": "limbs.super_palm",
                  "basic.copy": "basic.super_copy",
                  "biped.arm": "",
                  "biped.leg": "",
                  "finger": "",
                  "neck_short": "",
                  "misc.delta": "",
                  "spine": ""
                  }


def get_rigify_type(pose_bone: PoseBone) -> str:
    rigify_type = pose_bone.rigify_type  # noqa
    return rigify_type.replace(" ", "")


def get_rigify_params(pose_bone: PoseBone) -> Any:
    return pose_bone.rigify_parameters  # noqa


def get_rigify_colors(arm: Armature) -> IdPropSequence['RigifyColorSet']:
    return arm.rigify_colors  # noqa


def get_rigify_target_rig(arm: Armature) -> Optional[ArmatureObject]:
    return arm.rigify_target_rig  # noqa


def get_rigify_rig_basename(arm: Armature) -> str:
    return arm.rigify_rig_basename  # noqa


def get_rigify_mirror_widgets(arm: Armature) -> bool:
    return arm.rigify_mirror_widgets  # noqa


def get_rigify_force_widget_update(arm: Armature) -> bool:
    return arm.rigify_force_widget_update  # noqa


def get_rigify_finalize_script(arm: Armature) -> Optional[bpy.types.Text]:
    return arm.rigify_finalize_script  # noqa


def is_rig_base_bone(obj: Object, name):
    return bool(get_rigify_type(obj.pose.bones[name]))


def metarig_needs_upgrade(obj):
    return bool(obj.data.get("rigify_layers"))


def is_valid_metarig(context, *, allow_needs_upgrade=False):
    obj = context.object
    if not context.object:
        return False
    if obj.type != 'ARMATURE' or obj.data.get("rig_id") is not None:
        return False
    return allow_needs_upgrade or not metarig_needs_upgrade(context.object)


def upgrade_metarig_types(metarig: Object, revert=False):
    """
    Replaces rigify_type properties from old versions with their current names.

    metarig: rig to update.
    revert: revert types to previous version (if old type available)
    """

    if revert:
        vals = list(outdated_types.values())
        rig_defs = {v: k for k, v in outdated_types.items() if vals.count(v) == 1}
    else:
        rig_defs = outdated_types

    for bone in metarig.pose.bones:
        rig_type = bone.rigify_type
        if rig_type in rig_defs:
            bone.rigify_type = rig_defs[rig_type]

            parameters = get_rigify_params(bone)

            if 'leg' in rig_type:
                parameters.limb_type = 'leg'
            if 'arm' in rig_type:
                parameters.limb_type = 'arm'
            if 'paw' in rig_type:
                parameters.limb_type = 'paw'
            if rig_type == "basic.copy":
                parameters.make_widget = False


def resolve_layer_names(layers):
    """Combine full layer names if some buttons use fragments with parentheses only."""

    ui_rows = defaultdict(list)
    name_counts = defaultdict(int)

    for i, layer in enumerate(layers):
        if name := layer.get("name", "").strip():
            name_counts[name] += 1
            ui_rows[layer.get("row", 1)].append(name)

    def needs_rename(n):
        return name_counts[n] > 1 or n.startswith("(")

    def clean_stem(raw_name):
        return re.sub(r"\s*\(.*\)$", "", raw_name)

    def search_left(my_row, col_idx):
        while col_idx > 0:
            col_idx -= 1
            if not needs_rename(my_row[col_idx]):
                return clean_stem(my_row[col_idx])

    def search_up(my_row, row_idx, col_idx):
        while row_idx > 1:
            row_idx -= 1
            prev_row = ui_rows[row_idx]
            if len(prev_row) != len(my_row):
                return None
            if not needs_rename(prev_row[col_idx]):
                return clean_stem(prev_row[col_idx])

    names = []

    for i, layer in enumerate(layers):
        name: str = layer.get("name", "").strip()

        if name and needs_rename(name):
            row = layer.get("row", 1)

            cur_row = ui_rows[row]
            cur_col = cur_row.index(name)

            if stem := search_left(cur_row, cur_col):
                name = stem + " " + name
            elif stem := search_up(cur_row, row, cur_col):
                name = stem + " " + name

        names.append(name)

    return names


def upgrade_metarig_layers(metarig: ArmatureObject):
    from .layers import (REFS_LIST_SUFFIX, DEF_COLLECTION, MCH_COLLECTION, ORG_COLLECTION, ROOT_COLLECTION,
                         ensure_collection_uid)

    arm = metarig.data

    # Find layer collections
    coll_table = {}

    for coll in arm.collections_all:
        if m := re.match(r'^Layer (\d+)', coll.name):
            coll_table[int(m[1]) - 1] = coll

    # Assign names to special layers if they exist
    special_layers = {28: ROOT_COLLECTION, 29: DEF_COLLECTION, 30: MCH_COLLECTION, 31: ORG_COLLECTION}

    for idx, name in special_layers.items():
        if coll := coll_table.get(idx):
            coll.name = name

    # Apply existing layer metadata
    if layers := arm.get("rigify_layers"):
        names = resolve_layer_names(layers)

        # Enforce the special names
        for idx, name in special_layers.items():
            if idx < len(names) and names[idx]:
                names[idx] = name

        cur_idx = 0

        for i, layer in enumerate(layers):
            coll = coll_table.get(i)

            old_name = layer.get("name", "").strip()
            new_name = names[i]

            if new_name:
                if not coll:
                    coll = arm.collections.new(new_name)
                    coll_table[i] = coll
                else:
                    coll.name = new_name

            if coll:
                coll_idx = find_index(arm.collections_all, coll)
                arm.collections.move(coll_idx, cur_idx)
                cur_idx += 1

                coll.rigify_ui_row = layer.get("row", 1)

                if old_name and old_name != coll.name:
                    coll.rigify_ui_title = old_name

                coll.rigify_sel_set = layer.get("selset", False)
                coll.rigify_color_set_id = layer.get("group_prop", 0)

        del arm["rigify_layers"]

    arm.collections.active_index = 0

    # Remove empty rows, and ensure the root button position is at the bottom
    root_bcoll = coll_table.get(28)

    used_rows = set()
    for bcoll in arm.collections_all:
        if bcoll != root_bcoll and bcoll.rigify_ui_row > 0:
            used_rows.add(bcoll.rigify_ui_row)

    row_map = {}
    for i in range(1, max(used_rows) + 1):
        if i in used_rows:
            row_map[i] = len(row_map) + 1

    for bcoll in arm.collections_all:
        if bcoll == root_bcoll:
            bcoll.rigify_ui_row = len(row_map) + 3
        elif bcoll.rigify_ui_row > 0:
            bcoll.rigify_ui_row = row_map[bcoll.rigify_ui_row]

    # Convert the layer references in rig component parameters
    default_layers = [i == 1 for i in range(32)]
    default_map = {
        'faces.super_face': ['primary', 'secondary'],
        'limbs.arm': ['fk', 'tweak'],
        'limbs.front_paw': ['fk', 'tweak'],
        'limbs.leg': ['fk', 'tweak'],
        'limbs.paw': ['fk', 'tweak'],
        'limbs.rear_paw': ['fk', 'tweak'],
        'limbs.simple_tentacle': ['tweak'],
        'limbs.super_finger': ['tweak'],
        'limbs.super_limb': ['fk', 'tweak'],
        'spines.basic_spine': ['fk', 'tweak'],
    }

    for pose_bone in metarig.pose.bones:
        params = get_rigify_params(pose_bone)

        # Work around the stupid legacy default where one layer is implicitly selected
        for name_stem in default_map.get(get_rigify_type(pose_bone), []):
            prop_name = name_stem + "_layers"
            if prop_name not in params and name_stem + REFS_LIST_SUFFIX not in params:
                params[prop_name] = default_layers

        for prop_name, prop_value in list(params.items()):
            if prop_name.endswith("_layers") and isinstance(prop_value, IDPropertyArray) and len(prop_value) == 32:
                entries = []

                for i, show in enumerate(prop_value.to_list()):
                    if show:
                        coll = coll_table.get(i)
                        uid = ensure_collection_uid(coll) if coll else i
                        name = coll.name if coll else f"Layer {i + 1}"
                        entries.append({"uid": uid, "name": name})

                params[prop_name[:-7] + REFS_LIST_SUFFIX] = entries

                del params[prop_name]


##############################################
# Misc
##############################################

def rig_is_child(rig: 'BaseRig', parent: Optional['BaseRig'], *, strict=False):
    """
    Checks if the rig is a child of the parent.
    Unless strict is True, returns true if the rig and parent are the same.
    """
    if parent is None:
        return True

    if rig and strict:
        rig = rig.rigify_parent

    while rig:
        if rig is parent:
            return True

        rig = rig.rigify_parent

    return False


def get_parent_rigs(rig: 'BaseRig') -> list['BaseRig']:
    """Returns a list containing the rig and all of its parents."""
    result = []
    while rig:
        result.append(rig)
        rig = rig.rigify_parent
    return result


def get_resource(resource_name):
    """ Fetches a rig module by name, and returns it.
    """
    module = importlib.import_module(resource_name)
    importlib.reload(module)
    return module


def connected_children_names(obj: ArmatureObject, bone_name: str) -> list[str]:
    """ Returns a list of bone names (in order) of the bones that form a single
        connected chain starting with the given bone as a parent.
        If there is a connected branch, the list stops there.
    """
    bone = obj.data.bones[bone_name]
    names = []

    while True:
        connects = 0
        con_name = ""

        for child in bone.children:
            if child.use_connect:
                connects += 1
                con_name = child.name

        if connects == 1:
            names += [con_name]
            bone = obj.data.bones[con_name]
        else:
            break

    return names


def has_connected_children(bone: Bone):
    """ Returns true/false whether a bone has connected children or not.
    """
    t = False
    for b in bone.children:
        t = t or b.use_connect
    return t


def _list_bone_names_depth_first_sorted_rec(result_list: list[str], bone: Bone):
    result_list.append(bone.name)

    for child in sorted(list(bone.children), key=lambda b: b.name):
        _list_bone_names_depth_first_sorted_rec(result_list, child)


def list_bone_names_depth_first_sorted(obj: ArmatureObject):
    """Returns a list of bone names in depth first name sorted order."""
    result_list = []

    for bone in sorted(list(obj.data.bones), key=lambda b: b.name):
        if bone.parent is None:
            _list_bone_names_depth_first_sorted_rec(result_list, bone)

    return result_list


def _get_property_value(obj, name: str):
    """Retrieve the attribute value, converting from Blender to python types."""
    value = getattr(obj, name, None)
    if isinstance(value, bpy_prop_array):
        value = tuple(value)
    return value


def _format_property_value(prefix: str, value: Any, *, limit=90, indent=4) -> list[str]:
    """Format a property value assignment to lines, wrapping if too long."""

    if isinstance(value, tuple):
        return wrap_list_to_lines(prefix, '()', map(repr, value), limit=limit, indent=indent)

    if isinstance(value, list):
        return wrap_list_to_lines(prefix, '[]', map(repr, value), limit=limit, indent=indent)

    return [prefix + repr(value)]


def _generate_properties(lines, prefix, obj: bpy_struct, base_class: type, *,
                         defaults: Optional[dict[str, Any]] = None,
                         objects: Optional[dict[Any, str]] = None):
    obj_rna: bpy.types.Struct = type(obj).bl_rna  # noqa
    base_rna: bpy.types.Struct = base_class.bl_rna  # noqa

    defaults = defaults or {}
    block_props = set(prop.identifier for prop in base_rna.properties) - set(defaults.keys())

    for prop in obj_rna.properties:
        if prop.identifier not in block_props and not prop.is_readonly:
            cur_value = _get_property_value(obj, prop.identifier)

            if prop.identifier in defaults:
                if cur_value == defaults[prop.identifier]:
                    continue

            if isinstance(cur_value, bpy_struct):
                if objects and cur_value in objects:
                    lines.append('%s.%s = %s' % (prefix, prop.identifier, objects[cur_value]))
            else:
                lines += _format_property_value('%s.%s = ' % (prefix, prop.identifier), cur_value)


def write_metarig_widgets(obj: Object):
    from .widgets import write_widget

    widget_set = set()

    for pbone in obj.pose.bones:
        if pbone.custom_shape:
            widget_set.add(pbone.custom_shape)

    id_set = set()
    widget_map = {}
    code = []

    for widget_obj in widget_set:
        ident = re.sub("[^0-9a-zA-Z_]+", "_", widget_obj.name)

        if ident in id_set:
            for i in count(1):
                if ident + '_' + str(i) not in id_set:
                    break

        id_set.add(ident)
        widget_map[widget_obj] = ident

        code.append(write_widget(widget_obj, name=ident, use_size=False))

    return widget_map, code


def write_metarig(obj: ArmatureObject, layers=False, func_name="create",
                  groups=False, widgets=False):
    """
    Write a metarig as a python script, this rig is to have all info needed for
    generating the real rig with rigify.
    """
    from .layers import REFS_LIST_SUFFIX, is_collection_ref_list_prop

    code = [
        "import bpy\n",
        "from rna_prop_ui import rna_idprop_ui_create\n",
        "from mathutils import Color\n\n",
    ]

    # Widget object creation functions if requested
    if widgets:
        widget_map, widget_code = write_metarig_widgets(obj)

        if widget_map:
            code.append("from rigify.utils.widgets import widget_generator\n\n")
            code += widget_code
    else:
        widget_map = {}

    # Start of the metarig function
    code.append("def %s(obj):  # noqa" % func_name)
    code.append("    # generated by rigify.utils.write_metarig")
    bpy.ops.object.mode_set(mode='EDIT')
    code.append("    bpy.ops.object.mode_set(mode='EDIT')")
    code.append("    arm = obj.data")

    arm = obj.data

    # Rigify bone group colors info
    rigify_colors = get_rigify_colors(arm)

    if groups and len(rigify_colors) > 0:
        code.append("\n    for i in range(" + str(len(rigify_colors)) + "):")
        code.append("        arm.rigify_colors.add()\n")

        for i in range(len(rigify_colors)):
            name = rigify_colors[i].name
            active = rigify_colors[i].active
            normal = rigify_colors[i].normal
            select = rigify_colors[i].select
            standard_colors_lock = rigify_colors[i].standard_colors_lock
            code.append('    arm.rigify_colors[' + str(i) + '].name = "' + name + '"')
            code.append('    arm.rigify_colors[' + str(i)
                        + '].active = Color((%.4f, %.4f, %.4f))' % tuple(active[:]))
            code.append('    arm.rigify_colors[' + str(i)
                        + '].normal = Color((%.4f, %.4f, %.4f))' % tuple(normal[:]))
            code.append('    arm.rigify_colors[' + str(i)
                        + '].select = Color((%.4f, %.4f, %.4f))' % tuple(select[:]))
            code.append('    arm.rigify_colors[' + str(i)
                        + '].standard_colors_lock = ' + str(standard_colors_lock))

    # Rigify collection layout info
    if layers:
        collection_attrs = {
            'ui_row': 0, 'ui_title': '', 'sel_set': False, 'color_set_id': 0
        }

        code.append('\n    bone_collections = {}')

        code.append('\n    for bcoll in list(arm.collections_all):'
                    '\n        arm.collections.remove(bcoll)\n')

        args = ', '.join(f'{k}={repr(v)}' for k, v in collection_attrs.items())

        code.append(f"    def add_bone_collection(name, *, parent=None, {args}):")
        code.append(f"        new_bcoll = arm.collections.new(name, parent=bone_collections.get(parent))")
        for k, _v in collection_attrs.items():
            code.append(f"        new_bcoll.rigify_{k} = {k}")
        code.append("        bone_collections[name] = new_bcoll")

        code.append("""
    def assign_bone_collections(pose_bone, *coll_names):
        assert not len(pose_bone.bone.collections)
        for name in coll_names:
            bone_collections[name].assign(pose_bone)

    def assign_bone_collection_refs(params, attr_name, *coll_names):
        ref_list = getattr(params, attr_name + '_coll_refs', None)
        if ref_list is not None:
            for name in coll_names:
                ref_list.add().set_collection(bone_collections[name])
""")

        for bcoll in flatten_children(arm.collections):
            args = [repr(bcoll.name)]
            if bcoll.parent:
                args.append(f"parent={bcoll.parent.name!r}")
            for k, v in collection_attrs.items():
                value = getattr(bcoll, "rigify_" + k)
                if value != v:
                    args.append(f"{k}={repr(value)}")
            code.append(f"    add_bone_collection({', '.join(args)})")

    # write parents first
    bones = [(len(bone.parent_recursive), bone.name) for bone in arm.edit_bones]
    bones.sort(key=lambda item: item[0])
    bones = [item[1] for item in bones]

    code.append("\n    bones = {}\n")

    # noinspection SpellCheckingInspection
    extra_props = {
        'bbone_segments': 1,
        'bbone_mapping_mode': 'STRAIGHT',
        'bbone_easein': 1, 'bbone_easeout': 1,
        'bbone_rollin': 0, 'bbone_rollout': 0,
        'bbone_curveinx': 0, 'bbone_curveinz': 0,
        'bbone_curveoutx': 0, 'bbone_curveoutz': 0,
        'bbone_scalein': Vector((1, 1, 1)),
        'bbone_scaleout': Vector((1, 1, 1)),
    }

    for bone_name in bones:
        bone = arm.edit_bones[bone_name]
        code.append("    bone = arm.edit_bones.new(%r)" % bone.name)
        code.append("    bone.head = %.4f, %.4f, %.4f" % bone.head.to_tuple(4))
        code.append("    bone.tail = %.4f, %.4f, %.4f" % bone.tail.to_tuple(4))
        code.append("    bone.roll = %.4f" % bone.roll)
        code.append("    bone.use_connect = %s" % str(bone.use_connect))
        if bone.inherit_scale != 'FULL':
            code.append("    bone.inherit_scale = %r" % str(bone.inherit_scale))
        if bone.parent:
            code.append("    bone.parent = arm.edit_bones[bones[%r]]" % bone.parent.name)
        for prop, default in extra_props.items():
            value = getattr(bone, prop)
            if value != default:
                code.append(f"    bone.{prop} = {value!r}")
        code.append("    bones[%r] = bone.name" % bone.name)

    bpy.ops.object.mode_set(mode='OBJECT')
    code.append("")
    code.append("    bpy.ops.object.mode_set(mode='OBJECT')")

    if widgets and widget_map:
        code.append("    widget_map = {}")

    # Rig type and other pose properties
    for bone_name in bones:
        pbone = obj.pose.bones[bone_name]

        rigify_type = get_rigify_type(pbone)
        rigify_parameters = get_rigify_params(pbone)

        code.append("    pbone = obj.pose.bones[bones[%r]]" % bone_name)
        code.append("    pbone.rigify_type = %r" % rigify_type)
        code.append("    pbone.lock_location = %s" % str(tuple(pbone.lock_location)))
        code.append("    pbone.lock_rotation = %s" % str(tuple(pbone.lock_rotation)))
        code.append("    pbone.lock_rotation_w = %s" % str(pbone.lock_rotation_w))
        code.append("    pbone.lock_scale = %s" % str(tuple(pbone.lock_scale)))
        code.append("    pbone.rotation_mode = %r" % pbone.rotation_mode)
        if layers and len(pbone.bone.collections):
            args = ', '.join(f"'{bcoll.name}'" for bcoll in pbone.bone.collections)
            code.append(f"    assign_bone_collections(pbone, {args})")

        # Rig type parameters
        for param_name in rigify_parameters.keys():
            param = _get_property_value(rigify_parameters, param_name)

            if isinstance(param, bpy_prop_collection):
                if layers and param_name.endswith(REFS_LIST_SUFFIX) and is_collection_ref_list_prop(param):
                    bcoll_set = [item.find_collection() for item in param]
                    bcoll_set = [bcoll for bcoll in bcoll_set if bcoll is not None]
                    if len(bcoll_set) > 0:
                        args = ', '.join(f"'{bcoll.name}'" for bcoll in bcoll_set)
                        code.append(f"    assign_bone_collection_refs("
                                    f"pbone.rigify_parameters, '{param_name[:-10]}', {args})")
                continue

            if param is not None:
                code.append("    try:")
                code += _format_property_value(
                    f"        pbone.rigify_parameters.{param_name} = ", param)
                code.append("    except AttributeError:")
                code.append("        pass")

        # Custom properties
        custom_properties = {
            property_name: value for property_name, value in pbone.items()
            if property_name not in pbone.bl_rna.properties.keys()
            and type(pbone[property_name]) in (float, int)
        }

        if custom_properties:
            code.append('    # custom properties')

        for custom_property, current_value in custom_properties.items():
            props_data = pbone.id_properties_ui(custom_property).as_dict()
            code.append(f"    rna_idprop_ui_create(")
            code.append(f"        pbone,")
            code.append(f"        {custom_property!r},")
            code.append(f"        default={props_data['default']!r},")
            if 'min' in props_data:
                code.append(f"        min={props_data['min']},")
            if 'max' in props_data:
                code.append(f"        max={props_data['max']},")
            if 'soft_min' in props_data:
                code.append(f"        soft_min={props_data['soft_min']},")
            if 'soft_max' in props_data:
                code.append(f"        soft_max={props_data['soft_max']},")
            if 'subtype' in props_data:
                code.append(f"        subtype={props_data['subtype']!r},")
            if 'description' in props_data:
                code.append(f"        description={props_data['description']!r},")
            if 'precision' in props_data:
                code.append(f"        precision={props_data['precision']},")
            if 'step' in props_data:
                code.append(f"        step={props_data['step']},")
            code.append(f"    )")
            if props_data['default'] != current_value:
                code.append(f"    pbone[{custom_property!r}] = {current_value}")

        # Constraints
        for con in pbone.constraints:
            code.append("    con = pbone.constraints.new(%r)" % con.type)
            code.append("    con.name = %r" % con.name)
            # Add target first because of target_space handling
            if con.type == 'ARMATURE':
                for tgt in con.targets:
                    code.append("    tgt = con.targets.new()")
                    code.append("    tgt.target = obj")
                    code.append("    tgt.subtarget = %r" % tgt.subtarget)
                    code.append("    tgt.weight = %.3f" % tgt.weight)
            elif getattr(con, 'target', None) == obj:
                code.append("    con.target = obj")
            # Generic properties
            _generate_properties(
                code, "    con", con, Constraint,
                defaults={
                    'owner_space': 'WORLD', 'target_space': 'WORLD',
                    'mute': False, 'influence': 1.0,
                    'target': obj,
                },
                objects={obj: 'obj'},
            )
        # Custom widgets
        if widgets and pbone.custom_shape:
            widget_id = widget_map[pbone.custom_shape]
            code.append("    if %r not in widget_map:" % widget_id)
            code.append(("        widget_map[%r] = create_%s_widget(obj, pbone.name, "
                         "widget_name=%r, widget_force_new=True)")
                        % (widget_id, widget_id, pbone.custom_shape.name))
            code.append("    pbone.custom_shape = widget_map[%r]" % widget_id)

    code.append("\n    bpy.ops.object.mode_set(mode='EDIT')")
    code.append("    for bone in arm.edit_bones:")
    code.append("        bone.select = False")
    code.append("        bone.select_head = False")
    code.append("        bone.select_tail = False")

    code.append("    for b in bones:")
    code.append("        bone = arm.edit_bones[bones[b]]")
    code.append("        bone.select = True")
    code.append("        bone.select_head = True")
    code.append("        bone.select_tail = True")
    code.append("        bone.bbone_x = bone.bbone_z = bone.length * 0.05")
    code.append("        arm.edit_bones.active = bone")

    if not layers:
        code.append("        if bcoll := arm.collections.active:")
        code.append("            bcoll.assign(bone)")
    else:
        code.append("\n    arm.collections.active_index = 0")

    code.append("\n    return bones")

    code.append('\n\nif __name__ == "__main__":')
    code.append("    " + func_name + "(bpy.context.active_object)\n")

    return "\n".join(code)