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

import bpy
from bpy.types import PoseBone

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."""

    glue_head_mode: str
    glue_use_tail: bool

    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

    head_constraint_node: ControlQueryNode
    tail_position_node: 'PositionQueryNode'

    @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(cls, 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",
        )

        cls.add_relink_constraints_params(params)

        super().add_parameters(params)

    @classmethod
    def parameters_ui(cls, 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: PoseBone) -> str:
        return bone.name

    ####################################################
    # BONES

    bones: BaseSkinRig.ToplevelBones[
        str,
        'SimpleGlueRig.CtrlBones',
        'SimpleGlueRig.MchBones',
        str
    ]

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

    head_position_node: 'PositionQueryNode'

    @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: PoseBone) -> list[str]:
        # Still only bind to one bone
        return [bone.name]

    # Assign the 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'