#!/usr/bin/env python

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

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from settings import ExportMappingFemale, ExportMappingMale, ExportMappingMaleBabyShape, ExportMappingMaleEdited, ExportMappingFemaleEdited, ExportMappingMaleTestShapeEdited, ExportMappingMaleBaseShapeEdited

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

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')

#        track = new_armature.animation_data.nla_tracks.new()
#        track.name = action.name


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

for mapping in [ExportMappingFemale(), ExportMappingMale(), ExportMappingMaleBabyShape(), ExportMappingMaleEdited(), ExportMappingFemaleEdited(), ExportMappingMaleTestShapeEdited(), ExportMappingMaleBaseShapeEdited()]:
    if not os.path.exists(mapping.blend_path):
        print("Skipping mapping: " + mapping.blend_path)
        continue
    print("Processing mapping: from: " + mapping.blend_path + " to: " + mapping.gltf_path)
    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.blend_path, mapping.inner_path),
        directory=os.path.join(mapping.blend_path, mapping.inner_path),
        files=mapping.files)
    print("Append done...")

    prepare_armature(mapping)
    print("Armature done...")

    print("Remove junk...")
    for ob in bpy.data.objects:
        if ob.name.startswith("cs_"):
            bpy.data.objects.remove(ob)
        elif ob.name.startswith("Face") and ob.name != "Face":
            bpy.data.objects.remove(ob)

    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)
    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"
    bpy.ops.wm.save_as_mainfile(filepath=(basepath + "/assets/blender/scripts/" + mapping.outfile))

    os.makedirs(os.path.dirname(mapping.gltf_path), exist_ok=True)
    bpy.ops.export_scene.gltf(filepath=mapping.gltf_path,
            use_selection=False,
            check_existing=False,
#            export_format='GLTF_SEPARATE',
            export_format='GLB',
            export_texture_dir='textures', export_texcoords=True,
            export_animation_mode='NLA_TRACKS',
            export_normals=True,
            export_tangents=True,
            export_materials='EXPORT',
#            export_all_vertex_colors=True,
#            colors_type='SRGB',
#            export_vertex_colors=True,
            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=True,
            export_force_sampling=True,
            export_def_bones=False,
            export_current_frame=False,
            export_morph=True,
            export_morph_animation=False,
            export_morph_normal=True,
            export_morph_tangent=True,
            export_lights=False,
            export_skins=True)
    print("exported to: " + mapping.gltf_path)

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