"""Run inside Unreal (Python console / Execute Python Script).

Watches a disk folder for FBX + textures named by convention (Name.fbx,
Name_BC/_N/_M/_R/_E.png), imports everything without opening the import
dialog, and wires the textures into a Material Instance per mesh.
"""
import os, unreal

DEMO_ROOT = r"C:\Users\koleg\Documents\Unreal Projects\ExampleProject\DemoAssets"
DEST_ROOT = "/Game/ImportedAssets"
MASTER_MATERIAL = f"{DEST_ROOT}/M_Master"

# suffix -> (texture parameter name, material property, sampled channel)
PARAMS = {
    "_BC": ("BaseColor", unreal.MaterialProperty.MP_BASE_COLOR, "RGB"),
    "_N": ("Normal", unreal.MaterialProperty.MP_NORMAL, "RGB"),
    "_M": ("Metallic", unreal.MaterialProperty.MP_METALLIC, "R"),
    "_R": ("Roughness", unreal.MaterialProperty.MP_ROUGHNESS, "R"),
    "_E": ("Emissive", unreal.MaterialProperty.MP_EMISSIVE_COLOR, "RGB"),
}

asset_tools = unreal.AssetToolsHelpers.get_asset_tools()
mel = unreal.MaterialEditingLibrary


def get_or_create_master_material():
    if unreal.EditorAssetLibrary.does_asset_exist(MASTER_MATERIAL):
        return unreal.EditorAssetLibrary.load_asset(MASTER_MATERIAL)

    material = asset_tools.create_asset("M_Master", DEST_ROOT, unreal.Material, unreal.MaterialFactoryNew())
    for i, (suffix, (param, prop, channel)) in enumerate(PARAMS.items()):
        tex = mel.create_material_expression(material, unreal.MaterialExpressionTextureSampleParameter2D, -400, i * 200)
        tex.set_editor_property("parameter_name", param)
        mel.connect_material_property(tex, channel, prop)
    mel.recompile_material(material)
    return material


def import_task(filename, name, dest_path, options=None):
    task = unreal.AssetImportTask()
    task.filename = filename
    task.destination_path = dest_path
    task.destination_name = name
    task.automated = True
    task.save = True
    task.replace_existing = True
    if options:
        task.options = options
    asset_tools.import_asset_tasks([task])
    return unreal.load_asset(f"{dest_path}/{name}")


def import_fbx(filename, name, dest_path):
    options = unreal.FbxImportUI()
    options.mesh_type_to_import = unreal.FBXImportType.FBXIT_SKELETAL_MESH
    options.import_mesh = True
    options.import_as_skeletal = True
    options.import_animations = True
    options.create_physics_asset = False
    options.import_materials = False
    options.import_textures = False
    return import_task(filename, name, dest_path, options)


def assign_material(skeletal_mesh, material):
    existing_slots = skeletal_mesh.get_editor_property("materials")
    new_slots = []
    for slot in existing_slots:
        new_slot = unreal.SkeletalMaterial()
        new_slot.set_editor_property("material_interface", material)
        new_slot.set_editor_property("material_slot_name", slot.get_editor_property("material_slot_name"))
        new_slots.append(new_slot)
    if not new_slots:
        # FBX had no material assigned to the mesh, so no slot was created on import.
        fallback = unreal.SkeletalMaterial()
        fallback.set_editor_property("material_interface", material)
        fallback.set_editor_property("material_slot_name", "Material")
        new_slots.append(fallback)
    skeletal_mesh.modify()
    skeletal_mesh.set_editor_property("materials", new_slots)
    unreal.EditorAssetLibrary.save_asset(skeletal_mesh.get_path_name())


master = get_or_create_master_material()

asset_count = 0
for asset_folder in os.listdir(DEMO_ROOT):
    src_dir = os.path.join(DEMO_ROOT, asset_folder)
    if not os.path.isdir(src_dir):
        continue
    dest_path = f"{DEST_ROOT}/{asset_folder}"

    # group source files by base name: {"RobotPolice": {"fbx": "...", "_BC": "...", ...}}
    groups = {}
    for f in os.listdir(src_dir):
        stem, ext = os.path.splitext(f)
        if ext.lower() == ".fbx":
            groups.setdefault(stem, {})["fbx"] = f
            continue
        for suffix in PARAMS:
            if stem.endswith(suffix):
                groups.setdefault(stem[: -len(suffix)], {})[suffix] = f
                break

    for base_name, files in groups.items():
        if "fbx" not in files:
            continue

        mesh = import_fbx(os.path.join(src_dir, files["fbx"]), base_name, dest_path)

        mi = asset_tools.create_asset(f"MI_{base_name}", dest_path, unreal.MaterialInstanceConstant,
                                       unreal.MaterialInstanceConstantFactoryNew())
        mel.set_material_instance_parent(mi, master)

        textures = {}
        for suffix, (param, _, _) in PARAMS.items():
            if suffix not in files:
                continue
            texture = import_task(os.path.join(src_dir, files[suffix]), f"T_{base_name}{suffix}", dest_path)
            textures[suffix] = texture
            mel.set_material_instance_texture_parameter_value(mi, param, texture)

        if "_E" not in files and "_BC" in textures:
            mel.set_material_instance_texture_parameter_value(mi, "Emissive", textures["_BC"])

        unreal.EditorAssetLibrary.save_asset(mi.get_path_name())
        assign_material(mesh, mi)
        asset_count += 1

print(f"Imported {asset_count} asset(s) into {DEST_ROOT}")