"""
ta_asset_health_export.py - The Technical Artist - Unreal Asset Health Dashboard
================================================================================

Scrapes asset data from your Unreal project and exports a JSON report that you
drag & drop into the Asset Health Dashboard web tool:

    https://thetechnicalartist.example/tools/unreal-asset-health/

WHAT IT COLLECTS (read-only - this script never modifies your project):
  - Static Meshes ... triangle counts, LOD counts, Nanite flag, material slots, referencers
  - Textures ........ dimensions, compression format, sRGB, mips, LOD group, streaming, referencers
  - Materials ....... domain, blend mode, two-sided, texture dependency count, referencers
  - Blueprints ...... parent class, tick settings, component count, dependency count, referencers

HOW TO RUN
  1. Enable the "Python Editor Script Plugin" (Edit -> Plugins -> search "Python"), restart.
  2. Window -> Output Log -> switch the console dropdown from Cmd to Python.
  3. Run:   exec(open(r"C:/path/to/ta_asset_health_export.py").read())
     ... or add the folder to Project Settings -> Python -> Additional Paths and:
           import ta_asset_health_export
  4. The JSON is written to  <Project>/Saved/AssetHealth/asset_health_<date>.json
     Drag that file into the web dashboard. Nothing is uploaded anywhere -
     all analysis happens locally in your browser.

SCOPE
  By default the script scans everything under /Game. Narrow it by changing
  SCAN_ROOT below (e.g. "/Game/Environment").

Compatible with UE 5.0 - 5.5 (uses defensive property access throughout).
"""

import unreal
import json
import os
import datetime

SCAN_ROOT = "/Game"          # narrow to e.g. "/Game/Levels/Dungeon" if you like
MAX_ASSETS_PER_TYPE = 5000   # safety valve for gigantic projects

SCHEMA_VERSION = 1
TOOL_ID = "unreal-asset-health"


# -- helpers -----------------------------------------------------------------

def _safe(obj, prop, default=None):
    """Read an editor property without exploding across engine versions."""
    try:
        return obj.get_editor_property(prop)
    except Exception:
        return default


def _enum_name(value, default=""):
    try:
        return str(value).split(".")[-1]
    except Exception:
        return default


def _asset_registry():
    return unreal.AssetRegistryHelpers.get_asset_registry()


def _count_referencers(package_name):
    """How many other packages reference this asset?"""
    try:
        opts = unreal.AssetRegistryDependencyOptions()
        refs = _asset_registry().get_referencers(package_name, opts)
        return len(refs) if refs else 0
    except Exception:
        return -1  # unknown


def _count_dependencies(package_name):
    try:
        opts = unreal.AssetRegistryDependencyOptions()
        deps = _asset_registry().get_dependencies(package_name, opts)
        return len(deps) if deps else 0
    except Exception:
        return -1


def _list_assets_of_class(class_name, script_pkg="/Script/Engine"):
    ar = _asset_registry()
    flt = unreal.ARFilter(
        class_paths=[unreal.TopLevelAssetPath(script_pkg, class_name)],
        package_paths=[SCAN_ROOT],
        recursive_paths=True,
        recursive_classes=True,
    )
    return list(ar.get_assets(flt))[:MAX_ASSETS_PER_TYPE]


# -- collectors --------------------------------------------------------------

def collect_static_meshes(task):
    out = []
    assets = _list_assets_of_class("StaticMesh")
    task.enter_progress_frame(0, f"Static Meshes: {len(assets)}")
    for ad in assets:
        task.enter_progress_frame(1, f"Mesh: {ad.asset_name}")
        path = str(ad.package_name)
        mesh = unreal.EditorAssetLibrary.load_asset(str(ad.get_soft_object_path().to_string()) if hasattr(ad, "get_soft_object_path") else path)
        if not mesh:
            continue

        # triangle count for LOD0 - try the modern subsystem, then the old library
        tris = -1
        try:
            sms = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem)
            tris = sms.get_number_triangles(mesh, 0)
        except Exception:
            try:
                tris = unreal.EditorStaticMeshLibrary.get_number_triangles(mesh, 0)
            except Exception:
                pass

        lods = -1
        try:
            lods = mesh.get_num_lods()
        except Exception:
            try:
                sms = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem)
                lods = sms.get_lod_count(mesh)
            except Exception:
                pass

        nanite = False
        ns = _safe(mesh, "nanite_settings")
        if ns is not None:
            nanite = bool(_safe(ns, "enabled", False))

        mats = _safe(mesh, "static_materials", []) or []

        out.append({
            "name": str(ad.asset_name),
            "path": path,
            "tris_lod0": int(tris) if tris is not None else -1,
            "lods": int(lods) if lods is not None else -1,
            "nanite": nanite,
            "material_slots": len(mats),
            "refs": _count_referencers(path),
        })
    return out


def collect_textures(task):
    out = []
    assets = _list_assets_of_class("Texture2D")
    task.enter_progress_frame(0, f"Textures: {len(assets)}")
    for ad in assets:
        task.enter_progress_frame(1, f"Texture: {ad.asset_name}")
        path = str(ad.package_name)
        tex = unreal.EditorAssetLibrary.load_asset(path)
        if not tex:
            continue

        w, h = -1, -1
        try:
            w, h = int(tex.blueprint_get_size_x()), int(tex.blueprint_get_size_y())
        except Exception:
            w = int(_safe(tex, "size_x", -1) or -1)
            h = int(_safe(tex, "size_y", -1) or -1)

        mip_setting = _enum_name(_safe(tex, "mip_gen_settings"), "")
        out.append({
            "name": str(ad.asset_name),
            "path": path,
            "width": w,
            "height": h,
            "format": _enum_name(_safe(tex, "compression_settings"), "Unknown"),
            "srgb": bool(_safe(tex, "srgb", False)),
            "mips": 0 if "NoMipmaps" in mip_setting else 1,
            "group": _enum_name(_safe(tex, "lod_group"), "Unknown"),
            "never_stream": bool(_safe(tex, "never_stream", False)),
            "refs": _count_referencers(path),
        })
    return out


def collect_materials(task):
    out = []
    assets = _list_assets_of_class("Material") + _list_assets_of_class("MaterialInstanceConstant")
    task.enter_progress_frame(0, f"Materials: {len(assets)}")
    for ad in assets:
        task.enter_progress_frame(1, f"Material: {ad.asset_name}")
        path = str(ad.package_name)
        mat = unreal.EditorAssetLibrary.load_asset(path)
        if not mat:
            continue

        is_instance = isinstance(mat, unreal.MaterialInstance)
        base = mat
        parent = ""
        if is_instance:
            p = _safe(mat, "parent")
            parent = p.get_name() if p else ""
            base = p or mat

        out.append({
            "name": str(ad.asset_name),
            "path": path,
            "is_instance": is_instance,
            "parent": parent,
            "domain": _enum_name(_safe(base, "material_domain"), "Surface"),
            "blend_mode": _enum_name(_safe(base, "blend_mode"), "Opaque"),
            "two_sided": bool(_safe(base, "two_sided", False)),
            "textures_referenced": max(0, _count_dependencies(path)),
            "refs": _count_referencers(path),
        })
    return out


def collect_blueprints(task):
    out = []
    assets = _list_assets_of_class("Blueprint")
    task.enter_progress_frame(0, f"Blueprints: {len(assets)}")
    for ad in assets:
        task.enter_progress_frame(1, f"Blueprint: {ad.asset_name}")
        path = str(ad.package_name)

        parent_class = ""
        tick_enabled = None
        num_components = -1
        try:
            gen_class = unreal.EditorAssetLibrary.load_blueprint_class(path)
            if gen_class:
                cdo = unreal.get_default_object(gen_class)
                if cdo:
                    try:
                        parent_class = cdo.get_class().get_name()
                    except Exception:
                        pass
                    if isinstance(cdo, unreal.Actor):
                        try:
                            tick_enabled = bool(cdo.is_actor_tick_enabled())
                        except Exception:
                            tick = _safe(cdo, "primary_actor_tick")
                            if tick is not None:
                                tick_enabled = bool(_safe(tick, "start_with_tick_enabled", False))
                        try:
                            num_components = len(cdo.get_components_by_class(unreal.ActorComponent))
                        except Exception:
                            pass
        except Exception:
            pass

        out.append({
            "name": str(ad.asset_name),
            "path": path,
            "parent_class": parent_class,
            "tick_enabled": tick_enabled,      # None = unknown / not an Actor
            "num_components": num_components,
            "deps": _count_dependencies(path),
            "refs": _count_referencers(path),
        })
    return out


# -- main --------------------------------------------------------------------

def run():
    unreal.log("=" * 70)
    unreal.log("TA Asset Health Export - scanning " + SCAN_ROOT)
    unreal.log("=" * 70)

    # generous frame estimate; ScopedSlowTask shows a cancellable progress bar
    with unreal.ScopedSlowTask(20000, "Asset Health Export") as task:
        task.make_dialog(True)
        payload = {
            "tool": TOOL_ID,
            "schemaVersion": SCHEMA_VERSION,
            "source": "unreal-" + unreal.SystemLibrary.get_engine_version().split("-")[0],
            "project": unreal.SystemLibrary.get_game_name(),
            "scanRoot": SCAN_ROOT,
            "exportedAt": datetime.datetime.now().isoformat(timespec="seconds"),
            "assets": {
                "meshes": collect_static_meshes(task),
                "textures": collect_textures(task),
                "materials": collect_materials(task),
                "blueprints": collect_blueprints(task),
            },
        }

    out_dir = os.path.join(unreal.SystemLibrary.get_project_saved_directory(), "AssetHealth")
    os.makedirs(out_dir, exist_ok=True)
    stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    out_path = os.path.join(out_dir, f"asset_health_{stamp}.json")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=1)

    a = payload["assets"]
    unreal.log("-" * 70)
    unreal.log(f"  Meshes:     {len(a['meshes'])}")
    unreal.log(f"  Textures:   {len(a['textures'])}")
    unreal.log(f"  Materials:  {len(a['materials'])}")
    unreal.log(f"  Blueprints: {len(a['blueprints'])}")
    unreal.log("-" * 70)
    unreal.log("WROTE: " + out_path)
    unreal.log("Drag this file into the Asset Health Dashboard in your browser.")
    return out_path


run()