"""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

SRC_DIR = r"C:\Users\koleg\Documents\Unreal Projects\ExampleProject\DemoAssets\RobotPolice"
DEST_PATH = "/Game/ImportedAssets"
MASTER_MATERIAL = f"{DEST_PATH}/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_PATH, unreal.Material, unreal.MaterialFactoryNew())
    for i, (suffix, (param, prop, channel)) in enumerate(PARAMS.items()):
        tex = mel.new_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, 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):
    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, options)


master = get_or_create_master_material()

# 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

    import_fbx(os.path.join(SRC_DIR, files["fbx"]), base_name)

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

    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}")
        mel.set_material_instance_texture_parameter_value(mi, param, texture)

    unreal.EditorAssetLibrary.save_asset(mi.get_path_name())

print(f"Imported {len(groups)} asset(s) into {DEST_PATH}")