import bpy
import numpy as np
from time import perf_counter, time
from mathutils import Vector, Matrix

from gp_interpolate.utils import (get_gp_draw_plane,
                                  is_animated,
                                  create_plane,
                                  following_keys,
                                  place_object_to_ref_facing_cam,
                                  empty_at,
                                  attr_set)


class GP_OT_parent_layer(bpy.types.Operator):
    bl_idname = "gp.parent_layer"
    bl_label = "Parent Layer"
    bl_description = 'Parent active layer to object or bone\
                    \nBake intermediate parent object to compensate GP offset and moves'
    bl_options = {'REGISTER', 'UNDO'} 

    @classmethod
    def poll(cls, context):
        if context.object\
        and context.object.type == 'GPENCIL'\
        and context.object.data.layers.active:
            return True
        cls.poll_message_set("Need a Grease pencil object with an active layer")
        return False
    
    direct_parent : bpy.props.BoolProperty(default=False, options={'SKIP_SAVE'})

    def parent_and_compensate(self, obj, bone=None):
        '''If bone is None, parent layer only to object'''
        self.lay.parent = obj
        if bone:
            self.lay.parent_type = 'BONE'
            self.lay.parent_bone = bone.name
            print(f'Parent to {self.lay.parent.name} > bone {bone.name}')
        
        ## Offset parent
        if bone:
            bone_mat = self.lay.parent.matrix_world @ bone.matrix
            self.lay.matrix_inverse = bone_mat.inverted() @ self.gp.matrix_world
            
        else:
            print(f'Parent to {self.lay.parent.name}')
            self.lay.matrix_inverse = self.lay.parent.matrix_world.inverted() @ self.gp.matrix_world
        
        return {'FINISHED'}

    def execute(self, context):
        settings = context.scene.gp_interpo_settings
        scn = bpy.context.scene
        
        self.gp = context.object
        self.lay = self.gp.data.layers.active

        if not settings.target_rig:
            self.report({'ERROR'}, 'No object Selected')
            return {'CANCELLED'}

        self.bone = None
        if settings.target_rig.type == 'ARMATURE':
            if not settings.target_bone:
                self.report({'ERROR'}, 'No Bone Selected')
                return {'CANCELLED'}
            self.bone = settings.target_rig.pose.bones.get(settings.target_bone)
            if not self.bone:
                self.report({'ERROR'}, f'{settings.target_bone} not found in armature {settings.target_rig.name}')
                return {'CANCELLED'}


        if self.direct_parent:
            return self.parent_and_compensate(settings.target_rig, settings.target_bone)
            return {'FINISHED'}

        
        ## Place on a 2d intermediate
        parent_name = f'ref_{self.gp.name}_{self.lay.info}'
        parent_collec_name = 'gp_parents'
        included_cols = [c.name for c in self.gp.users_collection]
        included_cols.append(parent_collec_name)

        ## Ensure collection and parent exists
        # Get/create collection
        col = bpy.data.collections.get(parent_collec_name)
        if not col:
            col = bpy.data.collections.new(parent_collec_name)
        if col.name not in bpy.context.scene.collection.children:
            bpy.context.scene.collection.children.link(col)
            ## hide parent collection ?
            # col.hide_viewport = True
        
        # Get/create empty parent
        parent = bpy.data.objects.get(parent_name)
        if not parent:
            parent = empty_at(name=parent_name, size=0.1, show_name=True)
            parent.select_set(False)

        ## Prepare context manager
        store_list = [
            (context.tool_settings, 'use_keyframe_insert_auto', False),
            ]
        
        for vlc in context.view_layer.layer_collection.children:
            store_list.append(
                (vlc, 'exclude', vlc.name not in included_cols),
                # (),
                # (vlc, 'hide_viewport', vlc.name not in included_cols), # viewport viz
            )
        
        ## Offset at curent frame to compensate for object, GP (and GP layer ?) transformations
        ## If GP object is animated, animate parent obj to compensate
        
        ## How to smart-test if self.gp is animated in space ?
        # if not is_animated(self.gp):

        with attr_set(store_list):
            parent.animation_data_clear() # Kill animation data to redo it
            for i in range(scn.frame_start, scn.frame_end+1):
                scn.frame_set(i)
                place_object_to_ref_facing_cam(parent,
                    ref_ob=settings.target_rig,
                    bone=self.bone,
                    set_rotation=settings.use_bone_rotation)
                parent.keyframe_insert('location')
        
                parent.keyframe_insert('rotation_euler')
                parent.keyframe_insert('scale')
        

        self.parent_and_compensate(parent)
        print('Parent Done')
        return {'FINISHED'}


classes = (
    GP_OT_parent_layer,
)

def register():
    for c in classes:
        bpy.utils.register_class(c)
    
        
def unregister():
    for c in reversed(classes):
        bpy.utils.unregister_class(c)