import bpy

def set_params(src, tgt, mod_to_node=True, org_modifier=None):
    # mod to node: est-ce qu'on copie les valeurs d'un modifier a une node, ou l'inverse
    if mod_to_node: #  syntax for node and modifier are slightly different
        tree = src.node_group.interface.items_tree
    else:
        tree = src.node_tree.interface.items_tree

    for param in tree:
        if param.socket_type == 'NodeSocketGeometry':
            continue
        if param.in_out == 'OUTPUT':
            continue

        # seulement en extract mode, src est une node donc on check si des parametres sont dans le modifier
        node_link_value = get_node_link_value(src, param.name, org_modifier)
        identifier = tree.get(param.name).identifier

        if mod_to_node:
            tgt.inputs[identifier].default_value = src[identifier]

        else:
            if node_link_value:
                tgt[identifier] = node_link_value
            else:
                tgt[identifier] = src.inputs[identifier].default_value

def set_group_inputs(target, objects, group):
    mod = target.modifiers[0]
    node_dct = {} # used for cleanup
    for key, inp in get_node_inputs(objects).items():
        # add the socket to the node group / modifier pannel
        sock = group.interface.new_socket(inp["label"],in_out="INPUT",socket_type=inp["socket"])
        mod[sock.identifier] = inp["data"]

        # inspect all nodes and add a group input node when that socket is used
        for node in parse_nodes(objects):
            for param in node.node_tree.interface.items_tree:
                nkey = get_input_socket_key(node, param)
                if not nkey: continue
                if nkey == key:
                    input_node = add_input_node(group, node, param.identifier, sock)

                    # on efface les parametres par defaut qui sont connectes
                    # ( pour ne pas garder de trace des anciens params / collections /object)
                    node.inputs[param.identifier].default_value = None

                    # add to dict for cleanup
                    if not node in node_dct.keys():
                        node_dct[node] = [input_node]
                    else:
                        node_dct[node].append(input_node)

        # on refait la meme chose pour les object info nodes car leur syntaxe est un peu differente
        for node in parse_nodes(objects, type = "OBJECT_INFO"):
            nkey = get_input_socket_key(node, param)
            if not nkey: continue
            if nkey == key:
                input_node = add_input_node(group, node, 'Object', sock)
                node.inputs['Object'].default_value = None

                # add to dict for cleanup
                if not node in node_dct.keys():
                    node_dct[node] = [input_node]
                else:
                    node_dct[node].append(input_node)

    # cleanup tree
    for input_nodes in node_dct.values():
        for offset, input_node in enumerate(input_nodes):
            input_node.location[1] += 50 * offset
            hide_sockets(input_node)

def get_node_link_value(node, param_name, org_mod):
    if not org_mod:
        return
    # est-ce que le param est connecté a une autre node ?
    if not node.inputs[param_name].links: 
        return
    socket_id = node.inputs[param_name].links[0].from_socket.identifier

    return org_mod[socket_id]

def get_geo_socket(node, input=True):
    if node.type != "GROUP": 
        return('Geometry')
    for param in node.node_tree.interface.items_tree:
        if param.socket_type != 'NodeSocketGeometry':
            continue
        if input and param.in_out == 'INPUT' : return param.identifier
        if not input and param.in_out == 'OUTPUT' : return param.identifier
    return None
def get_input_socket_key(node, param):
    if node.type == "GROUP":
        if param.in_out != 'INPUT':
            return False
        if not param.socket_type in ['NodeSocketObject','NodeSocketCollection']:
            return False
        tgt = node.inputs[param.identifier].default_value

        if not tgt:
            return False
        return f"{param.socket_type[10:][:3]} {tgt.name}"

    if node.type == "OBJECT_INFO":
        tgt = node.inputs['Object'].default_value
        if not tgt:
            return False
        return f"Object {tgt.name}"

def get_node_inputs(combined_nodes):
    # inputs["Col COL.name"] = {name = COL.name, data = COL, socket = "COLLECTION"}
    # inputs["Obj OBJ.name"] = {name = OBJ.name, data = OBJ, socket = "OBJECT"}
    inputs = {}
    for node in parse_nodes(combined_nodes):
        for param in node.node_tree.interface.items_tree:
            key = get_input_socket_key(node, param)
            if not key: 
                continue
            tgt = node.inputs[param.identifier].default_value
            inputs[key] = {'name': tgt.name, 'data': tgt, 'label': param.name , 'socket': param.socket_type}

    for node in parse_nodes(combined_nodes, type = "OBJECT_INFO"):
        key = get_input_socket_key(node, None)
        if not key:
            continue
        tgt = node.inputs['Object'].default_value
        inputs[key] = {'name': tgt.name, 'data': tgt, 'label': 'Source OB' , 'socket': "NodeSocketObject"}

    return inputs

def get_node_bounds(objects, mode=0, x=0, y=0):
    min_x = min_y = 10000000
    max_x = max_y = 0

    for ob in objects:
        for node in ob:
            co = node.location
            min_x = min(co[0],min_x)
            max_x = max(co[0],max_x)

            min_y = min(co[1],min_y)
            max_y = max(co[1],max_y)
    if mode == 0:
        return([max_x+x, (min_y+max_y)/2 ])

def get_collection(name):
    scn = bpy.context.scene
    col = None
    link = False

    # look for existing
    for c in bpy.data.collections:
        if c.name == name: col = c

    # create if needed
    if not col: col = bpy.data.collections.new(name)

    # link to scene if needed
    for c in scn.collection.children_recursive:
        if c.name == col.name: link = True
    if not link:
        scn.collection.children.link(col)
    return col

def get_mod_frames(grp):
    frames = []
    for node in grp.nodes:
        if node.type == "FRAME": frames.append(node)
    return(frames)

def get_frame_childrens(frame):
    childrens = []
    locs = {}

    for node in frame.id_data.nodes:
        if node.parent == frame:
            locs[node.location[0]] = node

    # sort nodes by their x location, je sais c'est mal ecris...
    entries = sorted(locs.keys())
    childrens = [locs[x] for x in entries]
    return childrens

def parse_nodes(combined_nodes, type = "GROUP"):
    nodes = []
    for frame in combined_nodes:
        for node in frame:
            if node.type == type: nodes.append(node)
    return nodes

def copy_source_ob(ob, col):
    # est-ce que l'objet a des data ? si oui on cree une copie ,
    # si non on renvois None
    new_ob = None
    if ob.type == "MESH" and len(ob.data.vertices) > 0:
        new_ob = ob.copy()
        new_ob.data = ob.data.copy()

    if ob.type == "CURVE" and len(ob.data.splines) > 0:
        new_ob = ob.copy()
        new_ob.data = ob.data.copy()

    if new_ob:
        for mod in new_ob.modifiers:
            new_ob.modifiers.remove(mod)

    if new_ob and col:
        col.objects.link(new_ob)
    return new_ob

def hide_sockets(node,collapse = True):
    for socket in node.outputs:
        if not socket.links: 
            socket.hide = True
    for socket in node.inputs:
        if not socket.links: 
            socket.hide = True
    if collapse: 
        node.hide = True

def add_input_node(group, node, param_id, socket):
    group_input_node = group.nodes.new('NodeGroupInput')
    group_input_node.location = node.location
    group_input_node.location[1] += 70
    group_input_node.label = socket.name
    group.links.new(group_input_node.outputs[socket.identifier],
                    node.inputs[param_id])
    return(group_input_node)

def add_material_node(ob, group, nodes):
    if not ob.material_slots:
        return nodes
    if not ob.material_slots[0].material:
        return nodes
    last_node = nodes[-1:][0]
    node = group.nodes.new('GeometryNodeSetMaterial')
    node.inputs['Material'].default_value = ob.material_slots[0].material
    node.location = last_node.location
    node.location[0] += 300
    nodes.append(node)
    return nodes

def join_nodes(group, nodes):
    prev = None
    for i , node in enumerate(nodes):
        if not prev:
            prev = node
            continue
        geo_in = get_geo_socket(node)
        geo_out = get_geo_socket(prev, input = False)

        if not geo_in or not geo_out:
            continue
        group.links.new(prev.outputs[geo_out], node.inputs[geo_in])
        prev = node

def frame_nodes(group, nodes, ob):
    nd = group.nodes.new('NodeFrame')
    # frame = nodes.new(type='NodeFrame')
    for n in nodes: 
        n.parent = nd

    nd.label = ob.name

def combine_ob(ob, group, y=0, col=None):
    nodes = []

    # object info node
    nd = group.nodes.new('GeometryNodeObjectInfo')
    nd.location[0] -= 300
    nd.location[1] = y * 800
    nd.transform_space = "RELATIVE"
    nd.inputs['Object'].default_value = copy_source_ob(ob, col) # si l'objet contient des data on crée une copie
    nodes.append(nd)

    # ob modifiers
    for x,md in enumerate(ob.modifiers):
        if md.type != "NODES" :
            print(abordage)
        if md.node_group == group:
            continue

        nd = group.nodes.new('GeometryNodeGroup')
        nd.label = md.name
        nd.width = 230
        nd.location[0] = x * 300
        nd.location[1] = y * 800

        nd.node_tree = md.node_group
        set_params(md, nd)
        nodes.append(nd)
    nodes = add_material_node(ob, group, nodes)
    join_nodes(group, nodes)
    frame_nodes(group, nodes, ob)
    return nodes

def gen_target_ob(group, col=None):
    ob = gen_empty_ob(group.name, col=col)
    mod = ob.modifiers.new(group.name, "NODES")
    mod.node_group = group

    ob.show_name = True
    bpy.context.view_layer.objects.active = ob
    return(ob)

def gen_empty_ob(name, col=None):
    scn = bpy.context.scene
    ob = bpy.data.objects.new(name, object_data=bpy.data.meshes.new(name))

    ob.data.materials.append(None)
    ob.material_slots[0].link = 'OBJECT'

    if not col:
        scn.collection.objects.link(ob)
    else:
        col.objects.link(ob)
    return(ob)

def assign_modifiers(ob, frame, org_modifier):
    for node in get_frame_childrens(frame):
        if node.type != "GROUP":
            continue
        mod = ob.modifiers.new(node.label, "NODES")
        mod.node_group = node.node_tree
        mod.show_expanded = False
        set_params(node, mod, mod_to_node=False, org_modifier=org_modifier)
        mod.node_group.interface_update(bpy.context)

def join_branches(objects, group):
    # join all trees and add an output node
    join = group.nodes.new('GeometryNodeJoinGeometry')
    out = group.nodes.new('NodeGroupOutput')
    out_sock = group.interface.new_socket("Geometry",in_out="OUTPUT",socket_type="NodeSocketGeometry")

    loc = get_node_bounds(objects, x=500)
    join.location = loc
    out.location = loc
    out.location[0] += 700

    for ob in objects:
        node = ob[-1:][0]
        group.links.new(node.outputs[get_geo_socket(node, input=False)],
                        join.inputs[get_geo_socket(join)])

    group.links.new(join.outputs[get_geo_socket(join, input=False)],
                    out.inputs[out_sock.identifier])

def gen_extracted_ob(name, frame, col, mod):
    ob = None
    for node in get_frame_childrens(frame):
        if node.type != "OBJECT_INFO":
            continue
        target = get_node_link_value(node, 'Object', mod)

        if target:
            ob = target.copy()
            ob.data = ob.data.copy()
            col.objects.link(ob)
    if not ob: ob = gen_empty_ob(name , col = col)

    # assign material
    for node in get_frame_childrens(frame):
        if node.type != "SET_MATERIAL":
            continue
        ob.material_slots[0].material = node.inputs['Material'].default_value
    return ob

def combine_objects(objs):
    name = f"NODEGROUP_combined"
    col = get_collection(name)
    group = bpy.data.node_groups.new(name=name, type='GeometryNodeTree')

    objects = []
    for y , ob in enumerate(objs):
        objects.append(combine_ob(ob, group, y=y, col=col))

    target = gen_target_ob(group, col = col)
    set_group_inputs(target, objects, group)
    join_branches(objects, group)

def extract_objects(object):
    mod = object.modifiers[0]
    grp = mod.node_group
    col = get_collection(grp.name)

    for frame in get_mod_frames(grp):
        name = f"{object.name} {frame.label}"
        ob = gen_extracted_ob(name, frame, col, mod)
        assign_modifiers(ob, frame, mod)

#combine_objects(bpy.context.selected_objects)
#extract_objects(bpy.context.active_object)
"""
TODO: extract copier les transform de l'objet original ...
OK ! combine: si un objet a un materiel on cree un node set material en fin de liste
OK ! extract: si on trouve un noeud set material on l'assigne
OK ! extract: si un socket est connecté on recup la valeur de la connection plutot que du socket
OK ! combine: effacer tout les parametres par defaut qui sont connectes ( pour ne pas garder de trace des anciens params / collections /object)
OK ! combine: mettre tout les objets crees/copiés dans une collection
OK ! combine: si un objet source a des mesh/curve data, on en fait une copie, remove les modifiers, et assign dans le object node source
OK ! combine: si un noeud object info n'est pas vide on expose son contenu dans l'interface
OK ! extract: si un noeud object info n'est pas vide on duplique son contenu au lieu de creer un mesh vide
"""