import bpy
from mathutils import Vector
from . import fn


def add_rlayer(layer_name, scene=None, location=None, color=None, node_name=None, width=400):
    '''create a render layer node
    if node_name is not specified, use passed layer name
    '''

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

    if not scene:
        scene=bpy.context.scene

    nodes = 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

def connect_render_layer(rlayer, ng=None, out=None, frame=None):
    if bpy.context.scene.name == 'Scene':
        scene = fn.get_render_scene()
    else:
        scene = bpy.context.scene

    nodes = scene.node_tree.nodes
    links = scene.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=scene.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

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

    else:
        ngroup = ng.node_tree
        ng_in = ngroup.nodes.get('Group Input')
        ng_out = ngroup.nodes.get('Group Output')

    # 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
    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 grp unlinked input {s.name}')
            # ng.inputs.remove(s)
            if bpy.app.version < (4,0,0):
                ng.node_tree.inputs.remove(s)
            else:
                ng.node_tree.interface.items_tree.remove(s)

    ## 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

        aa.mute = scene.use_aa # mute if native AA is used


    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

    slot_name = f'{lname}/{lname}_'
    out_name = f'OUT_{obname}' # or get output from frame
    if not out:
        out = nodes.get(out_name)
    if not out:
        # color = (0.2,0.3,0.5)
        out = fn.create_node('CompositorNodeOutputFile', tree=scene.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
        out.base_path = f'//render/{bpy.path.clean_name(obname)}'

    ## 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:
        out.file_slots.new(slot_name)
        out_input = out.inputs[-1] # assigning directly 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 get_set_viewlayer_from_gp(ob, l, scene=None):
    '''setup ouptut from passed gp obj > layer'''
    if not scene:
        if bpy.context.scene.name != 'Scene':
            scene = bpy.context.scene
        else:
            scene = fn.get_render_scene()

    # print('Set viewlayer Scene: ', scene.name)
    node_tree = 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
    # 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 like : {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, location=loc)
        cp.parent = frame
        # use same color as layer
        if fn.has_channel_color(l):
            cp.use_custom_color = True
            cp.color = l.channel_color

        connect_render_layer(cp, frame=frame)

        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, 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 = l.channel_color

    frame = [f for f in nodes if f.type == 'FRAME' and f.label == ob.name][0]
    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:
        top_loc = fn.get_frame_transform(frame[1], 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]

    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)

    # 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):
    # 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 any(x + '_' in l.info for x in exclude_list): # exclude hided ?
                print(f'Exclude export: {ob.name} : {l.info}')
                l.viewlayer_render = fn.get_view_layer('exclude', scene=scene).name # assign "exclude"
                continue

            _vl, _cp = get_set_viewlayer_from_gp(ob, l, scene=scene) # scene=fn.get_render_scene())