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)

    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]

        if not act.viewlayer_render:
            self.report({'ERROR'}, f'Active layer {act.info} has no viewlayer assigned')
            return {'CANCELLED'}

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

        # list layers and viewlayers
        vls = [rd_scn.view_layers.get(l.viewlayer_render) for l in layers
            if l.viewlayer_render and l.viewlayer_render != act.viewlayer_render and rd_scn.view_layers.get(l.viewlayer_render)]

        vl_names = [v.name for v in vls]

        for n in reversed(rd_scn.node_tree.nodes):
            if n.type == 'R_LAYERS' and n.layer in vl_names:
                for lnk in n.outputs[0].links:
                    grp = lnk.to_node
                    if grp.type != 'GROUP':
                        continue
                    if not grp.name.startswith('NG'):
                        continue
                    sockin = lnk.to_socket
                    sockout = grp.outputs.get(sockin.name)
                    if not sockout:
                        continue
                    
                    for grplink in sockout.links:
                        if grplink.to_node.type != 'OUTPUT_FILE':
                            continue
                        fo_socket = grplink.to_socket
                        fo = grplink.to_node
                        fo.file_slots.remove(fo_socket)
                    
                    # remove input and output from group
                    # grp.inputs.remove(sockin) # do not clear inside !!
                    # grp.outputs.remove(sockout) # do not clear inside !!
                    ngroup = grp.node_tree
                    for i in range(len(grp.inputs))[::-1]:
                        if grp.inputs[i].name == sockin.name:
                            ngroup.inputs.remove(ngroup.inputs[i])
                            break
                    for i in range(len(grp.outputs))[::-1]:
                        if grp.outputs[i].name == sockout.name:
                            ngroup.outputs.remove(ngroup.outputs[i])
                            break
                
                # remove render_layer node
                rd_scn.node_tree.nodes.remove(n)
        
        # assign view layer from active to selected
        for l in layers:
            l.viewlayer_render = act.viewlayer_render

        ## delete unused_vl

        # used_vl_name = [n.layer for n in rd_scn.node_tree.nodes if n.type == 'R_LAYERS' and n.layer]        
        for vl in vls:
            rd_scn.view_layers.remove(vl)            
            # if not vl.name in used_vl_name:
                # rd_scn.view_layers.remove(vl)

        return {"FINISHED"}

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):
        render = bpy.data.scenes.get('Render')
        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_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)