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

if bpy.app.version[0] == 3:
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + "/mixamo/3.6")
if bpy.app.version[0] == 4:
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + "/mixamo/4.3")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from vrm import rename
if bpy.app.version[0] == 3:
    from mixamo import mixamo_rig
from mixamo.lib.armature import *
from settings import VRMDataFemale, VRMDataMale, VRMDataMaleBabyShape, basepath
from geometry import tris2quads

imports = [VRMDataFemale(), VRMDataMale(), VRMDataMaleBabyShape()]

class mkrig:
    ik_arms = True
    ik_legs = True
    def __init__(self):
        mixamo_rig._make_rig(self)
        mixamo_rig.remove_temp_objects()
        mixamo_rig.clean_scene()
def get_anim(filepath, root_bone_name="Root", hip_bone_name="mixamorig:Hips", remove_prefix=False, name_prefix="mixamorig:",  insert_root=False, delete_armatures=False):
    old_objs = set(bpy.context.scene.objects)
    try:
        import_armature(filepath, root_bone_name, hip_bone_name, remove_prefix, name_prefix, insert_root, delete_armatures)
        imported_objects = set(bpy.context.scene.objects) - old_objs
        if delete_armatures and num_files > 1:
            deleteArmature(imported_objects)
    except Exception as e:
        log.error("[Mixamo Root] ERROR get_all_anims raised %s when processing %s" % (str(e), filepath))
    bpy.context.scene.frame_start = 0
    bpy.ops.object.mode_set(mode='OBJECT')
def create_root_bone(anim_obj):
    bpy.context.view_layer.objects.active = anim_obj
    bpy.ops.object.mode_set(mode='EDIT')
    root_bone = anim_obj.data.edit_bones.new("Root")
    root_bone.tail.y = 0.5
    anim_obj.data.edit_bones["mixamorig:Hips"].parent = anim_obj.data.edit_bones["Root"]
    bpy.ops.object.mode_set(mode='OBJECT')
def hips_to_root(anim_obj):
    for fc in anim_obj.animation_data.action.fcurves:
        if "mixamorig:Hips" in fc.data_path:
           if "location" in fc.data_path:
               print(fc, fc.data_path, fc.array_index)
               if fc.array_index in [0, 2]:
                   fc.data_path = fc.data_path.replace("mixamorig:Hips", "Root")
def hips_to_root(anim_obj):
    for fc in anim_obj.animation_data.action.fcurves:
        if "mixamorig:Hips" in fc.data_path:
           if "location" in fc.data_path:
               print(fc, fc.data_path, fc.array_index)
               if fc.array_index in [0, 2]:
                   fc.data_path = fc.data_path.replace("mixamorig:Hips", "Root")
def hips_to_root2(anim_obj):
    for fc in anim_obj.animation_data.action.fcurves:
        if "Ctrl_Hips" in fc.data_path:
           if "location" in fc.data_path:
               print(fc, fc.data_path, fc.array_index)
               if fc.array_index in [0, 2]:
                   data_path = fc.data_path[:]
                   fcr = anim_obj.animation_data.action.fcurves.new(data_path = data_path.replace("Ctrl_Hips", "Root"), index = fc.array_index)
                   keys = [0] * (2 * len(fc.keyframe_points))
                   fcr.keyframe_points.add(len(fc.keyframe_points))
                   fc.keyframe_points.foreach_get("co", keys)
                   fcr.keyframe_points.foreach_set("co", keys)
                   print("ROOT")

for imp in imports:
    bpy.ops.wm.read_homefile(use_empty=True)
    for o in bpy.data.objects:
        bpy.data.objects.remove(o)
    print(basepath + "/assets/vroid/" + imp.path)
    result = bpy.ops.import_scene.vrm(filepath=(basepath + "/assets/vroid/" + imp.path))
    if result != {"FINISHED"}:
        raise Exception(f"Failed to import vrm: {result}")
    for o in bpy.data.objects:
        if o.type == 'MESH':
            tris2quads.tris2quads(o)
    armature_count = 0
    for obj in bpy.data.objects:
        if (obj.type == "ARMATURE"):
            armature_count += 1
    if armature_count > 1:
        raise Exception("Bad scene data")
    main_armature = None
    for obj in bpy.data.objects:
        if (obj.type == "ARMATURE"):
            main_armature = obj
            break
    if main_armature == None:
        raise Exception("Bad scene data")
    print("Renaming...")
    main_armature.select_set(True)
    bpy.context.view_layer.objects.active = main_armature
    rename.rename_bones(main_armature)
    main_armature.select_set(False)
    main_armature.name = imp.armature_name
    main_armature["VRM_IMPORTED_NAME"] = "female"
    if main_armature.animation_data is None:
        main_armature.animation_data_create()
    main_armature.select_set(True)
    bpy.context.view_layer.objects.active = main_armature
    mkrig()
    main_armature.select_set(False)
    bpy.context.view_layer.objects.active = None

    for anim in imp.mixamo_animations:
        anim_path = imp.mixamo_animation_path + anim + ".fbx"
        print("Load BVH..." + anim_path)
        old_objs = set(bpy.context.scene.objects)
        bpy.ops.import_scene.fbx(filepath=anim_path, global_scale=imp.fbx_scale)
        imported_objects = set(bpy.context.scene.objects) - old_objs
        for anim_obj in imported_objects:
            if anim_obj == main_armature:
                continue
            if "VRM_IMPORTED_NAME" in anim_obj:
                continue
            if main_armature == anim_obj or anim_obj.type != "ARMATURE":
                continue
            if anim_obj.animation_data == None:
                print("Bad object: " + anim_obj.name + " " + anim_obj.type)
                bpy.data.objects.remove(anim_obj)
                continue
            if anim_obj.animation_data.action == None:
                print("Bad object: " + anim_obj.name + " " + anim_obj.type)
                bpy.data.objects.remove(anim_obj)
                continue
            print("importing: " + anim_obj.animation_data.action.name)
            mixamo_rig._import_anim(anim_obj, main_armature, True)
            print("===" + main_armature.animation_data.action.name)
            main_armature.animation_data.action.name = anim
            hips_to_root2(main_armature)
            track = main_armature.animation_data.nla_tracks.new()
            action = main_armature.animation_data.action
            track.name = action.name
            track.strips.new(action.name, int(action.frame_range[0]), action)
            track.mute = True
            track.lock = True
            bpy.data.objects.remove(anim_obj)
        imported_objects.clear()
        imported_objects = set(bpy.context.scene.objects) - old_objs
        for anim_obj in imported_objects:
            bpy.data.objects.remove(anim_obj)

    main_armature.animation_data.action = bpy.data.actions.new(name="default")
    default_action = main_armature.animation_data.action
    bpy.context.view_layer.objects.active = main_armature
    main_armature.select_set(True)
    layers = enable_all_armature_layers()
    bpy.ops.object.mode_set(mode='POSE')
    for b in main_armature.pose.bones:
        b.location = [0,0,0]
        b.rotation_euler = [0,0,0]
        b.rotation_quaternion = [1,0,0,0]
        b.scale = [1,1,1]
    bpy.ops.object.mode_set(mode='OBJECT')

    bpy.ops.object.mode_set(mode='POSE')
    for b in main_armature.data.bones:
        if b.name.find("Ctrl_") < 0 and b.name.find("ctrl_") < 0:
            continue
        main_armature.pose.bones[b.name].matrix_basis.identity()
        main_armature.pose.bones[b.name].keyframe_insert(data_path='location', frame=1)
        rotation_mode = main_armature.pose.bones[b.name].rotation_mode
        if rotation_mode == "QUATERNION":
            main_armature.pose.bones[b.name].keyframe_insert(data_path='rotation_quaternion', frame=1)
        elif rotation_mode == "AXIS_ANGLE":
            main_armature.pose.bones[b.name].keyframe_insert(data_path='rotation_axis_angle', frame=1)
        else:
            main_armature.pose.bones[b.name].keyframe_insert(data_path='rotation_euler', frame=1)
        main_armature.pose.bones[b.name].keyframe_insert(data_path='scale', frame=1)
    bpy.ops.object.mode_set(mode='OBJECT')
    track = main_armature.animation_data.nla_tracks.new()
    action = main_armature.animation_data.action
    track.name = action.name
    track.strips.new(action.name, int(action.frame_range[0]), action)
    track.mute = True
    track.lock = True
    restore_armature_layers(layers)
    main_armature.animation_data.action = default_action

    main_armature.select_set(False)
    bpy.context.view_layer.objects.active = None
    bpy.ops.wm.save_as_mainfile(filepath=(basepath + "/assets/blender/" + imp.outfile))

