"""
The player's cybernetic arm (module of the import runner, see import_assets.py).

Needs SourceArt/Arm/arm_manifest.json and the FBX files (Tools/export_arm.ps1 makes them). Makes, under /Game/Player/Arm:
  Materials/M_Arm_Opaque   master: BaseColor, Metallic, Roughness (every arm material is opaque, none emits light)
  Materials/MI_<name>      one instance per Blender material (7), parameters straight from the manifest
  SM_Arm_<Part>            one static mesh per segment (20), the joint as the asset origin, in the part's own local frame

Safe to re-run: the cheap things (material graph and instances) are rebuilt every time; a mesh is imported again only
when its FBX changed (SHA-256 in Saved/ImportState/arm_import.json), the import options changed (IMPORT_SCHEMA) or --force.
Imported bounds are compared with the exporter's prediction to 0.05 cm; a part that disagrees is deleted and the run fails.
"""
import json
import os

import unreal

GAME_DIR = "/Game/Player/Arm"
MAT_DIR = GAME_DIR + "/Materials"
STATE_NAME = "arm_import"
IMPORT_SCHEMA = 1   # bump when the import options below change, so existing meshes are re-imported
# Tight on purpose: the finger parts are asymmetric by only a couple of millimetres in Y, which is how a mirrored
# import is told from a correct one (see the Kazah module for the same idea at 0.5 cm).
BOUNDS_TOL_CM = 0.05


def vec_param(ctx, mat, name, default, x, y):
    n = ctx.mel.create_material_expression(mat, unreal.MaterialExpressionVectorParameter, x, y)
    n.set_editor_property("parameter_name", name)
    n.set_editor_property("default_value", unreal.LinearColor(default[0], default[1], default[2], 1.0))
    return n


def scalar_param(ctx, mat, name, default, x, y):
    n = ctx.mel.create_material_expression(mat, unreal.MaterialExpressionScalarParameter, x, y)
    n.set_editor_property("parameter_name", name)
    n.set_editor_property("default_value", float(default))
    return n


def build_master(ctx):
    MEL, EAL = ctx.mel, ctx.eal
    mat = ctx.get_or_create(MAT_DIR, "M_Arm_Opaque", unreal.Material, unreal.MaterialFactoryNew())
    MEL.delete_all_material_expressions(mat)
    mat.set_editor_property("blend_mode", unreal.BlendMode.BLEND_OPAQUE)
    mat.set_editor_property("shading_model", unreal.MaterialShadingModel.MSM_DEFAULT_LIT)
    mat.set_editor_property("two_sided", False)
    base = vec_param(ctx, mat, "BaseColor", (0.5, 0.5, 0.5), -700, -300)
    metal = scalar_param(ctx, mat, "Metallic", 0.0, -700, -100)
    rough = scalar_param(ctx, mat, "Roughness", 0.5, -700, 100)
    MEL.connect_material_property(base, "", unreal.MaterialProperty.MP_BASE_COLOR)
    MEL.connect_material_property(metal, "", unreal.MaterialProperty.MP_METALLIC)
    MEL.connect_material_property(rough, "", unreal.MaterialProperty.MP_ROUGHNESS)
    MEL.recompile_material(mat)
    EAL.save_loaded_asset(mat)
    return mat


def build_instances(ctx, manifest, master):
    MEL, EAL = ctx.mel, ctx.eal
    out = {}
    for name, m in manifest["materials"].items():
        mi_name = "MI_" + name[3:] if name.startswith("RA_") else "MI_" + name
        mi = ctx.get_or_create(MAT_DIR, mi_name, unreal.MaterialInstanceConstant, unreal.MaterialInstanceConstantFactoryNew())
        mi.set_editor_property("parent", master)
        bc = m["base_color"]
        MEL.set_material_instance_vector_parameter_value(mi, "BaseColor", unreal.LinearColor(bc[0], bc[1], bc[2], 1.0))
        MEL.set_material_instance_scalar_parameter_value(mi, "Metallic", float(m["metallic"]))
        MEL.set_material_instance_scalar_parameter_value(mi, "Roughness", float(m["roughness"]))
        MEL.update_material_instance(mi)
        EAL.save_loaded_asset(mi)
        out[name] = mi
        ctx.log("material %-16s -> %s (base %s, metallic %.2f, roughness %.2f)" % (name, mi_name, bc, m["metallic"], m["roughness"]))
    return out


def import_part(ctx, part, state, src_dir):
    EAL = ctx.eal
    fbx = os.path.join(src_dir, part["file"])
    if not os.path.isfile(fbx):
        raise RuntimeError("missing %s - run Tools/export_arm.ps1 first" % fbx)
    asset_path = GAME_DIR + "/" + part["asset"]
    digest = ctx.sha256(fbx)
    prev = state.get(part["asset"], {})
    if (not ctx.force and EAL.does_asset_exist(asset_path)
            and prev.get("sha256") == digest and prev.get("schema") == IMPORT_SCHEMA):
        return EAL.load_asset(asset_path), False

    # Always import into a fresh asset (re-importing over a mesh makes the classic FBX importer reuse its old settings).
    if EAL.does_asset_exist(asset_path):
        state.pop(part["asset"], None)
        if not EAL.delete_asset(asset_path):
            raise RuntimeError("could not delete the old " + asset_path)
    ui = unreal.FbxImportUI()
    ui.set_editor_property("import_mesh", True)
    ui.set_editor_property("import_as_skeletal", False)
    ui.set_editor_property("import_animations", False)
    ui.set_editor_property("import_materials", False)   # our own, built above
    ui.set_editor_property("import_textures", False)
    ui.set_editor_property("mesh_type_to_import", unreal.FBXImportType.FBXIT_STATIC_MESH)
    sm = ui.get_editor_property("static_mesh_import_data")
    sm.set_editor_property("combine_meshes", True)
    sm.set_editor_property("convert_scene", True)               # FBX (Blender forward X, up Z) -> Unreal axes
    sm.set_editor_property("force_front_x_axis", True)          # the arm lies along +X
    sm.set_editor_property("convert_scene_unit", True)
    sm.set_editor_property("import_uniform_scale", 1.0)
    sm.set_editor_property("transform_vertex_to_absolute", True)
    sm.set_editor_property("bake_pivot_in_vertex", False)
    sm.set_editor_property("normal_import_method", unreal.FBXNormalImportMethod.FBXNIM_IMPORT_NORMALS)
    sm.set_editor_property("remove_degenerates", True)
    sm.set_editor_property("generate_lightmap_u_vs", False)
    sm.set_editor_property("auto_generate_collision", False)
    sm.set_editor_property("build_nanite", False)               # 17 to 8,700 triangles per part: no Nanite

    task = unreal.AssetImportTask()
    task.set_editor_property("filename", fbx)
    task.set_editor_property("destination_path", GAME_DIR)
    task.set_editor_property("destination_name", part["asset"])
    task.set_editor_property("replace_existing", True)
    task.set_editor_property("replace_existing_settings", True)
    task.set_editor_property("automated", True)
    task.set_editor_property("save", False)
    task.set_editor_property("factory", unreal.FbxFactory())    # classic importer, not Interchange
    task.set_editor_property("options", ui)
    ctx.tools.import_asset_tasks([task])
    if not [str(p) for p in task.get_editor_property("imported_object_paths")]:
        raise RuntimeError("import of %s produced no asset" % part["file"])
    asset = EAL.load_asset(asset_path)
    if asset is None:
        raise RuntimeError("expected %s after import" % asset_path)
    state[part["asset"]] = {"sha256": digest, "schema": IMPORT_SCHEMA}
    return asset, True


def assign_materials(ctx, mesh, part, instances):
    slots = list(mesh.get_editor_property("static_materials"))
    wanted = part["materials"]
    if len(slots) != len(wanted):
        ctx.fail("%s has %d material slots, the manifest lists %d" % (part["asset"], len(slots), len(wanted)))
    for i, slot in enumerate(slots):
        slot_name = str(slot.get_editor_property("material_slot_name"))
        key = slot_name if slot_name in instances else (wanted[i] if i < len(wanted) else None)
        if key is None or key not in instances:
            ctx.fail("%s slot %d '%s' has no matching material" % (part["asset"], i, slot_name))
            continue
        if mesh.get_material(i) != instances[key]:
            mesh.set_material(i, instances[key])


def check_bounds(ctx, mesh, part):
    b = mesh.get_bounds()
    o, e = b.origin, b.box_extent
    got_min = [o.x - e.x, o.y - e.y, o.z - e.z]
    got_max = [o.x + e.x, o.y + e.y, o.z + e.z]
    exp = part["local_bounds_ue_cm"]
    ok = all(abs(got_min[i] - exp["min"][i]) <= BOUNDS_TOL_CM and abs(got_max[i] - exp["max"][i]) <= BOUNDS_TOL_CM for i in range(3))
    worst = max(max(abs(got_min[i] - exp["min"][i]), abs(got_max[i] - exp["max"][i])) for i in range(3))
    ctx.log("bounds %-22s worst deviation %.4f cm  %s" % (part["asset"], worst, "OK" if ok else
            "MISMATCH got min %s max %s expected min %s max %s" % (
                [round(v, 3) for v in got_min], [round(v, 3) for v in got_max], exp["min"], exp["max"])))
    if not ok:
        ctx.fail("%s imported with the wrong orientation, mirroring or scale (see the bounds line)" % part["asset"])
    return ok


def run(ctx):
    src_dir = os.path.join(ctx.project_dir, "SourceArt", "Arm")
    manifest_path = os.path.join(src_dir, "arm_manifest.json")
    if not os.path.isfile(manifest_path):
        raise RuntimeError("no manifest at %s - run Tools/export_arm.ps1 first" % manifest_path)
    with open(manifest_path, "r", encoding="utf-8") as f:
        manifest = json.load(f)
    for d in (GAME_DIR, MAT_DIR):
        ctx.eal.make_directory(d)
    instances = build_instances(ctx, manifest, build_master(ctx))

    state = ctx.load_state(STATE_NAME)
    imported_now = 0
    for part in manifest["parts"]:
        if not part.get("asset"):
            continue   # RA_Root is an empty
        mesh, fresh = import_part(ctx, part, state, src_dir)
        assign_materials(ctx, mesh, part, instances)
        if not check_bounds(ctx, mesh, part):
            state.pop(part["asset"], None)
            ctx.eal.delete_asset(GAME_DIR + "/" + part["asset"])
            continue
        ctx.eal.save_loaded_asset(mesh)
        imported_now += 1 if fresh else 0
    ctx.save_state(STATE_NAME, state)
    ctx.log("arm: %d meshes checked, %d imported this run" % (len(manifest["parts"]) - 1, imported_now))