bl_info = {
    "name": "Z-Display",
    "author": "Marcin Zieliński, Morgan Le Chénéchal, Z-Anatomy",
    "description": "Smart hiding tools for Blender objects and collections",
    "blender": (2, 80, 0),
    "version": (0, 0, 1),
    "location": "View3D > Sidebar > Z-Anatomy",
    "warning": "",
    "category": "Interface"
}

import bpy
from bpy_extras.object_utils import object_data_add

# Ensemble des éléments de label
label_elements = {".t", ".j", ".s", ".i"}

def family_all(object):
    ''' Object + Grand children without ancestors '''
    # Initialisation de la famille avec l'objet et ses enfants directs
    family = [object.children[:] + (object,)]

    # Fonction récursive pour ajouter tous les enfants et petits-enfants
    def rec(object, family):
        family[0] += object.children
        for ob in object.children:
            rec(ob, family)

    # Appel récursif pour chaque enfant de l'objet
    for ob in object.children:
        rec(ob, family)

    # Retourne la famille complète
    return family.pop()

def family(object):
    ''' Object + Grand children without ancestors (labels only) '''
    # Fonction interne pour filtrer les enfants qui sont des labels
    def inner_elements(children):
        return tuple(ob for ob in children if any(x in ob.name for x in label_elements))
        # return [ob for ob in children if any(x in ob.name for x in label_elements)]

    # Initialisation des labels avec les enfants directs de l'objet
    labels = inner_elements(object.children)
    family = [labels + (object,)]

    # Fonction récursive pour ajouter tous les labels enfants
    def rec(object, family):
        labels = inner_elements(object.children)
        family[0] += labels
        for ob in labels:
            rec(ob, family)

    # Appel récursif pour chaque enfant de l'objet
    for ob in object.children:
        rec(ob, family)

    # Retourne la famille complète
    return family.pop()


class OBJECT_OT_hide_wrapper(bpy.types.Operator):
    """Hide"""
    bl_idname = "object.hide_wrapper"  # Identifiant unique de l'opérateur
    bl_label = "Hide"  # Nom de l'opérateur
    bl_options = {'REGISTER', 'UNDO'}  # Options de l'opérateur (enregistrement et annulation)

    # Propriétés de l'opérateur
    unselected: bpy.props.BoolProperty(default=True, name="Hide unselected")  # Masquer les objets non sélectionnés
    unselected_in_layer: bpy.props.BoolProperty(default=True, name="From active layer")  # Masquer les objets non sélectionnés dans la couche active
    follow_parent: bpy.props.BoolProperty(default=True, name="Follow parent's visibility")  # Suivre la visibilité du parent

    @classmethod
    def poll(cls, context):
        # Vérifie si l'opérateur peut être exécuté
        return context.mode == 'OBJECT' and context.object is not None

    def execute(self, context):
        # Fonction pour trouver la racine d'un objet
        def findRoot(ob):
            while ob and ob.parent:
                if ob in objects_to:
                    objects_to.remove(ob)
                # if '%' in ob.parent.name:
                if not any(x in ob.name for x in label_elements):
                    return ob
                if not ob.parent.parent and not self.unselected:
                    return ob
                ob = ob.parent
            return ob

        roots = []  # Liste pour stocker les racines des objets
        objects_to = set(context.selected_objects)  # Ensemble des objets sélectionnés

        # Trouver les racines des objets sélectionnés
        while objects_to:
            ob = objects_to.pop()
            roots.append(findRoot(ob))

        if self.unselected and not self.unselected_in_layer:
            # Masquer les objets non sélectionnés dans la scène
            all_visible_objects = set(context.visible_objects)
            lights = {l for l in context.view_layer.objects if l.type == 'LIGHT'}

            family_obs = (ob for r in roots for ob in family(r))
            for ob in all_visible_objects - lights - set(family_obs):
                ob.hide_set(True)
        elif self.unselected and self.unselected_in_layer:
            # Masquer les objets non sélectionnés dans la couche active
            all_layer_objects = set()
            for c in context.scene.collection.children:
                if context.object.name in c.all_objects:
                    all_layer_objects = set(c.all_objects)
                    break
            lights = {l for l in context.view_layer.objects if l.type == 'LIGHT'}

            family_obs = (ob for r in roots for ob in family(r))
            for ob in all_layer_objects - set(family_obs) - lights:
                ob.hide_set(True)
        else:
            # Masquer les objets de la famille
            if self.follow_parent:
                family_obs = (ob for r in roots for ob in family_all(r))
            else:
                family_obs = (ob for r in roots for ob in family(r))

            for ob in family_obs:
                ob.hide_set(True)

        return {'FINISHED'}

class OBJECT_OT_hide_view_clear_wrapper(bpy.types.Operator):
    """Show Hidden Objects"""
    bl_idname = "object.hide_view_clear_wrapper"  # Identifiant unique de l'opérateur
    bl_label = "Show Hidden Objects"  # Nom de l'opérateur
    bl_options = {'REGISTER', 'UNDO'}  # Options de l'opérateur (enregistrement et annulation)

    # Propriétés de l'opérateur
    select: bpy.props.BoolProperty()  # Propriété pour sélectionner les objets
    active_layer: bpy.props.BoolProperty(default=True, name="Affect only active layer")  # Propriété pour affecter uniquement la couche active

    @classmethod
    def poll(cls, context):
        # Vérifie si l'opérateur peut être exécuté en mode objet
        return context.mode == 'OBJECT'

    def execute(self, context):
        if self.active_layer:
            # Si active_layer est vrai, affecter uniquement la couche active
            for col in context.scene.collection.children:
                if context.object.name in col.all_objects:
                    break

            # Rendre tous les objets de la couche active visibles
            for ob in col.all_objects:
                ob.hide_set(False)
        else:
            # Sinon, utiliser l'opérateur intégré pour rendre tous les objets visibles
            bpy.ops.object.hide_view_clear(select=self.select)

        # Masquer les objets avec ".t" dans leur nom, sauf l'objet actif
        for ob in (o for o in bpy.context.visible_objects if any(x in o.name for x in label_elements) and o != context.object):
            if not ob.parent == context.object:
                ob.hide_set(True)
                for child in ob.children:
                    child.hide_set(True)

        # Si l'option enable_group_labels n'est pas activée, mettre à jour la case à cocher du groupe de labels
        #if not context.scene.zanatomy.enable_group_labels:
        #    label_group_checkbox_update()

        return {'FINISHED'}


class OBJECT_OT_local_view_wrapper(bpy.types.Operator):
    """Local view (with lights)"""
    bl_idname = "view3d.localview_lights"  # Identifiant unique de l'opérateur
    bl_label = "Local View (w/ lights)"  # Nom de l'opérateur
    # bl_options = {'REGISTER', 'UNDO'}

    frame_selected: bpy.props.BoolProperty(default=True, name="Frame Selected")  # Propriété pour encadrer les objets sélectionnés

    @classmethod
    def poll(cls, context):
        # Vérifie si l'opérateur peut être exécuté en mode objet et si un objet est sélectionné
        return context.mode == 'OBJECT' and context.object is not None

    def execute(self, context):
        def findRoot(ob):
            # Trouve la racine d'un objet en remontant la hiérarchie des parents
            while ob and ob.parent:
                if ob in objects_to:
                    objects_to.remove(ob)
                if not any(x in ob.name for x in label_elements):
                    return ob
                if not ob.parent.parent:
                    return ob
                ob = ob.parent
            return ob

        roots = []  # Liste pour stocker les racines des objets
        objects_to = set(context.selected_objects)  # Ensemble des objets sélectionnés
        selected_objects = list(objects_to)  # Liste des objets sélectionnés
        while objects_to:
            ob = objects_to.pop()
            roots.append(findRoot(ob))

        lights = [o for o in context.visible_objects if o.type == 'LIGHT']  # Liste des lumières visibles
        hide_selects = []  # Liste des objets avec hide_select activé
        hidden = []  # Liste des objets masqués
        obj_family = []  # Liste des objets de la famille
        for obj in roots:
            obj_family += list(family(obj))
        obj_family += lights
        print(obj_family)

        # Rendre tous les objets de la famille visibles et sélectionnables
        for ob in obj_family:
            if ob.hide_get():
                ob.hide_set(False)
                hidden.append(ob)
            if ob.hide_select:
                ob.hide_select = False
                hide_selects.append(ob)
            ob.select_set(True)

        # Activer la vue locale
        bpy.ops.view3d.localview('INVOKE_DEFAULT', frame_selected=self.frame_selected)

        # Restaurer l'état des objets
        for ob in obj_family:
            ob.select_set(False)
        #obj.select_set(True)
        for ob in hide_selects:
            ob.hide_select = True
        for ob in hidden:
            ob.hide_set(True)

        # Sélectionner les objets d'origine
        for ob in selected_objects:
            ob.select_set(True)
        bpy.ops.view3d.view_selected()

        return {"FINISHED"}

class ZANATOMY_PT_display_panel(bpy.types.Panel):
    bl_label = "Z-display"
    bl_idname = "VIEW3D_PT_z_display"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = 'Z-Anatomy'
    bl_options = {'DEFAULT_CLOSED'}

    def draw(self, context):
        layout = self.layout
        layout.prop(context.scene, "z_display_enable_addon")
        if context.scene.z_display_enable_addon:
            layout.operator("object.hide_wrapper", text="Isolate on Layer")
            layout.operator("object.hide_view_clear_wrapper", text="Unhide Layer")
            layout.operator("view3d.localview_lights", text="Local View with Lights")

from bpy.app.handlers import persistent
import mathutils  
import math
  
# Fonction persistante pour suivre la vue 3D
@persistent
def z_anatomy_load_post(scene=None):
    def refresh():
        # Parcourir toutes les zones de l'écran pour trouver les zones de vue 3D
        for area in bpy.context.screen.areas:
            if area.type == 'VIEW_3D':
                area3D = area
                viewport_orientation = area3D.spaces[0].region_3d.view_rotation
                
                # Parcourir tous les objets visibles
                for obj in bpy.context.visible_objects:
                    # Vérifier si l'objet a un suffixe spécifique ou se termine par '...'
                    # we only consider ".t" labels, because mirrored ones (".s") have negative scaled parent and they do not align properly with the viewport because of this
                    if ".t" in obj.name or obj.name.endswith('...'):
                        # Si la rotation de l'objet ne correspond pas à l'orientation de la vue, la mettre à jour
                        if obj.rotation_quaternion != viewport_orientation:
                            if obj.rotation_mode != 'QUATERNION':
                                obj.rotation_mode = 'QUATERNION'
                            obj.rotation_quaternion = viewport_orientation

                    # Si l'objet a 'always_show' dans son nom et est masqué, le rendre visible
                    if 'always_show' in obj.name and obj.hide_get():
                        obj.hide_set(False)
                        obj.hide_viewport = False
        return 0.0165

    # Enregistrer la fonction de rafraîchissement pour qu'elle soit appelée périodiquement
    bpy.app.timers.register(refresh, first_interval=0.01)

    # S'abonner aux messages RNA pour les changements d'objets actifs et d'espaces de travail
    bpy.msgbus.subscribe_rna(
        key=subscribe_to,
        owner=owner,
        args=(),
        notify=msgbus_callback,
    )

    bpy.msgbus.subscribe_rna(
        key=(bpy.types.Window, "workspace"),
        owner=owner,
        args=(),
        notify=msgbus_callback_workspace,
    )

    # Utiliser un timer pour obtenir le contexte correct après l'enregistrement
    #bpy.app.timers.register(only_once, first_interval=0.01)
    
# Fonction pour synchroniser la sélection avec l'éditeur de texte
def msgbus_callback_workspace(*args):
    workspace = bpy.context.workspace
    if workspace.name == "Biomechanics":
        bpy.context.window.view_layer = bpy.data.scenes['Scene'].view_layers["Biomechanics"]
        bpy.context.view_layer.objects.active = bpy.data.objects['Armature']
        bpy.ops.object.mode_set(mode='POSE')
    elif workspace.name == "Anatomy":
        bpy.ops.object.mode_set(mode='OBJECT')
        bpy.context.window.view_layer = bpy.data.scenes['Scene'].view_layers["Anatomy"]

# Fonction pour synchroniser la sélection avec l'éditeur de texte
def msgbus_callback(*args):
    active_object = bpy.context.active_object
    if not active_object or not active_object.data: return

    basename, _ = clean_name(active_object.data.name)
    text_editor_area = None
    for area in bpy.context.screen.areas:
        if area.type == "TEXT_EDITOR":
            text_editor_area = area
            break

    if text_editor_area and basename in bpy.data.texts:
        text_editor_area.spaces[0].text = bpy.data.texts[basename]
        text_editor_area.spaces[0].top = 0
        text_editor_area.spaces[0].text.select_set(0, 0, 0, 0)

    # Seulement le label de l'objet visible doit être visible
    if any(x in active_object.name for x in label_elements) or ".st" in active_object.name:
        return

    select_family = family(active_object)
        
    for child in select_family:
        child.hide_set(False)

    for ob in (c for c in bpy.context.visible_objects if c.name.endswith('.g') or (c.name.endswith('-line') and c.parent and c.parent.name.endswith('.g'))):
        ob.hide_set(True)
        
    for ob in (o for o in set(bpy.context.visible_objects) - set(active_object.users_collection[0].all_objects) if o.name.endswith('.g')):
        ob.hide_set(True)

        for child in (c for c in bpy.context.visible_objects if c.parent == ob and c.name.endswith('-line')):
            child.hide_set(True)

    for ob in (o for o in active_object.users_collection[0].all_objects if not o.visible_get() and o.name.endswith('.g')):
        ob.hide_set(False)

        for child in (c for c in bpy.context.scene.objects if c.parent == ob and not c.visible_get() and c.name.endswith('-line')):
            child.hide_set(False)

    #manage labels visiblity (do not display ".s" label, i.e mirrored ones with parent with negative scale, they do not align properly with viewport)
    for ob in (o for o in bpy.context.visible_objects if any(x in o.name for x in label_elements)):
        if (not ob.parent == active_object and not ob in select_family) or ".s" in ob.name:
            ob.hide_set(True)
            for child in ob.children:
                child.hide_set(True)
            #manage the cat case with inverted hierarchy between line and text label
            #hide the line of the label (it is parent of the label for the cat)
            if ".s" in ob.name:
                if ".i" in ob.parent.name or ".j" in ob.parent.name:
                    ob.parent.hide_set(True)

    for ob in (o for o in bpy.context.visible_objects if o.name.endswith('...')):
        if ob == active_object:
            for child in ob.children:
                child.hide_set(False)
        else:
            for child in ob.children:
                child.hide_set(True)

    if active_object.name.endswith('.g'):
        bpy.ops.object.select_grouped('INVOKE_DEFAULT', type='CHILDREN_RECURSIVE')
        active_object.select_set(True)

subscribe_to = bpy.types.LayerObjects, "active"
owner = object()

classes = (
    OBJECT_OT_hide_wrapper,
    OBJECT_OT_hide_view_clear_wrapper,
    OBJECT_OT_local_view_wrapper,
    ZANATOMY_PT_display_panel
)

font_info = {
    "handler": None,
}

# Fonction pour nettoyer les noms en enlevant les suffixes spécifiques
def clean_name(name):
    for ending in ('.r', '.l', '.t', '.st', '.r.t', '.l.t', '.r-line', '.l-line', '.g', '-line', '.o', '.e', ''):
        if ending == '':
            return name, ending
        elif name.endswith(ending):
            clean_name = name[:-len(ending)]
            return clean_name, ending

import blf

# Fonction de rappel pour dessiner le nom de l'objet actif dans la vue 3D
def draw_callback_px(self, context):
    context = bpy.context
    if not hasattr(context, 'area') or not context.object: return

    font_size = 24
    blf.color(0, 255, 255, 255, 255)
    blf.size(0, font_size)#, 72)
    blf.position(0, 55, context.area.height-70 - font_size, 0)
    blf.draw(0, f'{clean_name(context.object.name)[0]}')

def register():
    z_anatomy_load_post()
    bpy.app.handlers.load_post.append(z_anatomy_load_post)

    for c in classes:
        bpy.utils.register_class(c)
    bpy.types.Scene.z_display_enable_addon = bpy.props.BoolProperty(
        default=True, name="Enable Add-on", description="Enable or disable the add-on functionality"
    )
    
    font_info["handler"] = bpy.types.SpaceView3D.draw_handler_add(
        draw_callback_px, (None, None), 'WINDOW', 'POST_PIXEL')

def unregister():
    for c in classes:
        bpy.utils.unregister_class(c)
    del bpy.types.Scene.z_display_enable_addon
    
    bpy.app.handlers.load_post.remove(z_anatomy_load_post)

    bpy.msgbus.clear_by_owner(owner)

    bpy.types.SpaceView3D.draw_handler_remove(font_info["handler"], 'WINDOW')

if __name__ == "__main__":
    register()