"""Static / Skeletal Mesh ops — info, LODs, sockets, collision.

Ops:
  static_info             {path}                — LODs, materials, collision, sockets, bounds, vertex/tri count
  skeletal_info           {path}                — LODs, bones, sockets, skeleton, physics asset
  list_sockets            {path}                — works for both static and skeletal
  add_socket              {path, name, bone?, location?, rotation?, scale?}
  remove_socket           {path, name}
  list_lods               {path}
  set_lod_screen_size     {path, lod_index, screen_size}
  remove_lod              {path, lod_index}
  set_collision_complexity{path, complexity:'Default|UseSimpleAsComplex|UseComplexAsSimple'}
  add_simple_collision    {path, shape:'Box|Sphere|Capsule', extents_or_radius}
"""
from __future__ import annotations

import json
import os
import traceback

try:
    import unreal  # type: ignore
except ImportError:
    unreal = None


def _load(p): return unreal.EditorAssetLibrary.load_asset(p)
def _vec(v, d=(0, 0, 0)):
    if v is None: v = d
    return unreal.Vector(float(v[0]), float(v[1]), float(v[2]))
def _rot(v, d=(0, 0, 0)):
    if v is None: v = d
    return unreal.Rotator(float(v[0]), float(v[1]), float(v[2]))


def static_info(args):
    m = _load(args["path"])
    if m is None or not isinstance(m, unreal.StaticMesh):
        return {"ok": False, "error": "not a StaticMesh"}
    out = {"ok": True, "path": args["path"]}
    # UE 5.4+ moved methods from EditorStaticMeshLibrary to StaticMeshEditorSubsystem.
    smes = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem) if hasattr(unreal, "StaticMeshEditorSubsystem") else None
    legacy = unreal.EditorStaticMeshLibrary if hasattr(unreal, "EditorStaticMeshLibrary") else None

    def _call(*sources, **kw):
        method = kw["method"]; args_ = kw.get("args", ())
        for src in sources:
            if src is None: continue
            fn = getattr(src, method, None)
            if fn:
                try: return fn(m, *args_)
                except Exception: pass
        return None

    try:
        v = _call(smes, legacy, method="get_lod_count")
        if v is not None: out["lod_count"] = v
        v = _call(smes, legacy, method="get_number_verts", args=(0,))
        if v is not None: out["vertex_count"] = v
        v = _call(smes, legacy, method="get_number_triangles", args=(0,))
        if v is not None: out["triangle_count"] = v
        slots = _call(smes, legacy, method="get_lod_material_slot_names", args=(0,))
        if slots: out["material_slots"] = [str(n) for n in slots]
        b = _call(smes, legacy, method="get_bounds")
        if b is not None:
            try: out["bounds"] = list(b.box_extent)
            except Exception: pass
    except Exception as e:
        out["info_error"] = str(e)
    try:
        smes = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem) if hasattr(unreal, "StaticMeshEditorSubsystem") else None
        out["sockets"] = [str(n) for n in (smes.get_sockets(m) if smes else [])]
    except Exception:
        out["sockets"] = []
    try:
        out["materials"] = [{"slot": str(s.material_slot_name),
                             "material": s.material_interface.get_path_name() if s.material_interface else None}
                            for s in (m.get_editor_property("static_materials") or [])]
    except Exception:
        pass
    return out


def skeletal_info(args):
    m = _load(args["path"])
    if m is None:
        return {"ok": False, "error": "not found"}
    out = {"ok": True, "path": args["path"], "class": m.get_class().get_name()}
    try:
        sk = m.get_editor_property("skeleton")
        out["skeleton"] = sk.get_path_name().split(".")[0] if sk else None
    except Exception:
        pass
    try:
        pa = m.get_editor_property("physics_asset")
        out["physics_asset"] = pa.get_path_name().split(".")[0] if pa else None
    except Exception:
        pass
    try:
        ssel = unreal.EditorSkeletalMeshLibrary if hasattr(unreal, "EditorSkeletalMeshLibrary") else None
        if ssel:
            out["lod_count"] = ssel.get_lod_count(m)
    except Exception:
        pass
    try:
        out["sockets"] = [str(s.socket_name) for s in (m.get_editor_property("sockets") or [])]
    except Exception:
        out["sockets"] = []
    try:
        mats = m.get_editor_property("materials") or []
        out["materials"] = [(str(x.material_slot_name), x.material_interface.get_path_name() if x.material_interface else None) for x in mats]
    except Exception:
        pass
    return out


def list_sockets(args):
    m = _load(args["path"])
    if m is None:
        return {"ok": False, "error": "not found"}
    if isinstance(m, unreal.StaticMesh):
        out = []
        # Sockets array is protected — use StaticMeshEditorSubsystem or find_socket fallback.
        names = []
        smes = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem) if hasattr(unreal, "StaticMeshEditorSubsystem") else None
        if smes and hasattr(smes, "get_sockets"):
            try: names = [str(n) for n in (smes.get_sockets(m) or [])]
            except Exception: names = []
        for n in names:
            try:
                s = m.find_socket(n)
                if s:
                    out.append({"name": n,
                                "location": list(s.get_editor_property("relative_location")),
                                "rotation": list(s.get_editor_property("relative_rotation"))})
            except Exception:
                out.append({"name": n})
    else:
        out = []
        for s in (m.get_editor_property("sockets") or []):
            try:
                out.append({"name": str(s.get_editor_property("socket_name")),
                            "bone": str(s.get_editor_property("bone_name")),
                            "location": list(s.get_editor_property("relative_location")),
                            "rotation": list(s.get_editor_property("relative_rotation"))})
            except Exception:
                pass
    return {"ok": True, "sockets": out}


def add_socket(args):
    m = _load(args["path"])
    if m is None:
        return {"ok": False, "error": "not found"}
    if isinstance(m, unreal.StaticMesh):
        # The Sockets array is a protected UPROPERTY in Python — fall back to
        # StaticMeshEditorSubsystem which exposes set_socket / set_sockets
        # depending on UE version. If neither is exposed, instruct user to use UI.
        smes = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem) if hasattr(unreal, "StaticMeshEditorSubsystem") else None
        sock = unreal.StaticMeshSocket()
        sock.set_editor_property("socket_name", args["name"])
        sock.set_editor_property("relative_location", _vec(args.get("location")))
        sock.set_editor_property("relative_rotation", _rot(args.get("rotation")))
        sock.set_editor_property("relative_scale", _vec(args.get("scale"), (1, 1, 1)))
        added = False
        if smes:
            for fn_name in ("set_socket", "add_socket"):
                fn = getattr(smes, fn_name, None)
                if fn:
                    try: fn(m, sock); added = True; break
                    except Exception: pass
        if not added:
            return {"ok": False, "error": "StaticMesh sockets array is protected; StaticMeshEditorSubsystem has no set_socket/add_socket on this build. Use Static Mesh Editor UI."}
    else:
        sock = unreal.SkeletalMeshSocket()
        sock.set_editor_property("socket_name", args["name"])
        sock.set_editor_property("bone_name", args.get("bone") or "")
        sock.set_editor_property("relative_location", _vec(args.get("location")))
        sock.set_editor_property("relative_rotation", _rot(args.get("rotation")))
        sockets = list(m.get_editor_property("sockets") or [])
        sockets.append(sock)
        m.set_editor_property("sockets", sockets)
    unreal.EditorAssetLibrary.save_asset(m.get_path_name())
    return {"ok": True}


def remove_socket(args):
    m = _load(args["path"])
    if m is None:
        return {"ok": False, "error": "not found"}
    target = args["name"]
    if isinstance(m, unreal.StaticMesh):
        smes = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem) if hasattr(unreal, "StaticMeshEditorSubsystem") else None
        removed = False
        if smes:
            for fn_name in ("remove_socket",):
                fn = getattr(smes, fn_name, None)
                if fn:
                    try: fn(m, target); removed = True; break
                    except Exception: pass
        if not removed:
            return {"ok": False, "error": "StaticMesh sockets array is protected; no remove_socket on subsystem in this build. Use UI."}
    else:
        m.set_editor_property("sockets",
            [s for s in (m.get_editor_property("sockets") or [])
             if str(s.get_editor_property("socket_name")) != target])
    unreal.EditorAssetLibrary.save_asset(m.get_path_name())
    return {"ok": True}


def list_lods(args):
    m = _load(args["path"])
    if m is None:
        return {"ok": False, "error": "not found"}
    out = []
    if isinstance(m, unreal.StaticMesh):
        smes = unreal.get_editor_subsystem(unreal.StaticMeshEditorSubsystem) if hasattr(unreal, "StaticMeshEditorSubsystem") else None
        legacy = unreal.EditorStaticMeshLibrary if hasattr(unreal, "EditorStaticMeshLibrary") else None
        def _c(method, *a):
            for src in (smes, legacy):
                if src is None: continue
                fn = getattr(src, method, None)
                if fn:
                    try: return fn(m, *a)
                    except Exception: pass
            return None
        n = _c("get_lod_count") or 0
        for i in range(n):
            lod = {"index": i}
            ss = _c("get_lod_screen_size", i)
            if ss is not None: lod["screen_size"] = ss
            tr = _c("get_number_triangles", i)
            if tr is not None: lod["tris"] = tr
            out.append(lod)
    else:
        ssel = unreal.EditorSkeletalMeshLibrary if hasattr(unreal, "EditorSkeletalMeshLibrary") else None
        sses = unreal.get_editor_subsystem(unreal.SkeletalMeshEditorSubsystem) if hasattr(unreal, "SkeletalMeshEditorSubsystem") else None
        n = 0
        for src in (sses, ssel):
            if src and hasattr(src, "get_lod_count"):
                try: n = src.get_lod_count(m); break
                except Exception: pass
        for i in range(n):
            out.append({"index": i})
    return {"ok": True, "lods": out}


def set_lod_screen_size(args):
    m = _load(args["path"])
    if m is None or not isinstance(m, unreal.StaticMesh):
        return {"ok": False, "error": "static mesh only"}
    smel = unreal.EditorStaticMeshLibrary
    try:
        smel.set_lod_reduction_settings if False else None
        # No direct setter; use BuildSettings via set_lod_screen_size if exposed.
        if hasattr(smel, "set_lod_screen_size"):
            smel.set_lod_screen_size(m, int(args["lod_index"]), float(args["screen_size"]))
            unreal.EditorAssetLibrary.save_asset(m.get_path_name())
            return {"ok": True}
        return {"ok": False, "error": "set_lod_screen_size not exposed in this UE build"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def remove_lod(args):
    m = _load(args["path"])
    if m is None or not isinstance(m, unreal.StaticMesh):
        return {"ok": False, "error": "static mesh only"}
    try:
        unreal.EditorStaticMeshLibrary.remove_lods(m)
        return {"ok": True, "note": "all generated LODs removed (UE API limitation)"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def set_collision_complexity(args):
    m = _load(args["path"])
    if m is None or not isinstance(m, unreal.StaticMesh):
        return {"ok": False, "error": "static mesh only"}
    mode = args["complexity"]
    enum_map = {
        "Default":               unreal.CollisionTraceFlag.CTF_USE_DEFAULT,
        "UseSimpleAsComplex":    unreal.CollisionTraceFlag.CTF_USE_SIMPLE_AS_COMPLEX,
        "UseComplexAsSimple":    unreal.CollisionTraceFlag.CTF_USE_COMPLEX_AS_SIMPLE,
    }
    if mode not in enum_map:
        return {"ok": False, "error": f"unknown mode; valid: {list(enum_map)}"}
    try:
        bs = m.get_editor_property("body_setup")
        if bs:
            bs.set_editor_property("collision_trace_flag", enum_map[mode])
            unreal.EditorAssetLibrary.save_asset(m.get_path_name())
            return {"ok": True}
        return {"ok": False, "error": "body_setup missing"}
    except Exception as e:
        return {"ok": False, "error": str(e)}


def add_simple_collision(args):
    m = _load(args["path"])
    if m is None or not isinstance(m, unreal.StaticMesh):
        return {"ok": False, "error": "static mesh only"}
    smel = unreal.EditorStaticMeshLibrary
    shape = args["shape"]
    try:
        if shape == "Box":
            smel.add_simple_collisions(m, unreal.ScriptingCollisionShapeType.BOX)
        elif shape == "Sphere":
            smel.add_simple_collisions(m, unreal.ScriptingCollisionShapeType.SPHERE)
        elif shape == "Capsule":
            smel.add_simple_collisions(m, unreal.ScriptingCollisionShapeType.CAPSULE)
        else:
            return {"ok": False, "error": "shape must be Box|Sphere|Capsule"}
        unreal.EditorAssetLibrary.save_asset(m.get_path_name())
        return {"ok": True}
    except Exception as e:
        return {"ok": False, "error": str(e)}


OPS = {
    "static_info":              static_info,
    "skeletal_info":            skeletal_info,
    "list_sockets":             list_sockets,
    "add_socket":               add_socket,
    "remove_socket":            remove_socket,
    "list_lods":                list_lods,
    "set_lod_screen_size":      set_lod_screen_size,
    "remove_lod":               remove_lod,
    "set_collision_complexity": set_collision_complexity,
    "add_simple_collision":     add_simple_collision,
}


def entrypoint(payload_json):
    try:
        p = json.loads(payload_json)
    except Exception as e:
        return _tee({"ok": False, "error": f"bad json: {e}"}, {})
    op = p.get("op"); fn = OPS.get(op)
    if not fn:
        return _tee({"ok": False, "error": f"unknown op {op!r}", "available": list(OPS)}, p)
    try:
        return _tee(fn(p), p)
    except Exception as e:
        return _tee({"ok": False, "error": str(e), "traceback": traceback.format_exc()}, p)


def _tee(result, p):
    txt = json.dumps(result)
    try:
        if unreal is not None:
            sd = unreal.Paths.project_saved_dir() + "MCP/"
            os.makedirs(sd, exist_ok=True)
            rid = p.get("__rid") or "last"
            with open(sd + "last.json", "w", encoding="utf-8") as f: f.write(txt)
            if rid != "last":
                with open(sd + f"{rid}.json", "w", encoding="utf-8") as f: f.write(txt)
            unreal.log("[MCP-MESH] " + txt[:1500])
    except Exception:
        pass
    return txt