import bpy
from mathutils import Vector, Color
from . import fn

def add_rlayer(layer_name, scene=None, node_scene=None, location=None, color=None, node_name=None, width=400):
    '''Create a render layer node

    args:
    layer_name (str): Name of the viewlayer
    scene (Scene): Scene holding the viewlayer
    node_scene (Scene): scene where to add render layer
        if None: fallback scene pointed by compo scene property
        if compo scene property is empty: fallback to render scene.
    location (Vector2): Location of the node
    color (tuple): Color of the node
    node_name (str): if not specified, use layer_name
    width (int): width of the node
    '''

    if not node_name:
        node_name = layer_name # 'RL_' +

    if not node_scene:
        node_scene=fn.get_compo_scene(create=False)
    if not node_scene:
        node_scene=fn.get_render_scene(create=False)

    nodes = node_scene.node_tree.nodes

    comp = nodes.get(node_name)
    if comp:
        if comp.layer == node_name:
            return comp
        else:
            # TODO : delete rlayer with bad VL name !
            pass

    comp = nodes.new('CompositorNodeRLayers')
    comp.name = node_name
    comp.scene = scene
    comp.layer = layer_name
    comp.label = layer_name
    if location:
        comp.location = location
    if color:
        comp.color = color

    if width:
        comp.width = width
    comp.show_preview = False
    return comp

# FIXME (Maybe just use "base_path")

def connect_render_layer(rlayer, ng=None, out=None, frame=None,
                         base_path=None, file_slot=None, layer_slot=None
                         ):
    '''Connect a render layer node to a fileoutput
    Return existing or created nodegroup and file output nodes

    Args:
        rlayer (node): the renderlayuer node to connect
        ng (node, optional): Nodegroup to connect to if given
        out (node, optional): Fileoutput node to connect to if given
        frame (node, optional): frame node to use as parent if given

        base_path (str, optional): Template for base path when used with EXR
        file_slot (str, optional): Template for slots when used with EXR
        layer_slot (str, optional): Template for slots when used with Multilayer EXR
        File output template strings keyword :
            {object} : Set object name
            {gplayer} : Set Gp layer name

    Return:
        tuple(node, node) Nodegroup node, file_output node
    '''

    multi_base_path = base_path or '//render/{object}/{object}_'
    base_path = base_path or '//render/{object}'
    file_slot = file_slot or '{gplayer}/{gplayer}_'
    layer_slot = layer_slot or '{gplayer}'


    node_tree = rlayer.id_data # get node_tree from rlayer

    nodes = node_tree.nodes
    links = node_tree.links

    vl_name = rlayer.layer
    if not vl_name or vl_name == 'View Layer':
        print(f'Bad layer for node {rlayer.name}')

    if not ' / ' in vl_name:
        print(f'no slash (" / ") separator in vl_name {vl_name}, should be "obj.name / layer_name"')
        return

    obname, lname = vl_name.split(' / ')
    lname = bpy.path.clean_name(lname)

    if not frame:
        if rlayer.parent:
            frame=rlayer.parent
        else:
            print(f'render_layer has not parent frame: {rlayer.name}')
            frame=None

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


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

    # get set nodegroup from vlayer name
    if not ng:
        ng = nodes.get(ng_name)

    if not ng:
        ngroup = bpy.data.node_groups.get(ng_name)
        # full clear True if exists but not used
        if ngroup and ngroup.users == 0:
            ngroup = None
            fn.clear_nodegroup(ng_name, full_clear=True)

        if not ngroup:
            # delete and recreate ?
            print(f'create nodegroup {ng_name}')
            ngroup = bpy.data.node_groups.new(ng_name, 'CompositorNodeTree')

        ng = fn.create_node('CompositorNodeGroup', tree=node_tree, location=(fn.real_loc(rlayer)[0] + 600, fn.real_loc(rlayer)[1]), width=400) # (rlayer.location[0] + 600, rlayer.location[1])
        if frame:
            ng.parent= frame
        ng.node_tree = ngroup
        ng.name = ngroup.name

    else:
        ngroup = ng.node_tree
    
    if not (ng_in := ngroup.nodes.get('Group Input')):
        ng_in = fn.create_node('NodeGroupInput', tree=ngroup, location=(-600,0))

    if not (ng_out := ngroup.nodes.get('Group Output')):
        ng_out = fn.create_node('NodeGroupOutput', tree=ngroup, location=(600,0))

    # Connect rlayer to nodegroup
    if not rlayer.outputs['Image'].is_linked:
        sockin = ng.inputs.get(vl_name)
        if not sockin:  
            if bpy.app.version < (4,0,0):
                sockin = ng.node_tree.inputs.new('NodeSocketColor', vl_name)
            else:
                sockin = ng.node_tree.interface.new_socket(vl_name, in_out='INPUT', socket_type='NodeSocketColor')
            sockin = ng.inputs[-1]

        links.new(rlayer.outputs['Image'], sockin)

    ## get nodes from frame
    # rl_nodes = [n for n in nodes if n.type == 'R_LAYERS' and n.layer != 'View Layer' and n.parent == frame]

    # auto clean : if an input exists but is not linked and name not exists in rlayers of current frame

    if bpy.app.version < (4,0,0):
        for s in reversed(ng.inputs):
            if not s.is_linked: # and not any(x.layer == s.name for x in rl_nodes)
                print(f'Removing unlinked input {s.name}')
                ng.inputs.remove(s)
    else:
        g_inputs = [s for s in ngroup.interface.items_tree if s.in_out == 'INPUT']
        # g_outputs = [s for s in ngroup.interface.items_tree if s.in_out == 'OUTPUT']
        for i in range(len(ng.inputs))[::-1]:
            if not ng.inputs[i].is_linked:
                print(f'Removing unlinked input {ng.inputs[i].name}')
                ngroup.interface.remove(g_inputs[i])

    ## get nodes from linked NG inputs ??? maybe more clear...
    # rl_nodes = [s.links[0].from_node for s in ng.inputs if s.links and s.links[0].from_node and s.links[0].from_node.type == 'R_LAYERS']

    ## reorder
    fn.reorder_inputs(ng)
    ng.update()

    # CREATE NG outsocket (individual, without taking merge)
    connected = False

    if ng_in.outputs[vl_name].is_linked:
        # check if connect to the other side
        socket = fn.connect_to_group_output(ng_in.outputs[vl_name].links[0].to_node) #if ng_in.outputs[vl_name].links[0].to_node.type == 'ALPHAOVER':
        if socket:
            connected = True
            groupout = ng.outputs.get(socket.name)

    ng.update()

    if not connected:
        # add AA and connect
        aa = fn.create_aa_nodegroup(ngroup)# fn.new_aa_node(ngroup)
        groupout = ng.outputs.get(vl_name)
        if not groupout:
            if bpy.app.version < (4,0,0):
                ng.node_tree.outputs.new('NodeSocketColor', vl_name) # assigning direcly doesn't link well
            else:
                ng.node_tree.interface.new_socket(vl_name, in_out='OUTPUT', socket_type='NodeSocketColor')

            groupout = ng.outputs[-1]

        # print('ng_out.inputs.get(vl_name): ', ng_out.inputs.get(vl_name))
        # ng_in.outputs[vl_name]
        ngroup.links.new(ng_in.outputs[vl_name], aa.inputs[0]) # node_tree
        ngroup.links.new(aa.outputs[0], ng_out.inputs[vl_name]) # node_tree


        ## Get use_aa prop from render scene ?
        # rd_scn = fn.get_render_scene(create=False)
        # if rd_scn:
        #     aa.mute = rd_scn.gp_render_settings.use_aa # mute if native AA is used
        scene = next((s for s in bpy.data.scenes if s.node_tree == node_tree), None)
        if scene:
            print(f'set AA from scene {scene.name}')
            aa.mute = scene.gp_render_settings.use_aa # mute if native AA is used
        else:
            print('/!\ Scene could not be found to define if internal AA should be muted')

    fn.reorganise_NG_nodegroup(ng) # decorative

    # Clean outputs
    if bpy.app.version < (4,0,0):
        for o in reversed(ngroup.outputs):
            if not o.name in [o.name for o in ngroup.inputs]:
                print(f'removing group output {o.name} (name not exists in group inputs)')
                ngroup.outputs.remove(o)
    else:
        n_outputs = [s for s in ngroup.interface.items_tree if s.in_out == 'OUTPUT']
        n_inputs = [s for s in ngroup.interface.items_tree if s.in_out == 'INPUT']
        for o in reversed(n_outputs):
            if not o.name in [o.name for o in n_inputs]:
                print(f'Removing group output {o.name} (name not exists in group inputs)')
                ngroup.interface.remove(o)

    ng.update()
    # reorder output to match inputs
    fn.reorder_outputs(ng)

    ng.update()

    # Clear : delete orphan nodes that are not connected from ng_in
    for n in reversed(ngroup.nodes):
        if n.type in ('GROUP_INPUT', 'GROUP_OUTPUT'):
            continue
        if not fn.connect_to_group_input(n) and not fn.connect_to_group_output(n): # is disconnected from both side
            ngroup.nodes.remove(n)
        # TODO clear nodes that are disconnected from input side ?

    if groupout.links and groupout.links[0].to_node.type == 'OUTPUT_FILE':
        # if already connected to outfile just skip cause user might have customised the name
        return
    
    out_name = f'OUT_{obname}' # or get output from frame
    out_base = bpy.path.clean_name(obname)

    if not out:
        out = nodes.get(out_name)
    
    if not out:
        # color = (0.2,0.3,0.5)
        out = fn.create_node('CompositorNodeOutputFile', tree=node_tree, location=(fn.real_loc(ng)[0]+500, fn.real_loc(ng)[1]+50), width=600) # =(ng.location[0]+600, ng.location[1]+50)
        fn.set_file_output_format(out)

        out.name = out_name
        out.parent = frame
        
        if out.format.file_format == 'OPEN_EXR_MULTILAYER':
            out.base_path = multi_base_path.format(object=out_base, gplayer=lname)
            # out.base_path = f'//render/{out_base}/{out_base}_'
        else:
            out.base_path = base_path.format(object=out_base, gplayer=lname)
            # out.base_path = f'//render/{out_base}'

    if out.format.file_format == 'OPEN_EXR_MULTILAYER':
        # Direct name in multilayer
        # slot_name = lname
        slot_name = layer_slot.format(object=out_base, gplayer=lname)
    else:
        # base_path/ named_folder/image_####
        # slot_name = f'{lname}/{lname}_'
        slot_name = file_slot.format(object=out_base, gplayer=lname)


    ## out_input = out.inputs.get(slot_name) # Ok for non-numbered outputs

    out_input = None
    out_input = fn.get_numbered_output(out, slot_name)

    if not out_input:
        ## Assigning directly does not work
        # out.file_slots.new(file_slot.format(object=out_base, gplayer=lname))
        # out.layer_slots.new(layer_slot.format(object=out_base, gplayer=lname))

        out.file_slots.new('file_slots_temp_name')
        out.layer_slots.new('layer_slots_temp_name')

        fs = out.file_slots[-1]
        fs.path = file_slot.format(object=out_base, gplayer=lname)

        ls = out.layer_slots[-1]
        ls.name = layer_slot.format(object=out_base, gplayer=lname)
        # out.file_slots.new(slot_name)
        out_input = out.inputs[-1] # assigning directly as above doesn't link afterwards
        # print(f'new filouput entry: {out_input}')

    # link to FileOut
    links.new(groupout, out_input)

    # clean fileout
    fn.clear_disconnected(out) # maybe not disconnect ?
    fn.reorder_fileout(out, ng=ng)

    out.update()
    return ng, out


def clamp_color_value(color, clamp_value=0.65) -> Color:
    '''Return Color instance with clamped value component'''
    color = Color(color)
    color.v = min(color.v, clamp_value)
    return color

def get_set_viewlayer_from_gp(ob, l, scene=None, node_scene=None,
                              base_path=None, file_slot=None, layer_slot=None):
    '''setup ouptut from passed gp obj > layer
        scene: scene to set viewlayer
        node_scene: where to add compo node (use scene if not passed)
    
        base_path (str, optional) : File output Base Path template
        file_slot (str, optional) : File output slot template for individual files
        layer_slot (str, optional) : File output slot for multilayer EXR

    Return:
        viewlayer, render-layer node
    '''

    scene = scene or fn.get_render_scene()
    node_scene = node_scene or fn.get_compo_scene() or scene

    # print('Viewlayer Scene:', scene.name) #Dbg
    # print('Compo Scene:', node_scene.name) #Dbg

    ## If not passed, identical to scene holding viewlayers
    if not node_scene.use_nodes:
        node_scene.use_nodes = True
    node_tree = node_scene.node_tree
    nodes = node_tree.nodes

    in_rds = scene.collection.all_objects.get(ob.name)
    if not in_rds:
        # TODO : ? duplicate the object with name with a specific suffix '_renderdupe' to still parse it ?
        ## make single user if its a multiuser object ? maybe let the user do it
        if ob.data.users > 1:
            print(f'/!!\ {ob.name} data has multiple users ! ({ob.data.users})')
            # ob.data = ob.data.copy() # create duplicate (this will also affect the one in original scene !!!)
        scene.collection.objects.link(ob)
        ob.hide_viewport = ob.hide_render = False

        ## set object active in default viewlayer

        # if (avl := scene.view_layers.get('ViewLayer')):
        #     # This select the object in source scene
        #     avl.objects.active = ob
        #     # avl.objects.active.select_set(True)
        nob = scene.collection.objects.get(ob.name)
        if nob:
            nob.select_set(True)

    # Create viewlayer
    vl_name = f'{ob.name} / {l.info}'
    vl = fn.get_view_layer(vl_name, scene=scene)
    vl_name = vl.name

    ## To avoid potential error, transfer compo scene prop to renderscene.
    scn = bpy.context.scene
    if scn.gp_render_settings.node_scene and scn != scene:
        if scn.gp_render_settings.node_scene != scene.gp_render_settings.node_scene:
            print(f'Transfer compo scene target prop to render scene: "{scn.gp_render_settings.node_scene}"')
            scene.gp_render_settings.node_scene = scn.gp_render_settings.node_scene

    # Affect layer to this vl
    l.viewlayer_render = vl_name

    # Check if already exists
    rlayer_list = [n for n in nodes if n.type == 'R_LAYERS' and n.layer == vl_name]

    # Get frame object and their contents
    # dict model : {objname : [layer_nodeA, layer_nodeB,...]}
    frame_dic = {f.label: [n for n in nodes if n.type == 'R_LAYERS' and n.parent and n.parent.name == f.name and '/' in n.layer] # n.layer != 'View Layer'
        for f in nodes if f.type == 'FRAME'}

    # debug print
    # for k,v in frame_dic.items():
    #     print('-', k)
    #     for n in v:
    #         print('---', n.layer)

    if rlayer_list: # rlayer exists
        print(f'{len(rlayer_list)} nodes using {vl_name}')

        # affect only the one within an object frame
        framed_rl = [n for n in rlayer_list if n.parent and n.parent.label == ob.name]
        if framed_rl:
            if len(framed_rl) > 1:
                print(f'! More than one nodes using {vl_name} in a frame ({len(framed_rl)}) !')

            # sort top to bottom and take upper node
            framed_rl.sort(key=lambda x:x.location.y, reverse=True)
            cp = framed_rl[0]
            cp.select = True # select so the user see that it existed
            return vl, cp

    # Returned if existed and OK

    if not ob.name in frame_dic.keys(): # and len(frame_dic[ob.name])
        print(f'\n{ob.name} -> {l.info} (first generation)')
        # frame not exists, add the RL and frame at the very bottom of all render_layers
        # check position of frame type ? all type ?
        all_frames = [n for n in nodes if n.type == 'FRAME']
        # all_rl_x = [n.location.x for n in nodes if n.type == 'R_LAYERS' and n.layer != 'View Layer']
        if all_frames:
            # all_frames.sort(key=lambda x: x.location.y, reverse=True)
            # loc.y -  dim.y
            y_loc = min(fn.get_frame_transform(f, node_tree)[0].y - fn.get_frame_transform(f, node_tree)[1].y for f in all_frames)
            loc = (0, y_loc)
        else:
            loc = (0,0)

        # create frame at new rl position
        frame = nodes.new('NodeFrame')
        frame.label = ob.name
        frame.label_size = 50
        frame.location = (loc[0], loc[1] + 20)

        cp = add_rlayer(vl_name, scene=scene, node_scene=node_scene, location=loc)
        cp.parent = frame
        # use same color as layer
        if fn.has_channel_color(l):
            cp.use_custom_color = True
            cp.color = clamp_color_value(l.channel_color)

        connect_render_layer(cp, frame=frame,
                             base_path=base_path, file_slot=file_slot, layer_slot=layer_slot)

        fn.rearrange_frames(node_tree)

        return vl, cp

    ## If ob already was used

    print(f'\n {ob.name} -> {l.info} (connect to existing)')

    ## object frame exists: get framing and insert
    cp = add_rlayer(vl_name, scene=scene, node_scene=node_scene, location=(0,0))
    if cp.layer != vl_name:
        print(f'problem with {cp}: {cp.layer} != {vl_name}')
        return

    if fn.has_channel_color(l):
        cp.use_custom_color = True
        cp.color = clamp_color_value(l.channel_color)

    frame = next((f for f in nodes if f.type == 'FRAME' and f.label == ob.name), None)
    rl_nodes = frame_dic[frame.label] # get nodes from
    if rl_nodes:
        # get nodes order to insert
        rl_nodes.sort(key=lambda n: fn.real_loc(n).y, reverse=True) # descending
        top_loc = fn.real_loc(rl_nodes[0])
    else:
        print('!! gen_viewlayer: No Render layers nodes !!')
        top_loc = fn.get_frame_transform(frame, node_tree) - 60

    # cp.location = (top_loc[0], top_loc[1] + 100) # temp location to adjust x loc


    # List of layer names in nodes order
    rl_names = [n.layer.split(' / ')[1] for n in rl_nodes] # Get True layer name from rl
    
    ## Names with the right order WITH the new layer included
    # names = [lay.info for lay in ob.data.layers if lay.info in rl_names or lay == l] # <- Character limit problem
    ## Consider viewlayer name length of 63 char max
    char_limit = 63 - len(ob.name + ' / ')
    names = [lay.info[:char_limit] for lay in ob.data.layers if lay.info[:char_limit] in rl_names or lay == l]
    
    rl_nodes.append(cp)

    # filter by getting index(layer_name)
    cp.parent = frame
    rl_nodes.sort(key=lambda x : names.index(x.layer.split(' / ')[1])) # Sort True layer name from rl

    offset = 0
    # print(f'number of nodes in frame: {len(rl_nodes)}')
    ref_node = rl_nodes[0]

    # print('ref_node: ', ref_node.name, ref_node.location)
    for n in rl_nodes:
        # set x loc from first node in list (maybe use leftmost ?)
        n.location = Vector((fn.real_loc(ref_node)[0], top_loc[1] - offset)) - n.parent.location
        offset += 180
        n.update()

    # reorder render layers nodes within frame
    connect_render_layer(cp, frame=frame,
                        base_path=base_path, file_slot=file_slot, layer_slot=layer_slot)

    # re-arrange all frames (since the offset probably overlapped)
    fn.rearrange_frames(node_tree)

    return vl, cp

def export_gp_objects(oblist, exclude_list=[], scene=None, node_scene=None, base_path=None, file_slot=None, layer_slot=None):
    # Skip layer containing element in exclude list
    if not isinstance(oblist, list):
        oblist = [oblist]
    if isinstance(exclude_list, str):
        exclude_list = [p.strip() for p in exclude_list.split(',')]

    # print('exclude_list: ', exclude_list)

    for ob in oblist:
        for l in ob.data.layers:
            # if l.hide:
            #     continue
            if l.hide or l.opacity == 0 or any(x + '_' in l.info for x in exclude_list):
                print(f'Exclude export: {ob.name} : {l.info}')
                # Assign "exclude" layer
                l.viewlayer_render = fn.get_view_layer('exclude', scene=scene).name
                continue

            get_set_viewlayer_from_gp(ob, l, scene=scene, node_scene=node_scene,
                                      base_path=base_path, file_slot=file_slot, layer_slot=layer_slot)

def add_layer_to_render(ob, node_scene=None, base_path=None, file_slot=None, layer_slot=None):
    '''Send GP object to render layer

    return a tuple with report message'''

    # ob = ob or bpy.context.object
    layer = ob.data.layers.active
    if not layer:
        return ('ERROR', 'No active layer')

    node_scene = fn.get_compo_scene(scene_name=node_scene, create=True)

    ct = 0
    # send scene ?
    hidden = 0
    for l in ob.data.layers:
        if not l.select:
            if not l.viewlayer_render:
                # TODO : need to link, can raise error if object is not linked in Render scene yet
                l.viewlayer_render = fn.get_view_layer('exclude').name
            continue

        get_set_viewlayer_from_gp(ob, l, node_scene=node_scene, base_path=base_path, file_slot=file_slot, layer_slot=layer_slot)

        if l.hide:
            hidden += 1
        ct += 1

    if hidden:
        return ('WARNING', f'{hidden}/{ct} layers are hidden!')

    else:
        return ('INFO', f'{ct} layer(s) added')


def add_object_to_render(mode='ALL', scene='', node_scene='', base_path=None, file_slot=None, layer_slot=None):
    context = bpy.context

    if scene:
        scn = fn.get_render_scene(scene)
    else:
        scn = fn.get_render_scene()

    if node_scene:
        node_scn = fn.get_compo_scene(scene_name=node_scene, create=True)
        if not node_scn:
            return ('ERROR', f'/!\ Node Scene "{node_scene}" not found ! Abort "Add object to Render" !')
    else:
        # if not passed add in render scene
        node_scn = scn

    excludes = [] # ['MA', 'IN'] # Get list dynamically
    if mode == 'SELECTED':
        export_gp_objects([o for o in context.selected_objects if o.type == 'GPENCIL'], 
                                     exclude_list=excludes, scene=scn, node_scene=node_scn,
                                     base_path=base_path, file_slot=file_slot, layer_slot=layer_slot)

    elif mode == 'ALL':
        export_gp_objects([o for o in context.scene.objects if o.type == 'GPENCIL' and not o.hide_get() and fn.is_valid_name(o.name)], 
                                     exclude_list=excludes, scene=scn, node_scene=node_scn,
                                     base_path=base_path, file_slot=file_slot, layer_slot=layer_slot)