# SPDX-License-Identifier: GPL-2.0-or-later

import bpy

from ...utils.naming import make_derived_name
from ...utils.widgets_basic import create_cube_widget
from ...utils.mechanism import move_all_constraints

from ...base_rig import stage
from ...base_generate import SubstitutionRig

from .skin_nodes import ControlQueryNode
from .skin_rigs import BaseSkinRig

from ..basic.raw_copy import RelinkConstraintsMixin

from .basic_chain import Rig as BasicChainRig


class Rig(SubstitutionRig):
    """Skin rig component that injects constraints into a control generated by other rigs."""

    def substitute(self):
        # Deformation is implemented by inheriting from the chain rig, so
        # enabling it requires switching between two different classes.
        if self.params.skin_glue_head_mode == 'BRIDGE':
            return [self.instantiate_rig(BridgeGlueRig, self.base_bone)]
        else:
            return [self.instantiate_rig(SimpleGlueRig, self.base_bone)]


def add_parameters(params):
    SimpleGlueRig.add_parameters(params)
    BridgeGlueRig.add_parameters(params)


def parameters_ui(layout, params):
    if params.skin_glue_head_mode == 'BRIDGE':
        BridgeGlueRig.parameters_ui(layout, params)
    else:
        SimpleGlueRig.parameters_ui(layout, params)


class BaseGlueRig(BaseSkinRig, RelinkConstraintsMixin):
    """Base class for the glue rigs."""

    def initialize(self):
        super().initialize()

        self.glue_head_mode = self.params.skin_glue_head_mode

        self.glue_use_tail = self.params.relink_constraints and self.params.skin_glue_use_tail
        self.relink_unmarked_constraints = self.glue_use_tail

    ####################################################
    # QUERY NODES

    @stage.initialize
    def init_glue_nodes(self):
        bone = self.get_bone(self.base_bone)

        self.head_constraint_node = ControlQueryNode(
            self, self.base_bone, point=bone.head
        )

        if self.glue_use_tail:
            self.tail_position_node = PositionQueryNode(
                self, self.base_bone, point=bone.tail,
                needs_reparent=self.params.skin_glue_tail_reparent,
            )

    ####################################################
    # GLUE CONSTRAINTS

    def rig_glue_constraints(self):
        org = self.base_bone
        ctrl = self.head_constraint_node.control_bone

        self.relink_bone_constraints(org)

        # Add the built-in constraint
        if self.glue_use_tail:
            target = self.tail_position_node.output_bone
            add_mode = self.params.skin_glue_add_constraint
            inf = self.params.skin_glue_add_constraint_influence

            if add_mode == 'COPY_LOCATION':
                self.make_constraint(
                    ctrl, 'COPY_LOCATION', target, insert_index=0,
                    owner_space='LOCAL', target_space='LOCAL',
                    use_offset=True, influence=inf
                )
            elif add_mode == 'COPY_LOCATION_OWNER':
                self.make_constraint(
                    ctrl, 'COPY_LOCATION', target, insert_index=0,
                    owner_space='LOCAL', target_space='LOCAL_OWNER_ORIENT',
                    use_offset=True, influence=inf
                )

        move_all_constraints(self.obj, org, ctrl)

    def find_relink_target(self, spec, old_target):
        if self.glue_use_tail and (spec == 'TARGET' or spec == '' == old_target):
            return self.tail_position_node.output_bone

        return super().find_relink_target(spec, old_target)

    ####################################################
    # SETTINGS

    @classmethod
    def add_parameters(self, params):
        params.skin_glue_head_mode = bpy.props.EnumProperty(
            name='Glue Mode',
            items=[('CHILD', 'Child Of Control',
                    "The glue bone becomes a child of the control bone"),
                   ('MIRROR', 'Mirror Of Control',
                    "The glue bone becomes a sibling of the control bone with Copy Transforms"),
                   ('REPARENT', 'Mirror With Parents',
                    "The glue bone keeps its parent, but uses Copy Transforms to group both local and parent induced motion of the control into local space"),
                   ('BRIDGE', 'Deformation Bridge',
                    "Other than adding glue constraints to the control, the rig acts as a one segment basic deform chain")],
            default='CHILD',
            description="Specifies how the glue bone is rigged to the control at the bone head location",
        )

        params.skin_glue_use_tail = bpy.props.BoolProperty(
            name='Use Tail Target',
            default=False,
            description='Find the control at the bone tail location and use it to relink TARGET or any constraints without an assigned subtarget or relink spec'
        )

        params.skin_glue_tail_reparent = bpy.props.BoolProperty(
            name='Target Local With Parents',
            default=False,
            description='Include transformations induced by target parents into target local space'
        )

        params.skin_glue_add_constraint = bpy.props.EnumProperty(
            name='Add Constraint',
            items=[('NONE', 'No New Constraint',
                    "Don't add new constraints"),
                   ('COPY_LOCATION', 'Copy Location (Local)',
                    "Add a constraint to copy Local Location with Offset. If the owner and target control " +
                    "rest orientations are different, the global movement direction will change accordingly"),
                   ('COPY_LOCATION_OWNER', 'Copy Location (Local, Owner Orientation)',
                    "Add a constraint to copy Local Location (Owner Orientation) with Offset. Even if the owner and " +
                    "target controls have different rest orientations, the global movement direction would be the same")],
            default='NONE',
            description="Add one of the common constraints linking the control to the tail target",
        )

        params.skin_glue_add_constraint_influence = bpy.props.FloatProperty(
            name="Influence",
            default=1.0, min=0, max=1,
            description="Influence of the added constraint",
        )

        self.add_relink_constraints_params(params)

        super().add_parameters(params)

    @classmethod
    def parameters_ui(self, layout, params):
        layout.prop(params, "skin_glue_head_mode")
        layout.prop(params, "relink_constraints")

        if params.relink_constraints:
            col = layout.column()
            col.prop(params, "skin_glue_use_tail")

            col2 = col.column()
            col2.active = params.skin_glue_use_tail
            col2.prop(params, "skin_glue_tail_reparent")

            col = layout.column()
            col.active = params.skin_glue_use_tail
            col.prop(params, "skin_glue_add_constraint", text="Add")

            col3 = col.column()
            col3.active = params.skin_glue_add_constraint != 'NONE'
            col3.prop(params, "skin_glue_add_constraint_influence", slider=True)

        layout.label(text="All constraints are moved to the control bone.", icon='INFO')

        super().parameters_ui(layout, params)


class SimpleGlueRig(BaseGlueRig):
    """Normal glue rig that only does glue."""

    def find_org_bones(self, bone):
        return bone.name

    ####################################################
    # QUERY NODES

    @stage.initialize
    def init_glue_nodes(self):
        super().init_glue_nodes()

        bone = self.get_bone(self.base_bone)

        self.head_position_node = PositionQueryNode(
            self, self.base_bone, point=bone.head,
            rig_org=self.glue_head_mode != 'CHILD',
            needs_reparent=self.glue_head_mode == 'REPARENT',
        )

    ##############################
    # ORG chain

    @stage.parent_bones
    def parent_org_bone(self):
        if self.glue_head_mode == 'CHILD':
            self.set_bone_parent(self.bones.org, self.head_position_node.output_bone)

    @stage.rig_bones
    def rig_org_bone(self):
        # This executes before head_position_node owned a by generator plugin
        self.rig_glue_constraints()


class BridgeGlueRig(BaseGlueRig, BasicChainRig):
    """Glue rig that also behaves like a deformation chain rig."""

    def find_org_bones(self, bone):
        # Still only bind to one bone
        return [bone.name]

    # Assign lowest priority
    chain_priority = -20

    # Orientation is irrelevant since controls should be merged into others
    use_skin_control_orientation_bone = False

    ####################################################
    # QUERY NODES

    @stage.prepare_bones
    def prepare_glue_nodes(self):
        # Verify that all nodes of the chain have been merged into others
        for node in self.control_nodes:
            if node.is_master_node:
                self.raise_error('glue control {} was not merged', node.name)

    ##############################
    # ORG chain

    @stage.rig_bones
    def rig_org_chain(self):
        # Move the user constraints away before the chain adds new ones
        self.rig_glue_constraints()

        super().rig_org_chain()


class PositionQueryNode(ControlQueryNode):
    """Finds the position of the highest layer control and rig reparent and/or org bone"""

    def __init__(self, rig, org, *, point=None, needs_reparent=False, rig_org=False):
        super().__init__(rig, org, point=point, find_highest_layer=True)

        self.needs_reparent = needs_reparent
        self.rig_org = rig_org

    @property
    def output_bone(self):
        if self.rig_org:
            return self.org
        elif self.needs_reparent:
            return self.reparent_bone
        else:
            return self.control_bone

    def initialize(self):
        if self.needs_reparent:
            self.build_parent(reparent=not self.rig_org)

    def parent_bones(self):
        if self.rig_org:
            if self.needs_reparent:
                parent = self.node_parent.output_bone
            else:
                parent = self.get_bone_parent(self.control_bone)

            self.set_bone_parent(self.org, parent, inherit_scale='AVERAGE')

    def apply_bones(self):
        if self.rig_org:
            self.get_bone(self.org).matrix = self.merged_master.matrix

    def rig_bones(self):
        if self.rig_org:
            self.make_constraint(self.org, 'COPY_TRANSFORMS', self.control_bone)


def create_sample(obj):
    from rigify.rigs.basic.super_copy import create_sample as inner
    obj.pose.bones[inner(obj)["Bone"]].rigify_type = 'skin.glue'
