"""Animation ops — AnimBlueprint, AnimSequence, AnimMontage, Notifies, Skeleton sockets.

Ops:
  create_anim_blueprint   {path, skeleton}
  create_anim_montage     {path, sequence?}
  create_blend_space      {path, skeleton, dim?:1|2}
  get_sequence_info       {path}                 — length, fps, frames, additive type, skeleton
  list_anim_assets        {folder?, skeleton?}
  list_sockets            {skeleton}
  add_socket              {skeleton, name, bone, location?, rotation?}
  remove_socket           {skeleton, name}
  list_montage_sections   {path}
  add_montage_section     {path, name, start_time?}
  list_notifies           {sequence}
  add_notify              {sequence, name, time, notify_class?}
  remove_notify           {sequence, index}
  list_state_machines     {anim_bp}              — best-effort (graph names)
  retarget_animation      {src, dst_skeleton, dst_path}
"""
from __future__ import annotations

import json
import os
import traceback

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


def _at():
    return unreal.AssetToolsHelpers.get_asset_tools()


def _load(path: str):
    return unreal.EditorAssetLibrary.load_asset(path)


def _split_pkg(path: str):
    folder, name = path.rsplit("/", 1)
    return folder, name


def _vec(v, default=(0, 0, 0)):
    if v is None:
        v = default
    return unreal.Vector(float(v[0]), float(v[1]), float(v[2]))


def _rot(v, default=(0, 0, 0)):
    if v is None:
        v = default
    return unreal.Rotator(float(v[0]), float(v[1]), float(v[2]))


def create_anim_blueprint(args: dict) -> dict:
    skel = _load(args["skeleton"])
    if skel is None:
        return {"ok": False, "error": f"skeleton not found: {args['skeleton']}"}
    folder, name = _split_pkg(args["path"])
    factory = unreal.AnimBlueprintFactory()
    factory.set_editor_property("target_skeleton", skel)
    asset = _at().create_asset(name, folder, unreal.AnimBlueprint, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    unreal.EditorAssetLibrary.save_asset(asset.get_path_name())
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_anim_montage(args: dict) -> dict:
    folder, name = _split_pkg(args["path"])
    factory = unreal.AnimMontageFactory()
    seq_path = args.get("sequence")
    if seq_path:
        seq = _load(seq_path)
        if seq is None:
            return {"ok": False, "error": f"sequence not found: {seq_path}"}
        if hasattr(factory, "set_editor_property"):
            try:
                factory.set_editor_property("preview_skeletal_mesh", seq.get_editor_property("skeleton"))
            except Exception:
                pass
    asset = _at().create_asset(name, folder, unreal.AnimMontage, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def create_blend_space(args: dict) -> dict:
    skel = _load(args["skeleton"])
    if skel is None:
        return {"ok": False, "error": "skeleton not found"}
    folder, name = _split_pkg(args["path"])
    dim = int(args.get("dim") or 2)
    factory_cls = getattr(unreal, "BlendSpaceFactory1D" if dim == 1 else "BlendSpaceFactoryNew", None)
    if factory_cls is None:
        return {"ok": False, "error": "BlendSpace factory unavailable in this UE build"}
    factory = factory_cls()
    if hasattr(factory, "target_skeleton"):
        factory.target_skeleton = skel
    asset_cls = unreal.BlendSpace1D if dim == 1 else unreal.BlendSpace
    asset = _at().create_asset(name, folder, asset_cls, factory)
    if asset is None:
        return {"ok": False, "error": "create_asset returned None"}
    return {"ok": True, "path": asset.get_path_name().split(".")[0]}


def get_sequence_info(args: dict) -> dict:
    seq = _load(args["path"])
    if seq is None:
        return {"ok": False, "error": "not found"}
    out = {"ok": True, "path": args["path"], "class": seq.get_class().get_name()}
    for attr in ("get_play_length", "sequence_length"):
        try:
            v = getattr(seq, attr)
            out["length"] = v() if callable(v) else float(v)
            break
        except Exception:
            pass
    try:
        out["num_frames"] = int(seq.get_editor_property("number_of_sampled_keys") or 0)
    except Exception:
        pass
    try:
        out["frame_rate"] = float(seq.get_editor_property("target_frame_rate") or 0)
    except Exception:
        pass
    try:
        skel = seq.get_editor_property("skeleton")
        if skel:
            out["skeleton"] = skel.get_path_name().split(".")[0]
    except Exception:
        pass
    try:
        out["additive_type"] = str(seq.get_editor_property("additive_anim_type"))
    except Exception:
        pass
    return out


def list_anim_assets(args: dict) -> dict:
    folder = args.get("folder") or "/Game"
    skel_filter = args.get("skeleton")
    ar = unreal.AssetRegistryHelpers.get_asset_registry()
    filt = unreal.ARFilter(package_paths=[folder], recursive_paths=True,
                           class_names=["AnimSequence", "AnimMontage", "BlendSpace", "BlendSpace1D", "AnimBlueprint"])
    items = []
    for d in ar.get_assets(filt):
        path = str(d.package_name)
        if skel_filter:
            try:
                asset = _load(path)
                sk = asset.get_editor_property("skeleton") if asset else None
                if not sk or sk.get_path_name().split(".")[0] != skel_filter:
                    continue
            except Exception:
                continue
        items.append({"path": path, "class": str(d.asset_class)})
    return {"ok": True, "count": len(items), "items": items}


def list_sockets(args: dict) -> dict:
    skel = _load(args["skeleton"])
    if skel is None:
        return {"ok": False, "error": "skeleton not found"}
    try:
        sockets = skel.get_editor_property("sockets") or []
    except Exception:
        sockets = []
    out = []
    for s in sockets:
        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, "count": len(out), "sockets": out}


def add_socket(args: dict) -> dict:
    skel = _load(args["skeleton"])
    if skel is None:
        return {"ok": False, "error": "skeleton not found"}
    sock = unreal.SkeletalMeshSocket()
    sock.set_editor_property("socket_name", args["name"])
    sock.set_editor_property("bone_name", args["bone"])
    sock.set_editor_property("relative_location", _vec(args.get("location")))
    sock.set_editor_property("relative_rotation", _rot(args.get("rotation")))
    sockets = list(skel.get_editor_property("sockets") or [])
    sockets.append(sock)
    skel.set_editor_property("sockets", sockets)
    unreal.EditorAssetLibrary.save_asset(skel.get_path_name())
    return {"ok": True, "name": args["name"]}


def remove_socket(args: dict) -> dict:
    skel = _load(args["skeleton"])
    if skel is None:
        return {"ok": False, "error": "skeleton not found"}
    target = args["name"]
    sockets = [s for s in (skel.get_editor_property("sockets") or [])
               if str(s.get_editor_property("socket_name")) != target]
    skel.set_editor_property("sockets", sockets)
    unreal.EditorAssetLibrary.save_asset(skel.get_path_name())
    return {"ok": True, "removed": target}


def list_montage_sections(args: dict) -> dict:
    m = _load(args["path"])
    if m is None:
        return {"ok": False, "error": "not found"}
    try:
        secs = m.get_editor_property("composite_sections") or []
    except Exception:
        secs = []
    out = []
    for s in secs:
        try:
            out.append({"name": str(s.get_editor_property("section_name")),
                        "start_time": float(s.get_editor_property("segment_begin_time") or 0)})
        except Exception:
            pass
    return {"ok": True, "sections": out}


def add_montage_section(args: dict) -> dict:
    m = _load(args["path"])
    if m is None:
        return {"ok": False, "error": "not found"}
    sec = unreal.CompositeSection() if hasattr(unreal, "CompositeSection") else None
    if sec is None:
        return {"ok": False, "error": "CompositeSection not exposed in this UE build"}
    sec.set_editor_property("section_name", args["name"])
    sec.set_editor_property("segment_begin_time", float(args.get("start_time") or 0))
    secs = list(m.get_editor_property("composite_sections") or [])
    secs.append(sec)
    m.set_editor_property("composite_sections", secs)
    unreal.EditorAssetLibrary.save_asset(m.get_path_name())
    return {"ok": True}


def list_notifies(args: dict) -> dict:
    seq = _load(args["sequence"])
    if seq is None:
        return {"ok": False, "error": "not found"}
    try:
        notifies = seq.get_editor_property("notifies") or []
    except Exception:
        notifies = []
    out = []
    for n in notifies:
        try:
            cls = n.get_editor_property("notify")
            out.append({
                "name": str(n.get_editor_property("notify_name")),
                "time": float(n.get_editor_property("trigger_time_offset") or 0) + float(n.get_editor_property("link_value") or 0),
                "class": cls.get_class().get_name() if cls else None,
            })
        except Exception:
            pass
    return {"ok": True, "notifies": out}


def add_notify(args: dict) -> dict:
    seq = _load(args["sequence"])
    if seq is None:
        return {"ok": False, "error": "not found"}
    ev = unreal.AnimNotifyEvent() if hasattr(unreal, "AnimNotifyEvent") else None
    if ev is None:
        return {"ok": False, "error": "AnimNotifyEvent not exposed; use editor UI"}
    ev.set_editor_property("notify_name", args["name"])
    ev.set_editor_property("link_value", float(args["time"]))
    notifies = list(seq.get_editor_property("notifies") or [])
    notifies.append(ev)
    seq.set_editor_property("notifies", notifies)
    unreal.EditorAssetLibrary.save_asset(seq.get_path_name())
    return {"ok": True}


def remove_notify(args: dict) -> dict:
    seq = _load(args["sequence"])
    if seq is None:
        return {"ok": False, "error": "not found"}
    notifies = list(seq.get_editor_property("notifies") or [])
    idx = int(args["index"])
    if idx < 0 or idx >= len(notifies):
        return {"ok": False, "error": "index out of range"}
    notifies.pop(idx)
    seq.set_editor_property("notifies", notifies)
    unreal.EditorAssetLibrary.save_asset(seq.get_path_name())
    return {"ok": True}


def list_state_machines(args: dict) -> dict:
    bp = _load(args["anim_bp"])
    if bp is None:
        return {"ok": False, "error": "not found"}
    out = []
    try:
        for g in bp.get_editor_property("ubergraph_pages") or []:
            out.append(g.get_name())
    except Exception:
        pass
    try:
        for g in bp.get_editor_property("function_graphs") or []:
            out.append(g.get_name())
    except Exception:
        pass
    return {"ok": True, "graphs": out}


def retarget_animation(args: dict) -> dict:
    src = _load(args["src"])
    dst_skel = _load(args["dst_skeleton"])
    if src is None or dst_skel is None:
        return {"ok": False, "error": "src or dst_skeleton not found"}
    folder, name = _split_pkg(args["dst_path"])
    try:
        new_asset = unreal.AnimationLibrary.duplicate_animation_asset(src, name, folder) \
            if hasattr(unreal, "AnimationLibrary") and hasattr(unreal.AnimationLibrary, "duplicate_animation_asset") \
            else unreal.EditorAssetLibrary.duplicate_asset(src.get_path_name(), folder + "/" + name)
    except Exception as e:
        return {"ok": False, "error": f"duplicate failed: {e}"}
    try:
        new_asset.set_editor_property("skeleton", dst_skel)
    except Exception:
        pass
    unreal.EditorAssetLibrary.save_asset(new_asset.get_path_name())
    return {"ok": True, "path": new_asset.get_path_name().split(".")[0]}


OPS = {
    "create_anim_blueprint":  create_anim_blueprint,
    "create_anim_montage":    create_anim_montage,
    "create_blend_space":     create_blend_space,
    "get_sequence_info":      get_sequence_info,
    "list_anim_assets":       list_anim_assets,
    "list_sockets":           list_sockets,
    "add_socket":             add_socket,
    "remove_socket":          remove_socket,
    "list_montage_sections":  list_montage_sections,
    "add_montage_section":    add_montage_section,
    "list_notifies":          list_notifies,
    "add_notify":             add_notify,
    "remove_notify":          remove_notify,
    "list_state_machines":    list_state_machines,
    "retarget_animation":     retarget_animation,
}


def entrypoint(payload_json: str) -> str:
    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: dict, p: dict) -> str:
    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-ANIM] " + txt[:1500])
    except Exception:
        pass
    return txt