#!/usr/bin/env python

import os, time
import bpy
from math import pi
import glob
from mathutils import Vector, Matrix
from math import radians, pi
import json

class ExportMappingFemale:
    char_blend_path = "assets/blender/vroid1-female.blend"
    cloth_blend_paths = ["assets/blender/female-coat2.blend"]
    gltf_path = "godot/clothes/female-coat.gltf"
    inner_path = "Object"
    objs = ["skeleton", "body", "privates", "ref-bodyskirt-noimp", "ref-bodydress-noimp", "ref-bodycap-noimp", "ref-bodyshoes-noimp", "ref-dress-noimp", "ref-topskirt-noimp", "ref-skirt4-noimp", "ref-skirt3-noimp"]
    objs_remove = []
    armature_name = "skeleton"
    outfile = "tmp-female-cloth.blend"
    default_action = 'default'
    sex = "female"
    def __init__(self):
        self.files = []
        for fobj in self.objs:
            self.files.append({"name": fobj})

class ExportMappingMale:
    char_blend_path = "assets/blender/vroid1-man-animate.blend"
    cloth_blend_paths = ["assets/blender/male-coat2.blend"]
    gltf_path = "godot/clothes/male-coat.gltf"
    inner_path = "Object"
    objs = ["pxy", "body", "ref-topbottom-noimp", "ref-bodytopbottom-noimp"]
    objs_remove = []
    armature_name = "pxy"
    outfile = "tmp-male-cloth.blend"
    default_action = 'default'
    sex = "male"
    def __init__(self):
        self.files = []
        for fobj in self.objs:
            self.files.append({"name": fobj})


basepath = os.getcwd()
def check_bone(bname):
    ok = True
    baddie = ["ctrl_", "mch_", "MCH_"]
    for bd in baddie:
        if bname.startswith(bd):
            ok = False
            break
    return ok

def clamp_angle_deg(angle, min_angle_deg, max_angle_deg):
    min_angle = radians(min_angle_deg)
    max_angle = radians(max_angle_deg)
    if angle < min_angle:
        angle = min_angle
    if angle > max_angle:
        angle = max_angle
    return angle
def angle_to_linear(angle, divider):
    if angle < 0.0:
        return angle / divider
    else:
        return 0.0
def angle_to_linear_x(bone, angle):
    skel = bpy.data.objects["skeleton_orig"]
    left_base = "ctrl_base_upperleg.L.001"
    right_base = "ctrl_base_upperleg.R.001"
    base = ""
    if base == "":
        for e in ["_R", ".R"]:
            if bone.name.endswith(e):
                base = right_base
                break
    if base == "":
        for e in ["_L", ".L"]:
            if bone.name.endswith(e):
                base = left_base
                break
    if base == "":
        for e in ["_R.", ".R."]:
            if bone.name.find(e) >= 0:
                base = right_base
                break
    if base == "":
        for e in ["_L.", ".L."]:
            if bone.name.find(e) >= 0:
                base = left_base
                break
    mul = skel.pose.bones[base]["to_linear_x_base"]
    offset = skel.pose.bones[base]["angle_offset"]
#    print("bone: ", bone.name, "base: ", base, "mul: ", mul)
#    print("angle: ", angle, " angle_offset: ", offset, " angle_sum: ", angle + offset)
    print("offset: ", mul * (angle + offset), "bone: ", base, "angle: ", angle)
    return (angle + offset) * mul
def extra_linear(angle, offset):
    ret = 0.0
    offt = offset * angle * 2.0 / -radians(-90)
    if angle * 2.0 < -radians(65):
        if angle * 2.0 > -radians(65):
            ret += offset
        else:
            ret += offt
    return ret

def prepare_armature(mapping):
    print("Preparing...")
    bpy.ops.object.armature_add(enter_editmode=False)
    new_armature = bpy.context.object
    orig_armature = bpy.data.objects[mapping.armature_name]
    armature_name = orig_armature.name
    orig_armature.name = orig_armature.name + "_orig"
    new_armature.name = armature_name
    queue = []
    if new_armature.animation_data is None:
        new_armature.animation_data_create()
    bpy.context.view_layer.objects.active = new_armature
    bpy.ops.object.mode_set(mode='EDIT')
    for b in new_armature.data.edit_bones:
        new_armature.data.edit_bones.remove(b)
    bpy.context.view_layer.objects.active = orig_armature
    bpy.ops.object.mode_set(mode='EDIT')
    for b in orig_armature.data.edit_bones:
        print(b.name)
        if b.parent is None:
            queue.append(b.name)
    print("Copying bones...")
    while len(queue) > 0:
        item = queue.pop(0)
        print(item)
        itemb = orig_armature.data.edit_bones[item]
        if not itemb.use_deform and not check_bone(item):
            continue
        for cb in orig_armature.data.edit_bones:
            if cb.parent == itemb:
                queue.append(cb.name)
        nb = new_armature.data.edit_bones.new(item)
        nb.name = item
        nb.head = itemb.head
        nb.tail = itemb.tail
        nb.matrix = itemb.matrix
        nb.use_deform = itemb.use_deform
        if itemb.parent is not None:
            ok = True
            pname = itemb.parent.name
            while not check_bone(pname):
                bparent = itemb.parent.parent
                if bparent is None:
                    ok = False
                    break
                pname = bparent.name
            if ok:
                nb.parent = new_armature.data.edit_bones[itemb.parent.name]
            else:
                nb.parent = None
        else:
            nb.parent = None
        nb.use_connect = itemb.use_connect
#    drivers_data = new_armature.animation_data.drivers
    print("Creating constraints...")
    bpy.context.view_layer.objects.active = new_armature
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.context.view_layer.objects.active = orig_armature
    bpy.ops.object.mode_set(mode='OBJECT')
    for b in new_armature.pose.bones:
        print(b.name)
        c = b.constraints.new(type='COPY_TRANSFORMS')
        c.target = orig_armature
        c.subtarget = b.name
    for obj in bpy.data.objects:
        if obj.parent == orig_armature:
            obj.parent = new_armature
            for mod in obj.modifiers:
                if mod.type == 'ARMATURE':
                    mod.object = new_armature
    print("Baking actions...")
    bpy.context.view_layer.objects.active = new_armature
    bpy.ops.object.mode_set(mode='POSE')
    for track in orig_armature.animation_data.nla_tracks:
        print(track.name)
        for s in track.strips:
            action = s.action
            print(action.name)
            orig_armature.animation_data.action = action
            new_armature.animation_data.action = None
            bpy.context.view_layer.objects.active = new_armature
            firstFrame = int(s.action_frame_start)
            lastFrame = int(s.action_frame_end)
            bpy.ops.nla.bake(frame_start=firstFrame, frame_end=lastFrame, step=5, only_selected=False, visual_keying=True, clear_constraints=False, clean_curves=True, use_current_action=False, bake_types={'POSE'})
            aname = orig_armature.animation_data.action.name
            orig_armature.animation_data.action.name = "bake_" + aname
            new_armature.animation_data.action.name = aname
            track = new_armature.animation_data.nla_tracks.new()
            track.name = aname
            track.strips.new(track.name, int(new_armature.animation_data.action.frame_range[0]), new_armature.animation_data.action)
            track.mute = True
            track.lock = True
    print("Removing constraints...")
    for b in new_armature.pose.bones:
        for c in b.constraints:
            b.constraints.remove(c)
    new_armature.animation_data.action = bpy.data.actions[mapping.default_action]
    bpy.context.view_layer.objects.active = new_armature
    bpy.ops.object.mode_set(mode='OBJECT')


obj_data = {}
for mapping in [ExportMappingFemale(), ExportMappingMale()]:
    print("Initializing...")
    bpy.ops.wm.read_homefile(use_empty=True)
    print("Preparing driver setup...")
    bpy.app.driver_namespace["clamp_angle_deg"] = clamp_angle_deg
    bpy.app.driver_namespace["angle_to_linear"] = angle_to_linear
    bpy.app.driver_namespace["extra_linear"] = extra_linear
    bpy.app.driver_namespace["angle_to_linear_x"] = angle_to_linear_x
    print("Driver setup done...")

    bpy.ops.wm.append(
        filepath=os.path.join(mapping.char_blend_path, mapping.inner_path),
        directory=os.path.join(mapping.char_blend_path, mapping.inner_path),
        files=mapping.files)
    for obj in bpy.data.objects:
        if obj.name.startswith("bone-"):
            bpy.data.objects.remove(obj)
    print("Append character done...")
    prepare_armature(mapping)
    print("Armature done...")
    print("Removing original armature and actions...")
    orig_arm = bpy.data.objects[mapping.armature_name + '_orig']
    bpy.data.objects.remove(orig_arm)
    for act in bpy.data.actions:
        if act.name.startswith("bake_"):
            act.name = act.name + "-noimp"
    for act in bpy.data.actions:
        if act.name.startswith("bake_"):
            bpy.data.actions.remove(act)
    print("Removing original armature and actions done...")

    for filepath in mapping.cloth_blend_paths:
        with bpy.data.libraries.load(filepath) as (data_from, data_to):
            data_to.objects = [ob for ob in data_from.objects if not ob.endswith('noimp')]
        for obj in data_to.objects:
            if obj.type == 'MESH':
                bpy.context.scene.collection.objects.link(obj)
                obj.parent = bpy.data.objects[mapping.armature_name]
                arm_mod = obj.modifiers.new('armature', 'ARMATURE')
                arm_mod.object = bpy.data.objects[mapping.armature_name]
                for vg in obj.vertex_groups:
                    obj.vertex_groups.remove(vg)
            bpy.ops.object.select_all(action='DESELECT')
            bpy.data.objects[obj.name].select_set(True)
            src_name = "body"
            params = ""
            otype = "nearest"
            distance = 0.1
            if "src_obj" in obj:
                obname = obj["src_obj"]
                if obname.find(";") >= 0:
                    src_name, params = obname.split(";");
                    vpar, vdist = params.split("=")
                    otype = vpar
                    distance = float(vdist)
                else:
                    src_name = obname
                if src_name not in bpy.data.objects:
                    src_name = "ref-" + src_name + "-noimp"
            keywords = []
            keywords.append(mapping.sex)
            if "slot" in obj:
                keywords.append(obj["slot"])
            if "keywords" in obj:
                kw = obj["keywords"].split(",")
                for xtk in kw:
                    keywords.append(xtk.strip())
            if "state" in obj:
                keywords.append(obj["state"].strip())
            else:
                if obj.name.endswith("-damaged"):
                    keywords.append("damaged")
                elif obj.name.endswith("-revealing"):
                    keywords.append("revealing")
                else:
                    keywords.append("normal")
            if "speciality" in obj:
                keywords.append(obj["speciality"].strip())
            else:
                keywords.append("common")
            keywords.append(obj.name.replace("-damaged", "").replace("-revealing", ""))
            obj_data[obj.name] = {}
            obj_data[obj.name]["keywords"] = keywords
            sobj = src_name
            vert_mapping = "NEAREST"
            use_distance = True
            bpy.data.objects[src_name].select_set(True)
            bpy.context.view_layer.objects.active = bpy.data.objects[obj.name]
            bpy.ops.paint.weight_paint_toggle()
            if otype == "nearest":
                vert_mapping = "NEAREST"
            elif otype == "poly":
                vert_mapping = "POLYINTERP_NEAREST"
            if distance <= 0.0001:
                use_distance = False
            bpy.ops.object.data_transfer(use_reverse_transfer=True,
                    data_type='VGROUP_WEIGHTS', use_create=True,
                    vert_mapping = vert_mapping,
                    layers_select_src='NAME',
                    max_distance = distance,
                    use_max_distance = use_distance,
                    layers_select_dst='ALL')
            bpy.ops.paint.weight_paint_toggle()
    print("Append clothes done...")
    for obj in bpy.data.objects:
        if obj.name.endswith("noimp"):
            bpy.data.objects.remove(obj)
        elif obj.name == "body":
            obj.name = "body-noimp"
        else:
            bpy.context.view_layer.objects.active = obj
            for modifier in obj.modifiers:
                if modifier.type != 'ARMATURE':
                    bpy.ops.object.modifier_apply(modifier = modifier.name)


#    print("Removing original armature and actions...")
#    orig_arm = bpy.data.objects[mapping.armature_name + '_orig']
#    bpy.data.objects.remove(orig_arm)
#    for act in bpy.data.actions:
#        if act.name.startswith("bake_"):
#            bpy.data.actions.remove(act)
#    for obj in bpy.data.objects:
#        if obj.type == 'MESH':
#            if not obj.name in mapping.objs and obj.parent is None:
#                if not obj.name.endswith("-noimp"):
#                    obj.name = obj.name + "-noimp"
    for obj in bpy.data.objects:
        if obj.name in mapping.objs_remove:
            bpy.data.objects.remove(obj)

    bpy.ops.wm.save_as_mainfile(filepath=(basepath + "/assets/blender/scripts/" + mapping.outfile))

    bpy.ops.export_scene.gltf(filepath=mapping.gltf_path,
            use_selection=False,
            check_existing=False,
            export_format='GLTF_SEPARATE',
            export_texture_dir='textures', export_texcoords=True,
            export_normals=True,
            export_tangents=False,
            export_materials='EXPORT',
            export_colors=True,
            use_mesh_edges=False,
            use_mesh_vertices=False,
            export_cameras=False,
            use_visible=False,
            use_renderable=False,
            export_yup=True,
            export_animations=False,
            export_force_sampling=True,
            export_def_bones=False,
            export_current_frame=False,
            export_morph=True,
            export_morph_normal=True,
            export_lights=False)

with open(basepath + "/godot/clothes/clothes.json", "w") as write_file:
    json.dump(obj_data, write_file, indent=8)

bpy.ops.wm.read_homefile(use_empty=True)
time.sleep(2)
bpy.ops.wm.quit_blender()