import bpy
import re
from math import isclose
from . import fn
from . import gen_vlayer

# TODO : make a merge compatible with already merged nodegroup (or even other node type)
# --> need to delete/mute AA internal node

def merge_layers(rlayers, obname=None, active=None, disconnect=True, color=None):

    print(f'Merging {len(rlayers)} layers')
    print('->', [r.layer for r in rlayers])
    print()

    if not rlayers:
        return ('ERROR', 'No render layer sent to merge')

    # get node group
    # ng = rlayers[0].outputs[0].links[0].to_node

    # sort RL descending    
    rlayers.sort(key=lambda n: fn.real_loc(n).y, reverse=True)

    node_tree = rlayers[0].id_data
    nodes = node_tree.nodes
    links = node_tree.links

    if active:
        vl_name = active.layer
    else:
        vl_name = rlayers[-1].layer # -1 : bottom node == upper layer

    if ' / ' in vl_name:
        obname, lname = vl_name.split(' / ')
        lname = bpy.path.clean_name(lname)
        base_path = f'//render/{bpy.path.clean_name(obname)}'
        slot_name = f'{lname}/{lname}_'
    else:
        # directly use full vlname for both base output and subfolder ?? (or return error)
        obname = lname = bpy.path.clean_name(vl_name)
        base_path = f'//render/'
        slot_name = f'{lname}/{lname}_'
    

    # change colors of those nodes
    disconnected_groups = []
    if not color: 
        color = fn.random_color()
    for n in rlayers:
        n.use_custom_color = True
        n.color = color

        if disconnect:
            if n.outputs[0].is_linked:
                for lnk in reversed(n.outputs[0].links):
                    if lnk.to_node.name.startswith('NG_'):
                        disconnected_groups.append(lnk.to_node)
                        links.remove(lnk)
    
    disconnected_groups = list(set(disconnected_groups))

    ng_name = f'merge_NG_{obname}' # only object name

    ## clear unused nodes groups duplication
    fn.clear_nodegroup(ng_name, full_clear=False)

    ### always create a new nodegroup (nerve call an existing one)

    # need a unique nodegroup name
    # increment name while nodegroup exists
    while bpy.data.node_groups.get(ng_name): # nodes.get(ng_name)
        if not re.search(r'(\d+)$', ng_name):
            ng_name += '_02' # if not ending with a number add _02
        ng_name = re.sub(r'(\d+)(?!.*\d)', lambda x: str(int(x.group(1))+1).zfill(len(x.group(1))), ng_name)        

    # print(f'create merge nodegroup {ng_name}')
    ngroup = bpy.data.node_groups.new(ng_name, 'CompositorNodeTree')
    ng = fn.create_node('CompositorNodeGroup', tree=node_tree, location=(fn.real_loc(rlayers[0]).x + 1900, fn.real_loc(rlayers[0]).y - 200), width=400)
    ng.node_tree = ngroup
    ng.name = ngroup.name

    _ng_in = fn.create_node('NodeGroupInput', tree=ngroup, location=(-600,0))
    _ng_out = fn.create_node('NodeGroupOutput', tree=ngroup, location=(600,0))

    # Create inputs and links to node_group
    for rln in rlayers:
        rln.outputs['Image']
        sockin = ng.inputs.new('NodeSocketColor', rln.layer)
        sockin = ng.inputs[-1]
        links.new(rln.outputs['Image'], sockin)

    fn.nodegroup_merge_inputs(ng.node_tree)
    ng.update()
    # create dedicated fileout

    out = fn.create_node('CompositorNodeOutputFile', tree=node_tree, location=(ng.location[0]+450, ng.location[1]+50), width=600)
    fn.set_file_output_format(out)
    out_name = f'merge_OUT_{vl_name}' # or get output from frame
    out.name = out_name
    out.base_path = base_path
    out.file_slots.new(slot_name)
    links.new(ng.outputs[0], out.inputs[-1])
    
    fn.clear_disconnected(out)
    out.update()

    ## Clear node_group after disconnect
    # for dg in disconnected_groups:
    #     fn.clean_nodegroup_inputs(dg)
    #     # fn.clear_nodegroup_content_if_disconnected(dg.node_tree)

    bpy.context.scene.use_aa = False # trigger fn.scene_aa(toggle=False)

    return ng, out

class GPEXP_OT_merge_viewlayers_to_active(bpy.types.Operator):
    bl_idname = "gp.merge_viewlayers_to_active"
    bl_label = "Merge selected layers view_layers"
    bl_description = "Merge view layers of selected gp layers to on the active one"
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'GPENCIL'

    def execute(self, context):
        ob = bpy.context.object
        # layers = [l for l in ob.data.layers if l.select and not l.hide]
        act = ob.data.layers.active
        layers = [l for l in ob.data.layers if l.select and l != act]

        ## Tested in func
        # rd_scn = bpy.data.scenes.get('Render')
        # if not rd_scn:
        #     self.report({'ERROR'}, 'Viewlayers needs to be generated first!')
        #     return {'CANCELLED'}

        # if not act.viewlayer_render:
        #     self.report({'ERROR'}, f'Active layer {act.info} has no viewlayer assigned')
        #     return {'CANCELLED'}
        
        ret = fn.merge_gplayer_viewlayers(ob, act=act, layers=layers)
        if isinstance(ret, tuple):
            self.report(*ret)
        return {"FINISHED"}


class GPEXP_OT_auto_merge_adjacent_prefix(bpy.types.Operator):
    bl_idname = "gpexp.auto_merge_adjacent_prefix"
    bl_label = "Auto Merge Adjacent Prefix"
    bl_description = "Automatically merge viewlayer and renderlayer of grouped layer prefix"
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'GPENCIL'

    excluded_prefix : bpy.props.StringProperty(
        name='Excluded Prefix', default='GP,RG,PO',
        description='Exclude comma separated prefix from merging viewlayer')
    
    first_name : bpy.props.BoolProperty(name='Merge On Bottom Layer', 
        default=True, 
        description='Keep the viewlayer of the bottom layer in groups, else upper layer')

    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)
        return self.execute(context)

    def draw(self, context):
        layout = self.layout
        layout.label(text='Settings for auto-merge:')
        layout.prop(self, 'excluded_prefix')
        layout.prop(self, 'first_name')

    def execute(self, context):
        prefix_list = [p.strip() for p in self.excluded_prefix.split(',')]
        for ob in [o for o in context.selected_objects if o.type == 'GPENCIL']:
            fn.group_adjacent_layer_prefix_rlayer(ob, excluded_prefix=prefix_list, first_name=self.first_name)
        return {"FINISHED"}

# unused
class GPEXP_OT_merge_selected_dopesheet_layers(bpy.types.Operator):
    bl_idname = "gp.merge_selected_dopesheet_layers"
    bl_label = "Merge selected layers nodes"
    bl_description = "Merge view layers of selected gp layers to a new dedicated file output"
    bl_options = {"REGISTER"}

    @classmethod
    def poll(cls, context):
        return context.object and context.object.type == 'GPENCIL'

    disconnect : bpy.props.BoolProperty(default=True, options={'SKIP_SAVE'})

    def execute(self, context):
        ob = bpy.context.object
        layers = [l for l in ob.data.layers if l.select and not l.hide]
        act = ob.data.layers.active
        # merge_selected_layers() # function to merge from GP dopesheet
        if not act:
            self.report({'ERROR'}, f'An active layer is needed to set merge output name')
            return {"CANCELLED"}

        if len(layers) < 2:
            self.report({'ERROR'}, f'Should select multiple layers for merging')
            return {"CANCELLED"}

        render = bpy.data.scenes.get('Render')
        if render:
            nodes = render.node_tree.nodes
        
        clean_ob_name = bpy.path.clean_name(ob.name)
        rlayers = []
        for l in layers:
            idname = f'{clean_ob_name} / {l.info}'
            rlayer = rl = None
            # check the render layer that have a parent frame
            if not render:
                _vl, rl = gen_vlayer.get_set_viewlayer_from_gp(ob, l)
                render = bpy.data.scenes.get('Render')
                nodes = render.node_tree.nodes

            if not rl:
                rlayer = [n for n in nodes if n.type == 'R_LAYERS' and n.layer == idname and n.parent]
            if not rlayer:
                # send to function to generate the rlayer and connect
                _vl, rl = gen_vlayer.get_set_viewlayer_from_gp(ob, l)
                    
            else:
                rlayer.sort(key=lambda n: n.location.y, reverse=True)
                rl = rlayer[0]
            
            if act == l:
                nodes.active = rl # make it active so the merge use this one

            rlayers.append(rl)
        
        color = None
        if fn.has_channel_color(act): # and bpy.context.preferences.edit.use_anim_channel_group_colors
            color = act.channel_color
        merge_layers(rlayers, disconnect=self.disconnect, color=color)

        return {"FINISHED"}


class GPEXP_OT_merge_selected_viewlayer_nodes(bpy.types.Operator):
    bl_idname = "gp.merge_selected_viewlayer_nodes"
    bl_label = "Merge selected view_layers "
    bl_description = "Merge selected view layers to a new dedicated file output\nDisconnect single output unless using 'keep connect'"
    bl_options = {"REGISTER"}

    disconnect : bpy.props.BoolProperty(default=True, options={'SKIP_SAVE'})

    def execute(self, context):
        if context.scene.name == 'Scene':
            render = bpy.data.scenes.get('Render')
        else:
            render = context.scene

        if not render:
            self.report({'ERROR'}, 'No render scene')
            return {"CANCELLED"}

        nodes = render.node_tree.nodes
        selection = [n for n in nodes if n.select and n.type == 'R_LAYERS']

        if not nodes.active in selection:
            self.report({'ERROR'}, 'The active node not within the render layer selection (used to define out name)')
            return {'CANCELLED'}

        # should be from the same object:
        if not all(selection[0].layer.split('.')[0] == n.layer.split('.')[0] for n in selection):
            print('/!\ Merge -> Not every nodes start with the same object')

        color = None
        if nodes.active.use_custom_color and nodes.active.color:
            color = nodes.active.color
        merge_layers(selection, active=nodes.active, disconnect=self.disconnect, color=color)

        return {"FINISHED"}

classes=(
GPEXP_OT_merge_viewlayers_to_active,
GPEXP_OT_auto_merge_adjacent_prefix,
GPEXP_OT_merge_selected_dopesheet_layers,# unused
GPEXP_OT_merge_selected_viewlayer_nodes,
)

def register(): 
    for cls in classes:
        bpy.utils.register_class(cls)

def unregister():
    for cls in reversed(classes):
        bpy.utils.unregister_class(cls)