multi-user/utils.py

import logging
import sys
from uuid import uuid4
import json
import os

import bpy
import mathutils

from . import draw, environment
from .libs import dump_anything

# TODO: replace hardcoded values...
BPY_TYPES = {'Image': 'images', 'Texture': 'textures', 'Material': 'materials', 'GreasePencil': 'grease_pencils', 'Curve': 'curves', 'Collection': 'collections', 'Mesh': 'meshes', 'Object': 'objects',
             'Scene': 'scenes', 'Light': 'lights', 'SunLight': 'lights', 'SpotLight': 'lights', 'AreaLight': 'lights', 'PointLight': 'lights', 'Camera': 'cameras', 'Action': 'actions', 'Armature': 'armatures'}

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# UTILITY FUNCTIONS


def revers(d):
    l = []
    for i in d:
        l.append(i)

    return l[::-1]


def refresh_window():
    import bpy

    bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)


def get_armature_edition_context(armature):

    override = {}
    # Set correct area
    for area in bpy.data.window_managers[0].windows[0].screen.areas:
        if area.type == 'VIEW_3D':
            override = bpy.context.copy()
            override['area'] = area
            break

    # Set correct armature settings
    override['window'] = bpy.data.window_managers[0].windows[0]
    override['screen'] = bpy.data.window_managers[0].windows[0].screen
    override['mode'] = 'EDIT_ARMATURE'
    override['active_object'] = armature
    override['selected_objects'] = [armature]

    for o in bpy.data.objects:
        if o.data == armature:
            override['edit_object'] = o

            break

    return override


def get_selected_objects(scene):
    selected_objects = []
    for obj in scene.objects:
        if obj.select_get():
            selected_objects.append(obj.name)

    return selected_objects


#    LOAD HELPERS

def load_dict(src_dict, target):
    try:
        for item in src_dict:
            # attr =
            setattr(target, item, src_dict[item])

    except Exception as e:
        logger.error(e)
        pass


def load(key, value):
    target = resolve_bpy_path(key)
    target_type = key.split('/')[0]

    logger.debug("load{}, {}".format(target_type, key))
    if value == "None":
        return

    if target_type == 'Object':
        load_object(target=target, data=value,
                    create=True)
    elif target_type == 'Image':
        load_image(target=target, data=value)
    elif target_type == 'Mesh':
        load_mesh(target=target, data=value,
                  create=True)
    elif target_type == 'Collection':
        load_collection(target=target, data=value,
                        create=True)
    elif target_type == 'Material':
        load_material(target=target, data=value,
                      create=True)
    elif target_type == 'GreasePencil':
        load_gpencil(target=target, data=value,
                     create=True)
    elif target_type == 'Scene':
        load_scene(target=target, data=value,
                   create=True)
    elif 'Light' in target_type:
        load_light(target=target, data=value,
                   create=True)
    elif target_type == 'Camera':
        load_default(target=target, data=value,
                     create=True, type=target_type)
    elif target_type == 'Armature':
        load_armature(target=target, data=value,
                      create=True)
    elif target_type == 'Curve':
        load_curve(target=target, data=value,
                   create=True)
    elif target_type == 'Client':
        load_client(key.split('/')[1], value)


def resolve_bpy_path(path):
    """
    Get bpy property value from path
    """
    item = None

    try:
        path = path.split('/')
        item = getattr(bpy.data, BPY_TYPES[path[0]])[path[1]]

    except:
        pass

    return item


def load_client(client=None, data=None):
    C = bpy.context
    D = bpy.data
    net_settings = C.window_manager.session

    if client and data:
        if net_settings.enable_presence:
            draw.renderer.draw_client(data)
            draw.renderer.draw_client_selected_objects(data)


def load_image(target=None, data=None):
    try:
        if not target:
            image = bpy.data.images.new(
                name=data['name'],
                width=data['size'][0],
                height=data['size'][1]
            )
        else:
            image = target

        img_name = "{}.png".format(image.name)

        logger.info("updating {} cache file".format(image.name))
        img_path = os.path.join(environment.CACHE_DIR, img_name)

        file = open(img_path, 'wb')
        file.write(data["pixels"])
        file.close()

        image.source = 'FILE'
        image.filepath = img_path
        # dump_anything.load(target, data)
    except Exception as e:
        logger.error(e)


def load_armature(target=None, data=None, create=False):
    file = "cache_{}.json".format(data['name'])
    context = bpy.context

    if not target:
        target = bpy.data.armatures.new(data['name'])

        dump_anything.load(target, data)

        with open(file, 'w') as fp:
            json.dump(data, fp)
            fp.close()

        target.id = data['id']
    else:
        # Construct a correct execution context
        file = "cache_{}.json".format(target.name)

        with open(file, 'r') as fp:
            data = json.load(fp)

            if data:
                ob = None
                for o in bpy.data.objects:
                    if o.data == target:
                        ob = o
                if ob:
                    bpy.context.view_layer.objects.active = ob
                    bpy.ops.object.mode_set(mode='EDIT', toggle=False)
                    for eb in data['edit_bones']:
                        if eb in target.edit_bones.keys():
                            # Update the bone
                            pass
                        else:
                            # Add new edit bone and load it

                            target_new_b = target.edit_bones.new[eb]
                            dump_anything.load(target_new_b, data['bones'][eb])

                        logger.debug(eb)

                    bpy.ops.object.mode_set(mode='OBJECT', toggle=False)
            fp.close()
            import os
            os.remove(file)


def load_mesh(target=None, data=None, create=False):
    import bmesh

    if not target or not target.is_editmode:
        # 1 - LOAD GEOMETRY
        mesh_buffer = bmesh.new()

        for i in data["verts"]:
            v = mesh_buffer.verts.new(data["verts"][i]["co"])

        mesh_buffer.verts.ensure_lookup_table()

        for i in data["edges"]:
            verts = mesh_buffer.verts
            v1 = data["edges"][i]["verts"][0]
            v2 = data["edges"][i]["verts"][1]
            mesh_buffer.edges.new([verts[v1], verts[v2]])

        for p in data["faces"]:
            verts = []
            for v in data["faces"][p]["verts"]:
                verts.append(mesh_buffer.verts[v])

            if len(verts) > 0:
                f = mesh_buffer.faces.new(verts)

                uv_layer = mesh_buffer.loops.layers.uv.verify()

                f.material_index = data["faces"][p]['material_index']

                # UV loading
                for i, loop in enumerate(f.loops):
                    loop_uv = loop[uv_layer]
                    loop_uv.uv = data["faces"][p]["uv"][i]

        if target is None and create:
            target = bpy.data.meshes.new(data["name"])

        mesh_buffer.to_mesh(target)

        # mesh_buffer.from_mesh(target)

        # 2 - LOAD METADATA

        # uv's
        for uv_layer in data['uv_layers']:
            target.uv_layers.new(name=uv_layer)

        bevel_layer = mesh_buffer.verts.layers.bevel_weight.verify()
        skin_layer = mesh_buffer.verts.layers.skin.verify()

        # for face in mesh_buffer.faces:

        #     # Face metadata
        #     for loop in face.loops:
        #         loop_uv = loop[uv_layer]
        #         loop_uv.uv = data['faces'][face.index]["uv"]

        dump_anything.load(target, data)

        # 3 - LOAD MATERIAL SLOTS
        material_to_load = []
        material_to_load = revers(data["materials"])
        target.materials.clear()
        # SLots
        i = 0

        for m in data["material_list"]:
            target.materials.append(bpy.data.materials[m])

        target.id = data['id']
    else:
        logger.error("Mesh can't be loaded")


def load_object(target=None, data=None, create=False):
    try:
        if target is None and create:
            pointer = None

            # Object specific constructor...
            if data["data"] in bpy.data.meshes.keys():
                pointer = bpy.data.meshes[data["data"]]
            elif data["data"] in bpy.data.lights.keys():
                pointer = bpy.data.lights[data["data"]]
            elif data["data"] in bpy.data.cameras.keys():
                pointer = bpy.data.cameras[data["data"]]
            elif data["data"] in bpy.data.curves.keys():
                pointer = bpy.data.curves[data["data"]]
            elif data["data"] in bpy.data.armatures.keys():
                pointer = bpy.data.armatures[data["data"]]
            elif data["data"] in bpy.data.grease_pencils.keys():
                pointer = bpy.data.grease_pencils[data["data"]]
            elif data["data"] in bpy.data.curves.keys():
                pointer = bpy.data.curves[data["data"]]

            target = bpy.data.objects.new(data["name"], pointer)

        # Load other meshes metadata
        # dump_anything.load(target, data)

        target.matrix_world = mathutils.Matrix(data["matrix_world"])

        target.id = data['id']

        client = bpy.context.window_manager.session.username

        # Load modifiers
        if hasattr(target,'modifiers'):
            for local_modifier in target.modifiers:
                if local_modifier.name not in data['modifiers']:
                    target.modifiers.remove(local_modifier)
            for modifier in data['modifiers']:
                target_modifier = target.modifiers.get(modifier)

                if not target_modifier:
                    target_modifier = target.modifiers.new(data['modifiers'][modifier]['name'],data['modifiers'][modifier]['type'])
               
                dump_anything.load(target_modifier, data['modifiers'][modifier])



        if target.id == client or target.id == "Common":
            target.hide_select = False
        else:
            target.hide_select = True

    except Exception as e:
        logger.error("Object {} loading error: {} ".format(data["name"], e))


def load_curve(target=None, data=None, create=False):
    try:
        if target is None and create:
            target = bpy.data.curves.new(data["name"], 'CURVE')

        dump_anything.load(target, data)

        target.splines.clear()
        # load splines
        for spline in data['splines']:
            # Update existing..
            # if spline in target.splines.keys():

            new_spline = target.splines.new(data['splines'][spline]['type'])
            dump_anything.load(new_spline, data['splines'][spline])

            # Load curve geometry data
            for bezier_point_index in data['splines'][spline]["bezier_points"]:
                new_spline.bezier_points.add(1)
                dump_anything.load(
                    new_spline.bezier_points[bezier_point_index], data['splines'][spline]["bezier_points"][bezier_point_index])

            for point_index in data['splines'][spline]["points"]:
                new_spline.points.add(1)
                dump_anything.load(
                    new_spline.points[point_index], data['splines'][spline]["points"][point_index])
        target.id = data['id']
    except Exception as e:
        logger.error("curve loading error: {}".format(e))


def load_collection(target=None, data=None, create=False):
    try:
        if target is None and create:
            target = bpy.data.collections.new(data["name"])

        # Load other meshes metadata
        # dump_anything.load(target, data)

        # link objects
        for object in data["objects"]:
            if object not in target.objects.keys():
                target.objects.link(bpy.data.objects[object])

        for object in target.objects.keys():
            if object not in data["objects"]:
                target.objects.unlink(bpy.data.objects[object])

        # Link childrens
        for collection in data["children"]:
            if collection not in target.children.keys():
                if bpy.data.collections.find(collection) == -1:
                    target.children.link(
                        bpy.data.collections[collection])
                else:
                    logger.debug(target.name)

        for collection in target.children.keys():
            if collection not in data["children"]:
                target.collection.children.unlink(
                    bpy.data.collections[collection])

        target.id = data['id']

        client = bpy.context.window_manager.session.username

        if target.id == client or target.id == "Common":
            target.hide_select = False
        else:
            target.hide_select = True

    except Exception as e:
        logger.error("Collection loading error: {}".format(e))


def load_scene(target=None, data=None, create=False):
    try:
        if target is None and create:
            target = bpy.data.scenes.new(data["name"])

        # Load other meshes metadata
        dump_anything.load(target, data)

        # Load master collection
        for object in data["collection"]["objects"]:
            if object not in target.collection.objects.keys():
                target.collection.objects.link(bpy.data.objects[object])

        for object in target.collection.objects.keys():
            if object not in data["collection"]["objects"]:
                target.collection.objects.unlink(bpy.data.objects[object])

        # load collections
        # TODO: Recursive link
        logger.debug("check for scene childs")
        for collection in data["collection"]["children"]:
            logger.debug(collection)
            if collection not in target.collection.children.keys():
                target.collection.children.link(
                    bpy.data.collections[collection])

        logger.debug("check for scene child to remove")
        for collection in target.collection.children.keys():
            if collection not in data["collection"]["children"]:
                target.collection.children.unlink(
                    bpy.data.collections[collection])

        target.id = data['id']
        # Load annotation
        # if data["grease_pencil"]:
        #     target.grease_pencil = bpy.data.grease_pencils[data["grease_pencil"]["name"]]
        # else:
        #     target.grease_pencil = None

    except Exception as e:
        logger.error("Scene loading error: {}".format(e))


def load_material(target=None, data=None, create=False):
    try:
        if target is None:
            target = bpy.data.materials.new(data["name"])

        if data['is_grease_pencil']:
            if not target.is_grease_pencil:
                bpy.data.materials.create_gpencil_data(target)

            dump_anything.load(target.grease_pencil, data['grease_pencil'])

            load_dict(data['grease_pencil'], target.grease_pencil)

        elif data["use_nodes"]:
            if target.node_tree is None:
                target.use_nodes = True

            target.node_tree.nodes.clear()

            for node in data["node_tree"]["nodes"]:
                # fix None node tree error

                index = target.node_tree.nodes.find(node)

                if index is -1:
                    node_type = data["node_tree"]["nodes"][node]["bl_idname"]

                    target.node_tree.nodes.new(type=node_type)

                dump_anything.load(
                    target.node_tree.nodes[index], data["node_tree"]["nodes"][node])

                if data["node_tree"]["nodes"][node]['type'] == 'TEX_IMAGE':
                    target.node_tree.nodes[index].image = bpy.data.images[data["node_tree"]
                                                                          ["nodes"][node]['image']['name']]

                for input in data["node_tree"]["nodes"][node]["inputs"]:

                    try:
                        if hasattr(target.node_tree.nodes[index].inputs[input], "default_value"):
                            target.node_tree.nodes[index].inputs[input].default_value = data[
                                "node_tree"]["nodes"][node]["inputs"][input]["default_value"]
                    except Exception as e:
                        logger.error("Fail loading {} node value from {} ({}) ".format(
                            target.name, target.node_tree.nodes[index].inputs[input].default_value, e))
                        continue

            # Load nodes links
            target.node_tree.links.clear()

            for link in data["node_tree"]["links"]:
                current_link = data["node_tree"]["links"][link]
                input_socket = target.node_tree.nodes[current_link['to_node']
                                                      ['name']].inputs[current_link['to_socket']['name']]
                output_socket = target.node_tree.nodes[current_link['from_node']
                                                       ['name']].outputs[current_link['from_socket']['name']]

                target.node_tree.links.new(input_socket, output_socket)

        # Load other meshes metadata
        # dump_anything.load(target, data)

        target.id = data['id']

    except Exception as e:
        logger.error("Material loading error: {}".format(e))


def load_gpencil_layer(target=None, data=None, create=False):

    dump_anything.load(target, data)

    for frame in data["frames"]:
        try:
            tframe = target.frames[frame]
        except:
            tframe = target.frames.new(frame)
        dump_anything.load(tframe, data["frames"][frame])
        for stroke in data["frames"][frame]["strokes"]:
            try:
                tstroke = tframe.strokes[stroke]
            except:
                tstroke = tframe.strokes.new()
            dump_anything.load(
                tstroke, data["frames"][frame]["strokes"][stroke])

            for point in data["frames"][frame]["strokes"][stroke]["points"]:
                p = data["frames"][frame]["strokes"][stroke]["points"][point]

                tstroke.points.add(1)
                tpoint = tstroke.points[len(tstroke.points)-1]

                dump_anything.load(tpoint, p)


def load_gpencil(target=None, data=None, create=False):
    try:
        if target is None and create:
            target = bpy.data.grease_pencils.new(data["name"])

        for layer in target.layers:
            target.layers.remove(layer)

        if "layers" in data.keys():
            for layer in data["layers"]:
                if layer not in target.layers.keys():
                    gp_layer = target.layers.new(data["layers"][layer]["info"])
                else:
                    gp_layer = target.layers[layer]
                load_gpencil_layer(
                    target=gp_layer, data=data["layers"][layer], create=create)

        dump_anything.load(target, data)

        target.materials.clear()
        if "materials" in data.keys():
            for mat in data['materials']:
                target.materials.append(bpy.data.materials[mat])

        target.id = data['id']
    except:
        logger.error("default loadi\ng error")


def load_light(target=None, data=None, create=False, type=None):
    try:
        if target is None and create:
            target = bpy.data.lights.new(data["name"], data["type"])

        dump_anything.load(target, data)

        target.id = data['id']

        
    except Exception as e:
        logger.error("light loading error: {}".format(e))
        pass


def load_default(target=None, data=None, create=False, type=None):
    try:
        if target is None and create:
            target = getattr(bpy.data, BPY_TYPES[type]).new(data["name"])

        dump_anything.load(target, data)

        target.id = data['id']
    except Exception as e:
        logger.error("default loading error {}".format(e))

# DUMP HELPERS
def dump(key):
    target = resolve_bpy_path(key)
    target_type = key.split('/')[0]
    data = None

    if target_type == 'Image':
        data = {}
        data['pixels'] = dump_image(target)
        dump_datablock_attibutes(target, [], 2, data)
        data = dump_datablock_attibutes(
            target, 
            ["name", 'size', 'height', 'alpha', 'float_buffer', 'filepath', 'source'],
            2,
            data)
    elif target_type == 'Material':
        data = dump_datablock(target, 2)
        if target.node_tree:
            dump_datablock_attibutes(
                target.node_tree, ["nodes", "links"], 5, data['node_tree'])
        elif target.grease_pencil:
            dump_datablock_attibutes(target, ["grease_pencil"], 3, data)
    elif target_type == 'GreasePencil':
        data = dump_datablock(target, 2)
        dump_datablock_attibutes(
            target, ['layers'], 9, data)
    elif target_type == 'Camera':
        data = dump_datablock(target, 1)
    elif 'Light' in target_type:
        data = dump_datablock(target, 3)
    elif target_type == 'Mesh':
        data = dump_datablock(target, 2)
        data = dump_mesh(target, data)
        # dump_datablock_attibutes(
        # target, ['name', 'polygons', 'edges', 'vertices', 'id'], 6, data)

        # Fix material index
        m_list = []
        for m in target.materials:
            m_list.append(m.name)

        data['material_list'] = m_list
    elif target_type == 'Curve':
        data = dump_datablock(target, 1)
        dump_datablock_attibutes(
            target, ['splines'], 5, data)
        # for index, spline in enumerate(target.splines):
        #     data["splines"][index] = dump_datablock_attibutes(target.splines[index],"Curve/{}".format(index), ["bezier_points", "material_index", "points", "order_u", "order_v", "point_count_u", "point_count_v",
        #                                                               "radius_interpolation", "resolution_v", "use_bezier_u", "use_bezier_v", "use_cyclic_u", "use_cyclic_v", "use_endpoint_u", "use_endpoint_v"], 3)
    elif target_type == 'Object':
        data = dump_datablock(target, 1)

        if hasattr(target,'modifiers'):
            dump_datablock_attibutes(
                target, ['modifiers'], 3, data)
    elif target_type == 'Collection':
        data = dump_datablock(target, 4)
    elif target_type == 'Scene':
        data = dump_datablock_attibutes(
            target, ['name', 'collection', 'id', 'camera', 'grease_pencil'], 2)
        dump_datablock_attibutes(
            target, ['collection'], 4, data)

    # elif target_type == 'Armature':
    #     data = dump_datablock(target, 4)

    return data


def dump_datablock(datablock, depth):
    if datablock:
        dumper = dump_anything.Dumper()
        dumper.type_subset = dumper.match_subset_all
        dumper.depth = depth

        datablock_type = datablock.bl_rna.name
        key = "{}/{}".format(datablock_type, datablock.name)
        data = dumper.dump(datablock)

        return data


def dump_datablock_attibutes(datablock=None, attributes=[], depth=1, dickt=None):
    if datablock:
        dumper = dump_anything.Dumper()
        dumper.type_subset = dumper.match_subset_all
        dumper.depth = depth

        datablock_type = datablock.bl_rna.name
        key = "{}/{}".format(datablock_type, datablock.name)

        data = {}

        if dickt:
            data = dickt
        for attr in attributes:
            try:
                data[attr] = dumper.dump(getattr(datablock, attr))
            except:
                pass

        return data


def dump_image(image):
    pixels = None
    if image.source == "GENERATED":
        img_name = "{}.png".format(image.name)

        image.filepath_raw = os.path.join(environment.CACHE_DIR, img_name)
        image.file_format = "PNG"
        image.save()

    if image.source == "FILE":
        image.save()
        file = open(image.filepath_raw, "rb")
        pixels = file.read()
        logger.debug("Reading image file {}".format(image.name))
    else:
        logger.error("image format not supported")
    return pixels


def dump_mesh(mesh, data={}):
    import bmesh

    mesh_data = data
    mesh_buffer = bmesh.new()

    mesh_buffer.from_mesh(mesh)

    uv_layer = mesh_buffer.loops.layers.uv.verify()
    bevel_layer = mesh_buffer.verts.layers.bevel_weight.verify()
    skin_layer = mesh_buffer.verts.layers.skin.verify()

    verts = {}
    for vert in mesh_buffer.verts:
        v = {}
        v["co"] = list(vert.co)

        # vert metadata
        v['bevel'] = vert[bevel_layer]
        # v['skin'] = list(vert[skin_layer])

        verts[str(vert.index)] = v

    mesh_data["verts"] = verts

    edges = {}
    for edge in mesh_buffer.edges:
        e = {}
        e["verts"] = [edge.verts[0].index, edge.verts[1].index]

        # Edge metadata
        e["smooth"] = edge.smooth

        edges[edge.index] = e
    mesh_data["edges"] = edges

    faces = {}
    for face in mesh_buffer.faces:
        f = {}
        fverts = []
        for vert in face.verts:
            fverts.append(vert.index)

        f["verts"] = fverts
        f["material_index"] = face.material_index

        uvs = []
        # Face metadata
        for loop in face.loops:
            loop_uv = loop[uv_layer]

            uvs.append(list(loop_uv.uv))

        f["uv"] = uvs
        faces[face.index] = f

    mesh_data["faces"] = faces

    uv_layers = []
    for uv_layer in mesh.uv_layers:
        uv_layers.append(uv_layer.name)

    mesh_data["uv_layers"] = uv_layers
    return mesh_data


def init_client(key=None):
    client_dict = {}

    C = bpy.context
    Net = C.window_manager.session
    client_dict['uuid'] = str(uuid4())
    client_dict['location'] = [[0, 0, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]]
    client_dict['color'] = [Net.client_color.r,
                            Net.client_color.g, Net.client_color.b, 1]

    client_dict['active_objects'] = get_selected_objects(C.view_layer)

    return client_dict