import bpy

def popup_error(message: str):
    def _draw(self, context):
        self.layout.label(text=message)
    try:
        bpy.context.window_manager.popup_menu(_draw, title="HardHat", icon='ERROR')
    except Exception:
        print(f"ERROR: {message}")

SETTINGS_BONE = 'Settings'
PROP_NAME = 'HardHat'


def find_armature():
    """Robustly resolve the armature that owns the Settings bone.
    Priority:
      1) Armature parenting the hard-hat object
      2) Armature parenting any selected hair object
      3) Active object if it's an armature with a Settings bone
      4) Any armature in scene with a Settings bone
    """
    # 1) Parent of hard-hat
    hat = bpy.data.objects.get('hard-hat')
    if hat and hat.parent and hat.parent.type == 'ARMATURE':
        if SETTINGS_BONE in hat.parent.pose.bones:
            return hat.parent

    # 2) Parent of any selected hair object
    for o in bpy.context.selected_objects or []:
        p = o.parent
        while p:
            if p.type == 'ARMATURE' and SETTINGS_BONE in p.pose.bones:
                return p
            p = p.parent

    # 3) Active armature
    a = bpy.context.active_object
    if a and a.type == 'ARMATURE' and SETTINGS_BONE in a.pose.bones:
        return a

    # 4) Any armature with Settings bone
    for o in bpy.data.objects:
        if o.type == 'ARMATURE' and SETTINGS_BONE in o.pose.bones:
            return o

    return None


def selected_hair_meshes():
    names = (bpy.context.scene.get('HH_HairTargets') or '').split(';')
    names = [n for n in names if n]
    if names:
        objs = [bpy.data.objects.get(n) for n in names]
        return [o for o in objs if o and o.type == 'MESH']
    popup_error("No HardHat hair targets set. Click 'Set HardHat Hair Targets' first.")
    return []


def ensure_property(arm_obj):
    pb = arm_obj.pose.bones.get(SETTINGS_BONE)
    if not pb:
        raise Exception("Settings pose bone not found")
    # Default True (hat on by default, like Devices)
    if PROP_NAME not in pb:
        pb[PROP_NAME] = True
    try:
        ui = pb.id_properties_ui(PROP_NAME)
        ui.update(description="Toggle HardHat visibility & masking", default=True)
    except Exception:
        pass
    try:
        # Mark the pose bone custom property overridable (use exact Blender path syntax with double quotes)
        if hasattr(pb, 'property_overridable_library_set'):
            pb.property_overridable_library_set(f'["{PROP_NAME}"]', True)
        # Also mark the owning object path overridable (helps in some linked setups)
        if hasattr(arm_obj, 'property_overridable_library_set'):
            arm_obj.property_overridable_library_set(f'pose.bones["{SETTINGS_BONE}"]["{PROP_NAME}"]', True)
    except Exception:
        pass
    return pb


def add_mask_drivers(arm_obj, obj, prop_path):
    # MainMask: visible when HardHat=False -> invert
    for name, invert in (('MainMask', True), ('HHMask', False)):
        mod = next((m for m in obj.modifiers if m.type == 'MASK' and m.name == name), None)
        if not mod:
            continue
        for prop in ('show_render', 'show_viewport'):
            # Remove existing driver
            try:
                mod.driver_remove(prop)
            except Exception:
                pass
            fcu = mod.driver_add(prop)
            drv = fcu.driver
            drv.type = 'SUM'
            var = drv.variables.new()
            var.name = 'hh'
            var.type = 'SINGLE_PROP'
            var.targets[0].id_type = 'OBJECT'
            var.targets[0].id = arm_obj
            var.targets[0].data_path = prop_path
            if invert:
                gen = fcu.modifiers.new('GENERATOR')
                gen.mode = 'POLYNOMIAL'
                gen.poly_order = 1
                gen.coefficients = (1.0, -1.0)


def add_hardhat_visibility_drivers(arm_obj, prop_path):
    hat = bpy.data.objects.get('hard-hat')
    if not hat:
        print("Warning: 'hard-hat' object not found for visibility drivers")
        return

    for prop in ('hide_render', 'hide_viewport'):
        try:
            hat.driver_remove(prop)
        except Exception:
            pass
        fcu = hat.driver_add(prop)
        drv = fcu.driver
        drv.type = 'SUM'
        var = drv.variables.new()
        var.name = 'hh'
        var.type = 'SINGLE_PROP'
        var.targets[0].id_type = 'OBJECT'
        var.targets[0].id = arm_obj
        var.targets[0].data_path = prop_path
        # invert so: visible when HardHat=True, hidden when HardHat=False
        gen = fcu.modifiers.new('GENERATOR')
        gen.mode = 'POLYNOMIAL'
        gen.poly_order = 1
        gen.coefficients = (1.0, -1.0)
        # Make overridable like Devices
        try:
            hat.property_overridable_library_set(prop, True)
        except Exception:
            pass


def add_shapekey_driver(arm_obj, obj, prop_path):
    if not obj.data.shape_keys:
        return
    key = obj.data.shape_keys.key_blocks.get('HardHat')
    if not key:
        return
    # Remove existing
    try:
        key.driver_remove('value')
    except Exception:
        pass
    fcu = key.driver_add('value')
    drv = fcu.driver
    drv.type = 'SUM'
    var = drv.variables.new()
    var.name = 'hh'
    var.type = 'SINGLE_PROP'
    var.targets[0].id_type = 'OBJECT'
    var.targets[0].id = arm_obj
    var.targets[0].data_path = prop_path


def run():
    arm = find_armature()
    if not arm:
        print("Error: No armature selected/found")
        return
    try:
        ensure_property(arm)
    except Exception as e:
        print(f"Error: {e}")
        return

    hair_objs = selected_hair_meshes()
    if not hair_objs:
        # No targets set; exit without applying any drivers
        return

    prop_path = f'pose.bones["{SETTINGS_BONE}"]["{PROP_NAME}"]'
    for obj in hair_objs:
        add_mask_drivers(arm, obj, prop_path)
        add_shapekey_driver(arm, obj, prop_path)
        print(f"Applied drivers to '{obj.name}' (existing masks/shapekey only)")

    # Also wire up the hard-hat visibility (only when hair targets exist)
    add_hardhat_visibility_drivers(arm, prop_path)

    # Force update to evaluate freshly created drivers
    try:
        bpy.context.view_layer.update()
        bpy.context.evaluated_depsgraph_get().update()
    except Exception:
        pass

    print("hh_settings: Done.")


run()