"""
UE-side runtime for the Blueprint MCP server.

Reaches into the editor via `unreal.MCPGraphLibrary` — a C++ BlueprintFunction
library provided by Plugins/UEBlueprintMCP/Source. That library exposes K2 node
spawning, which is NOT available in stock UE Python.

Flow per apply():
  1. Resolve target Blueprint (by path, or via Content Browser selection).
  2. Resolve target graph (EventGraph by default, else named function).
  3. For each NodeSpec — call the matching MCPGraphLibrary spawn function and
     remember NodeGuid keyed by local id.
  4. Set pin defaults (params).
  5. For each EdgeSpec — call ConnectPins using NodeGuids.
  6. Compile + save Blueprint.
  7. Write the result envelope to <Project>/Saved/MCP/last.json and unreal.log.
"""
from __future__ import annotations

import json
import os
import traceback

import graph_layout

try:
    import unreal  # type: ignore
except ImportError:  # pragma: no cover — tests run outside UE
    unreal = None


def _lib():
    """Return the MCPGraphLibrary class, or raise if the plugin isn't built."""
    L = getattr(unreal, "MCPGraphLibrary", None)
    if L is None:
        raise RuntimeError(
            "unreal.MCPGraphLibrary not found — the UEBlueprintMCP C++ plugin "
            "is not built. Right-click your .uproject → Generate VS project files, "
            "then build the editor."
        )
    return L


# ---------------------------------------------------------------------------
# Blueprint / graph resolution
# ---------------------------------------------------------------------------
def _load_blueprint(path: str):
    asset = unreal.EditorAssetLibrary.load_asset(path)
    if asset is None:
        raise RuntimeError(f"Blueprint not found at {path}")
    if not isinstance(asset, unreal.Blueprint):
        raise RuntimeError(f"Asset at {path} is not a Blueprint ({type(asset).__name__})")
    return asset


def _active_blueprint():
    """Return (Blueprint, path) from Content Browser selection."""
    selected = []
    try:
        selected = list(unreal.EditorUtilityLibrary.get_selected_assets() or [])
    except Exception:
        selected = []
    bps = [a for a in selected if isinstance(a, unreal.Blueprint)]
    if len(bps) == 1:
        bp = bps[0]
        return bp, bp.get_path_name().split(".")[0]
    if len(bps) > 1:
        raise RuntimeError(f"{len(bps)} Blueprints selected — select exactly one or pass bp_path")
    raise RuntimeError(
        "auto-detect needs one Blueprint selected in Content Browser. "
        "Click your BP asset and retry, or pass bp_path explicitly."
    )


def _find_graph(bp, function_name: str):
    L = _lib()
    if function_name in (None, "", "EventGraph"):
        g = L.find_event_graph(bp)
        if g is None:
            raise RuntimeError("Blueprint has no EventGraph")
        return g
    g = L.find_function_graph(bp, function_name)
    if g is None:
        raise RuntimeError(f"Graph '{function_name}' not found on Blueprint")
    return g


# ---------------------------------------------------------------------------
# Function target parsing
# ---------------------------------------------------------------------------
def _resolve_function_owner(target: str):
    """Parse 'KismetSystemLibrary.PrintString' or '/Script/Engine.KismetSystemLibrary:PrintString'.
    Returns (UClass, FName-able str).
    """
    if ":" in target:
        cls_path, fn = target.split(":", 1)
    elif "." in target:
        cls_path, fn = target.rsplit(".", 1)
    else:
        raise RuntimeError(f"call_function target must be Class.Function, got {target!r}")

    cls = None
    if "/" in cls_path:
        cls = unreal.load_class(None, cls_path)
    else:
        for prefix in ("/Script/Engine.", "/Script/CoreUObject.", "/Script/UMG.",
                       "/Script/GameplayTags.", "/Script/UnrealEd.", "/Script/Kismet."):
            cls = unreal.load_class(None, prefix + cls_path)
            if cls is not None:
                break
    if cls is None:
        raise RuntimeError(f"Class not found for target {target!r}")
    return cls, fn


# ---------------------------------------------------------------------------
# Apply
# ---------------------------------------------------------------------------
def apply(graph: dict) -> dict:
    result = {
        "ok": False, "spawned": {}, "node_guids": {},
        "errors": [], "warnings": [], "compile": None,
    }
    if unreal is None:
        result["errors"].append("unreal module not importable — must run inside UE editor")
        return result

    try:
        L = _lib()
        bp_path = graph.get("bp_path") or ""
        if bp_path:
            bp = _load_blueprint(bp_path)
            result["resolved_bp"] = bp_path
        else:
            bp, bp_path = _active_blueprint()
            result["resolved_bp"] = bp_path
            result["auto_detected"] = True

        function_name = graph.get("function") or "EventGraph"
        target_graph = _find_graph(bp, function_name)

        # ---- mode handling ----
        mode = (graph.get("mode") or "append").lower()
        if mode in ("replace_function", "replace_event"):
            try:
                cleared = L.clear_graph(target_graph)
                result["cleared_nodes"] = int(cleared)
            except Exception as e:  # noqa: BLE001
                result["warnings"].append(f"clear_graph failed: {e}")

        layout = str(graph.get("layout") or "").lower()
        auto_layout = bool(graph.get("auto_layout", False))
        should_layout = auto_layout or layout in ("compact", "columns")

        node_specs = graph_layout.normalize_specs(
            graph.get("nodes", []),
            graph.get("edges", []),
            domain="blueprint",
        ) if should_layout else list(graph.get("nodes", []))

        # ---- spawn nodes ----
        guid_map: dict[str, str] = {}
        for spec in node_specs:
            try:
                guid = _spawn_one(target_graph, spec, bp, L)
                if not guid:
                    raise RuntimeError("spawn returned empty guid")
                guid_map[spec["id"]] = guid
                result["spawned"][spec["id"]] = guid
            except Exception as e:  # noqa: BLE001
                result["errors"].append(f"node {spec.get('id')!r}: {e}")

        # ---- post-spawn config (switch cases, format text, sequence pins are spawn-time) ----
        for spec in node_specs:
            local_id = spec.get("id")
            guid = guid_map.get(local_id)
            if not guid:
                continue
            p = spec.get("params") or {}
            kind = spec.get("kind")

            # FormatText: set Format & reconstruct so {arg} pins appear
            if kind == "format_text" and "Format" in p:
                try:
                    L.format_text_set_format(target_graph, guid, str(p["Format"]))
                except Exception as e:  # noqa: BLE001
                    result["warnings"].append(f"format_text_set_format: {local_id}: {e}")

            # Switch nodes: add case pins
            if kind in ("switch_int", "switch_string", "switch_enum"):
                cases = int(p.get("cases", 0))
                for _ in range(cases):
                    try:
                        L.add_case_pin_to_switch(target_graph, guid)
                    except Exception as e:  # noqa: BLE001
                        result["warnings"].append(f"add_case_pin_to_switch: {local_id}: {e}")

        # ---- set pin defaults ----
        SKIP_PARAMS = {"cases", "then_count", "width", "height", "text", "Format"}
        for spec in node_specs:
            local_id = spec.get("id")
            guid = guid_map.get(local_id)
            if not guid:
                continue
            for pin_name, value in (spec.get("params") or {}).items():
                if pin_name in SKIP_PARAMS:
                    continue
                try:
                    ok = L.set_pin_default(target_graph, guid, pin_name, str(value))
                    if not ok:
                        result["warnings"].append(f"set_pin_default failed: {local_id}.{pin_name}")
                except Exception as e:  # noqa: BLE001
                    result["warnings"].append(f"set_pin_default raised: {local_id}.{pin_name}: {e}")

        # ---- connect edges ----
        def _resolve_ref(ref: str):
            # "@guid:<NodeGuid>" -> raw guid (existing node)
            if isinstance(ref, str) and ref.startswith("@guid:"):
                return ref[len("@guid:"):]
            return guid_map.get(ref)

        for edge in graph.get("edges", []):
            try:
                (fa, fp) = edge["from"]
                (tb, tp) = edge["to"]
                ga, gb = _resolve_ref(fa), _resolve_ref(tb)
                if not ga or not gb:
                    raise RuntimeError("edge references unspawned/unknown node")
                ok = L.connect_pins(target_graph, ga, fp, gb, tp)
                if not ok:
                    raise RuntimeError(f"connect_pins returned false ({fa}.{fp} -> {tb}.{tp})")
            except Exception as e:  # noqa: BLE001
                result["errors"].append(f"edge {edge}: {e}")

        if result["errors"]:
            result["ok"] = False
        else:
            try:
                msgs = list(L.compile_with_messages(bp))
                result["compile_messages"] = msgs
                result["compile"] = "ok" if all(not m.startswith("ERROR|") for m in msgs) else "errors"
            except Exception:
                L.compile_and_save_blueprint(bp)
                result["compile"] = "ok"
            unreal.EditorAssetLibrary.save_loaded_asset(bp)
            result["ok"] = (result["compile"] == "ok")
    except Exception as e:  # noqa: BLE001
        result["errors"].append(str(e))
        result["traceback"] = traceback.format_exc()
    return result


def _load_class(path_or_name: str):
    if "/" in path_or_name:
        return unreal.load_class(None, path_or_name)
    for prefix in ("/Script/Engine.", "/Script/CoreUObject.", "/Script/UMG.",
                   "/Script/GameplayTags.", "/Script/UnrealEd.", "/Script/Kismet."):
        c = unreal.load_class(None, prefix + path_or_name)
        if c is not None:
            return c
    return None


def _load_struct(path: str):
    if "/" in path:
        return unreal.load_object(None, path)
    return unreal.load_object(None, "/Script/CoreUObject." + path)


def _load_enum(path: str):
    if "/" in path:
        return unreal.load_object(None, path)
    return unreal.load_object(None, "/Script/CoreUObject." + path)


def _spawn_one(graph, spec: dict, bp, L) -> str:
    kind = spec["kind"]
    pos = unreal.Vector2D(*(spec.get("position") or [0, 0]))
    target = spec.get("target") or ""

    # ---- Function / event ----
    if kind == "call_function":
        cls, fn = _resolve_function_owner(target)
        return L.spawn_call_function_node(bp, graph, cls, fn, pos)
    if kind == "event":
        return L.spawn_event_node(bp, graph, None, target, pos)
    if kind == "custom_event":
        return L.spawn_custom_event_node(bp, graph, target, pos)
    if kind == "component_event":
        if "." not in target:
            raise RuntimeError(f"component_event target must be 'Component.Delegate', got {target!r}")
        comp, delegate = target.split(".", 1)
        return L.spawn_component_bound_event(bp, graph, comp, delegate, pos)

    # ---- Variables ----
    if kind == "variable_get":
        return L.spawn_variable_get_node(bp, graph, target, pos)
    if kind == "variable_set":
        return L.spawn_variable_set_node(bp, graph, target, pos)
    if kind == "self":
        return L.spawn_self_node(bp, graph, pos)

    # ---- Flow control ----
    if kind == "branch":
        return L.spawn_branch_node(bp, graph, pos)
    if kind == "sequence":
        out_count = int((spec.get("params") or {}).get("then_count", 2))
        return L.spawn_sequence_node(bp, graph, pos, out_count)
    if kind == "switch_int":
        return L.spawn_switch_on_int_node(bp, graph, pos)
    if kind == "switch_string":
        return L.spawn_switch_on_string_node(bp, graph, pos)
    if kind == "switch_enum":
        e = _load_enum(target)
        if e is None:
            raise RuntimeError(f"switch_enum: enum not found at {target!r}")
        return L.spawn_switch_on_enum_node(bp, graph, e, pos)
    if kind == "macro":
        # target: macro name from StandardMacros (e.g. "ForLoop", "ForEachLoop",
        # "WhileLoop", "DoOnce", "Gate", "FlipFlop", "MultiGate", "IsValid")
        return L.spawn_standard_macro_node(bp, graph, target, pos)

    # ---- Cast ----
    if kind == "cast":
        cls = _load_class(target)
        if cls is None:
            raise RuntimeError(f"cast: class not found at {target!r}")
        return L.spawn_cast_node(bp, graph, cls, pos)

    # ---- Collections / structs ----
    if kind == "make_array":
        return L.spawn_make_array_node(bp, graph, pos)
    if kind == "make_struct":
        s = _load_struct(target)
        if s is None:
            raise RuntimeError(f"make_struct: struct not found at {target!r}")
        return L.spawn_make_struct_node(bp, graph, s, pos)
    if kind == "break_struct":
        s = _load_struct(target)
        if s is None:
            raise RuntimeError(f"break_struct: struct not found at {target!r}")
        return L.spawn_break_struct_node(bp, graph, s, pos)
    if kind == "set_fields_in_struct":
        s = _load_struct(target)
        if s is None:
            raise RuntimeError(f"set_fields_in_struct: struct not found at {target!r}")
        return L.spawn_set_fields_in_struct_node(bp, graph, s, pos)

    # ---- Misc ----
    if kind == "spawn_actor":
        return L.spawn_spawn_actor_from_class_node(bp, graph, pos)
    if kind == "format_text":
        return L.spawn_format_text_node(bp, graph, pos)
    if kind == "get_subsystem":
        cls = _load_class(target)
        if cls is None:
            raise RuntimeError(f"get_subsystem: class not found at {target!r}")
        return L.spawn_get_subsystem_node(bp, graph, cls, pos)
    if kind == "load_asset":
        return L.spawn_load_asset_node(bp, graph, pos)

    # ---- Delegates ----
    if kind in ("bind_delegate", "unbind_delegate", "assign_delegate"):
        if "." not in target:
            raise RuntimeError(f"{kind} target must be 'Component.Delegate', got {target!r}")
        comp, delegate = target.split(".", 1)
        fn = {
            "bind_delegate":   L.spawn_bind_delegate_node,
            "unbind_delegate": L.spawn_unbind_delegate_node,
            "assign_delegate": L.spawn_assign_delegate_node,
        }[kind]
        return fn(bp, graph, comp, delegate, pos)
    if kind == "call_delegate":
        return L.spawn_call_delegate_node(bp, graph, target, pos)

    # ---- Tier 2 nodes ----
    if kind == "knot":
        return L.spawn_knot_node(bp, graph, pos)
    if kind == "comment":
        p = spec.get("params") or {}
        size = unreal.Vector2D(float(p.get("width", 400)), float(p.get("height", 200)))
        text = str(p.get("text", target or "Comment"))
        return L.spawn_comment_node(bp, graph, pos, size, text)
    if kind == "enum_literal":
        e = _load_enum(target)
        if e is None:
            raise RuntimeError(f"enum_literal: enum not found at {target!r}")
        return L.spawn_enum_literal_node(bp, graph, e, pos)
    if kind == "get_class_defaults":
        cls = _load_class(target)
        if cls is None:
            raise RuntimeError(f"get_class_defaults: class not found at {target!r}")
        return L.spawn_get_class_defaults_node(bp, graph, cls, pos)
    if kind == "function_result":
        return L.spawn_function_result_node(bp, graph, pos)
    if kind == "get_array_item":
        return L.spawn_get_array_item_node(bp, graph, pos)
    if kind == "operator":
        # target = "+", "-", "*", etc.
        return L.spawn_promotable_operator_node(bp, graph, target, pos)

    raise RuntimeError(f"kind {kind!r} not supported")


# ---------------------------------------------------------------------------
# Validation (dry-resolve all targets without spawning)
# ---------------------------------------------------------------------------
def validate(graph: dict) -> dict:
    """Resolve all targets in NodeSpec without mutating UE.

    Returns {ok, missing: [{id, kind, target, reason}], resolved: [{id, ...}]}.
    """
    out = {"ok": True, "missing": [], "resolved": []}
    if unreal is None:
        out["ok"] = False
        out["missing"].append({"reason": "unreal module not importable"})
        return out
    for spec in graph.get("nodes", []):
        kind = spec.get("kind", "")
        target = spec.get("target") or ""
        try:
            if kind == "call_function":
                cls, fn = _resolve_function_owner(target)
                if not cls.find_function(fn):
                    raise RuntimeError(f"function {fn!r} not on {cls.get_name()}")
            elif kind in ("cast", "get_subsystem", "get_class_defaults"):
                if _load_class(target) is None:
                    raise RuntimeError(f"class not found: {target!r}")
            elif kind in ("make_struct", "break_struct", "set_fields_in_struct"):
                if _load_struct(target) is None:
                    raise RuntimeError(f"struct not found: {target!r}")
            elif kind in ("switch_enum", "enum_literal"):
                if _load_enum(target) is None:
                    raise RuntimeError(f"enum not found: {target!r}")
            out["resolved"].append({"id": spec.get("id"), "kind": kind, "target": target})
        except Exception as e:  # noqa: BLE001
            out["ok"] = False
            out["missing"].append({"id": spec.get("id"), "kind": kind, "target": target, "reason": str(e)})
    return out


# ---------------------------------------------------------------------------
# read_graph — dump existing graph via C++ DumpGraphToJson
# ---------------------------------------------------------------------------
def read_graph(bp_path: str, function_name: str = "EventGraph") -> dict:
    try:
        L = _lib()
        if bp_path:
            bp = _load_blueprint(bp_path)
        else:
            bp, bp_path = _active_blueprint()
        g = _find_graph(bp, function_name)
        raw = L.dump_graph_to_json(g)
        return {"ok": True, "bp_path": bp_path, "function": function_name, "graph": json.loads(raw)}
    except Exception as e:  # noqa: BLE001
        return {"ok": False, "error": str(e), "traceback": traceback.format_exc()}


# ---------------------------------------------------------------------------
# Entrypoint (called by the MCP server's bootstrap snippet)
# ---------------------------------------------------------------------------
def entrypoint(payload_json: str) -> str:
    try:
        graph = json.loads(payload_json)
    except Exception as e:  # noqa: BLE001
        result_json = json.dumps({"ok": False, "errors": [f"bad payload json: {e}"]})
        _tee(result_json, {})
        return result_json
    op = graph.get("__op", "apply")
    if op == "read":
        out = read_graph(graph.get("bp_path") or "", graph.get("function") or "EventGraph")
    elif op == "validate":
        out = validate(graph)
    elif op == "batch":
        out = batch_apply(graph.get("ops") or [])
    else:
        out = apply(graph)
    result_json = json.dumps(out)
    _tee(result_json, graph)
    return result_json


def batch_apply(ops: list) -> dict:
    """Run a list of apply()s atomically-ish — same transaction, single undo."""
    results = []
    for i, op in enumerate(ops):
        op_kind = op.get("op") or "apply"
        if op_kind == "apply":
            results.append({"index": i, "op": "apply", "result": apply(op)})
        elif op_kind == "validate":
            results.append({"index": i, "op": "validate", "result": validate(op)})
        elif op_kind == "read":
            results.append({"index": i, "op": "read", "result": read_graph(op.get("bp_path") or "", op.get("function") or "EventGraph")})
        else:
            results.append({"index": i, "op": op_kind, "result": {"ok": False, "error": f"unknown op kind {op_kind!r}"}})
    overall_ok = all(r["result"].get("ok") for r in results)
    return {"ok": overall_ok, "results": results}


def _tee(result_json: str, payload: dict) -> None:
    try:
        if unreal is None:
            return
        proj_dir = unreal.Paths.project_saved_dir()
        out_dir = proj_dir + "MCP/"
        try:
            os.makedirs(out_dir, exist_ok=True)
        except Exception:
            pass
        rid = (payload or {}).get("__rid") or "last"
        # Always write to last.json for human inspection, and to {rid}.json for race-free reads
        with open(out_dir + "last.json", "w", encoding="utf-8") as f:
            f.write(result_json)
        if rid != "last":
            with open(out_dir + f"{rid}.json", "w", encoding="utf-8") as f:
                f.write(result_json)
        unreal.log("[MCP] " + result_json[:1500])
    except Exception:
        pass