"""Bake3D for Unreal Engine 5: MCP client, generation, motion and Interchange GLB import.

Installation
  1. Edit > Plugins: enable "Python Editor Script Plugin" and "Editor Scripting Utilities". Keep the
     "Interchange Framework" plugin enabled; it provides the glTF translator used for import. Restart.
  2. Copy this file to <YourProject>/Content/Python/bake3d_unreal.py. Unreal adds that folder to sys.path.
  3. Create <YourProject>/Content/Python/init_unreal.py containing the single line
         import bake3d_unreal
     so the Tools > Bake3D menu appears on start, or run that line in the Output Log's Python console.
  4. Provide a key: start the editor with BAKE3D_API_KEY in its environment, or run
         bake3d_unreal.set_api_key("bk_live_...")
     in the Python console. Keys are held in memory for this editor session only; nothing is written
     to disk, to project settings or to the editor's configuration files.

Console usage (Output Log > Python, or Window > Python REPL)
  bake3d_unreal.account()
  bake3d_unreal.list_projects("fox")
  job = bake3d_unreal.generate(prompt="A ceramic fox standing on four legs")   # 35 credits
  job = bake3d_unreal.generate(image_url="https://example.com/fox.png")       # 35 credits
  bake3d_unreal.wait_for_job(job["jobId"])                                     # blocks the editor while polling
  bake3d_unreal.list_animations(job["projectId"])
  bake3d_unreal.list_presets()
  bake3d_unreal.import_character(project_id)                       # included initial motion
  bake3d_unreal.import_character(project_id, animation_id="...")   # a saved animation
  bake3d_unreal.import_character(project_id, preset_id="tail-wag") # a built-in preset (free)
  bake3d_unreal.import_character(project_id, rest_pose=True)       # rigged, unanimated

Each import downloads a GLB into Content/Bake3D and imports it through Interchange into /Game/Bake3D,
creating a Skeletal Mesh, a Skeleton, materials, textures and Animation Sequences. Interchange's glTF
translator converts metres to centimetres and Y-up to Z-up; this script leaves the pipeline's import
offset at 1.0 so that conversion is not applied twice. Requires Python 3.9 or newer (UE 5.3+).
"""
import json
import os
import re
import ssl
import time
import urllib.error
import urllib.parse
import urllib.request

try:
    import unreal
except ImportError:  # Allows the client to be imported and tested outside the editor.
    unreal = None

VERSION = "1.0.0"
DEFAULT_ORIGIN = "https://bake3d.ai"
PROTOCOL = "2025-11-25"
MAX_JSON_BYTES = 4 * 1024 * 1024
MAX_GLB_BYTES = 64 * 1024 * 1024
ASSET_PATH = "/Game/Bake3D"
_SESSION = {"api_key": "", "origin": DEFAULT_ORIGIN}


def _log(message):
    if unreal is not None:
        unreal.log("[Bake3D] " + message)
    else:
        print("[Bake3D] " + message)


def _warn(message):
    if unreal is not None:
        unreal.log_warning("[Bake3D] " + message)
    else:
        print("[Bake3D] " + message)


def _dialog(title, message):
    """Editor Scripting Utilities dialog when available; the Output Log otherwise."""
    dialog = getattr(unreal, "EditorDialog", None) if unreal is not None else None
    if dialog is not None and hasattr(dialog, "show_message"):
        dialog.show_message(title, message, unreal.AppMsgType.OK)
    else:
        _log(title + ": " + message)


def set_api_key(key):
    """Keep a bk_live_ key for this editor session. It is never persisted."""
    key = (key or "").strip()
    if not key.startswith("bk_live_") or any(c.isspace() for c in key):
        raise ValueError("Create an API key in Bake3D Account > API keys (Pro or Studio); it starts with bk_live_")
    _SESSION["api_key"] = key
    _log("API key set for this editor session.")


def set_origin(origin):
    """Development only: point the client at a local Bake3D server (http://localhost:4321)."""
    _SESSION["origin"] = origin.rstrip("/")


class Client:
    """Minimal MCP Streamable HTTP client: JSON responses, no persistent session."""

    def __init__(self, api_key=None, origin=None):
        self.origin = (origin or _SESSION["origin"]).rstrip("/")
        parsed = urllib.parse.urlparse(self.origin)
        if parsed.scheme != "https" and not (parsed.scheme == "http" and parsed.hostname in {"localhost", "127.0.0.1"}):
            raise ValueError("Use HTTPS for Bake3D, or localhost for development")
        self.api_key = api_key or _SESSION["api_key"] or os.environ.get("BAKE3D_API_KEY", "")
        if not self.api_key.startswith("bk_live_"):
            raise ValueError("No Bake3D API key. Run bake3d_unreal.set_api_key('bk_live_...') or set BAKE3D_API_KEY before starting the editor")
        self.context = ssl.create_default_context(cafile=os.environ.get("BAKE3D_CA_BUNDLE") or None)
        self.sequence = 0
        self._rpc("initialize", {"protocolVersion": PROTOCOL, "capabilities": {},
                                 "clientInfo": {"name": "bake3d-unreal", "version": VERSION}})

    def _rpc(self, method, params):
        self.sequence += 1
        request_id = self.sequence
        payload = json.dumps({"jsonrpc": "2.0", "id": request_id, "method": method, "params": params}).encode()
        req = urllib.request.Request(self.origin + "/api/v1/mcp", data=payload, headers={
            "Content-Type": "application/json", "Accept": "application/json, text/event-stream",
            "Authorization": "Bearer " + self.api_key, "MCP-Protocol-Version": PROTOCOL})
        try:
            with urllib.request.urlopen(req, timeout=180, context=self.context) as res:
                body = res.read(MAX_JSON_BYTES + 1)
        except urllib.error.HTTPError as error:
            raise RuntimeError("Bake3D HTTP " + str(error.code) + ": " + error.read(4096).decode(errors="replace")) from None
        if len(body) > MAX_JSON_BYTES:
            raise RuntimeError("Bake3D response exceeded the 4 MiB limit")
        data = json.loads(body)
        if not isinstance(data, dict) or data.get("id") != request_id:
            raise RuntimeError("Bake3D returned an invalid MCP response")
        if "error" in data:
            raise RuntimeError(str(data["error"].get("message", "MCP request failed")))
        result = data.get("result")
        if not isinstance(result, dict):
            raise RuntimeError("Bake3D returned no result")
        if result.get("isError"):
            raise RuntimeError("; ".join(c.get("text", "") for c in result.get("content", []) if isinstance(c, dict)))
        return result

    def call(self, name, **arguments):
        result = self._rpc("tools/call", {"name": name, "arguments": arguments})
        if isinstance(result.get("structuredContent"), dict):
            return result["structuredContent"]
        return json.loads(result["content"][0]["text"])

    def download(self, url, path):
        """Signed downloads never receive the API key and must come from the Bake3D file endpoint."""
        if not url.startswith(self.origin + "/api/files/"):
            raise ValueError("The export did not return a Bake3D download URL")
        count = 0
        try:
            with urllib.request.urlopen(url, timeout=120, context=self.context) as response, open(path, "wb") as output:
                while True:
                    chunk = response.read(1024 * 1024)
                    if not chunk:
                        break
                    count += len(chunk)
                    if count > MAX_GLB_BYTES:
                        raise ValueError("GLB exceeds the 64 MiB editor download limit")
                    output.write(chunk)
        except Exception:
            if os.path.exists(path):
                os.unlink(path)
            raise
        with open(path, "rb") as handle:
            header = handle.read(12)
        if len(header) < 12 or header[:4] != b"glTF" or int.from_bytes(header[8:12], "little") != count:
            os.unlink(path)
            raise ValueError("The download is not a complete glTF 2.0 binary; export again to refresh the link")
        return count


def client():
    return Client()


def account():
    """Credit balance, plan and action costs."""
    result = client().call("bake3d_account")
    _log("Credits: " + str(result.get("credits")) + " · plan " + str(result.get("plan")) + " · model " + str(result.get("model")))
    return result


def list_projects(query="", filter="all", limit=25, cursor=None):
    """Search your library. Free. filter: all, active, ready or needs-work."""
    args = {"q": query, "filter": filter, "sort": "updated", "limit": limit}
    if cursor:
        args["cursor"] = cursor
    result = client().call("bake3d_list_projects", **args)
    for project in result.get("projects", []):
        _log(str(project.get("id")) + "  " + str(project.get("title")) + "  " + str(project.get("status")) + ("  rigged" if project.get("hasRig") else ""))
    return result


def generate(prompt=None, image_url=None, title=None, initial_animation=True, quality="standard"):
    """Start one paid generation (25 mesh + 10 rig credits). Returns projectId and jobId."""
    if bool(prompt) == bool(image_url):
        raise ValueError("Pass exactly one of prompt or image_url")
    args = {"rig": True, "quality": quality, "initialAnimation": initial_animation}
    if prompt:
        args["prompt"] = prompt
    else:
        args["imageUrl"] = image_url
    if title:
        args["title"] = title
    result = client().call("bake3d_generate", **args)
    _log("Generation accepted: project " + str(result["projectId"]) + ", job " + str(result["job"]["id"]) + ", " + str(result.get("creditsCharged")) + " credits")
    return {"projectId": result["projectId"], "jobId": result["job"]["id"]}


def check_job(job_id):
    """One free status read."""
    job = client().call("bake3d_get_job", jobId=job_id)["job"]
    _log("Job " + str(job_id) + ": " + str(job.get("status")) + " · " + str(job.get("progress")) + "% · " + str(job.get("error") or job.get("stage") or ""))
    return job


def wait_for_job(job_id, timeout_seconds=900, interval_seconds=4):
    """Poll until the job ends. Blocks the editor; the slow-task dialog's Cancel stops polling, not the job."""
    if unreal is None:
        return _poll_job(job_id, timeout_seconds, interval_seconds, None)
    with unreal.ScopedSlowTask(100, "Bake3D job " + str(job_id)) as slow:
        slow.make_dialog(True)
        return _poll_job(job_id, timeout_seconds, interval_seconds, slow)


def _poll_job(job_id, timeout_seconds, interval_seconds, slow):
    mcp = client()
    started = time.monotonic()
    progress = 0
    while True:
        job = mcp.call("bake3d_get_job", jobId=job_id)["job"]
        if slow is not None:
            step = max(0, min(100, int(job.get("progress") or 0)) - progress)
            slow.enter_progress_frame(step, str(job.get("stage") or job.get("status")))
            progress += step
        if job.get("status") in {"succeeded", "failed", "canceled"}:
            _log("Job " + str(job_id) + " " + str(job["status"]) + (": " + str(job["error"]) if job.get("error") else ""))
            return job
        if slow is not None and slow.should_cancel():
            _warn("Stopped waiting; the job continues on Bake3D. Call check_job later.")
            return job
        if time.monotonic() - started > timeout_seconds:
            _warn("Timed out waiting; the job continues on Bake3D. Call check_job later.")
            return job
        time.sleep(interval_seconds)


def list_animations(project_id):
    """Saved animations of a project, newest first. Free."""
    result = client().call("bake3d_list_animations", projectId=project_id)
    for animation in result.get("animations", []):
        _log(str(animation.get("id")) + "  " + str(animation.get("name")) + "  " + str(animation.get("duration")) + "s  " + str(animation.get("kind")))
    return result


def list_presets():
    """Built-in motion presets. Exporting one is free; compatibility is checked against the rig."""
    result = client().call("bake3d_list_presets")
    for preset in result.get("presets", []):
        _log(str(preset.get("id")) + "  " + str(preset.get("name")) + ("  adaptive" if preset.get("adaptive") else ""))
    return result


def animate(project_id, prompt, duration_seconds=4, loop=True):
    """Describe new motion for the rig (5 credits). Returns the saved animation ID."""
    result = client().call("bake3d_animate", projectId=project_id, prompt=prompt, durationSeconds=duration_seconds, loop=loop)
    _log("Motion saved: " + str(result["animation"]["id"]))
    return result["animation"]["id"]


def content_folder():
    base = unreal.Paths.convert_relative_path_to_full(unreal.Paths.project_content_dir())
    folder = os.path.join(base, "Bake3D")
    os.makedirs(folder, exist_ok=True)
    return folder


def _unique_path(folder, stem):
    stem = re.sub(r"[^A-Za-z0-9_-]+", "_", stem).strip("_") or "character"
    path = os.path.join(folder, stem + ".glb")
    counter = 2
    while os.path.exists(path):
        path = os.path.join(folder, stem + "-" + str(counter) + ".glb")
        counter += 1
    return path


def export_glb(project_id, animation_id=None, preset_id=None, rest_pose=False, root_motion=False, fps=30, folder=None):
    """Request a free export and download it into Content/Bake3D. Returns the local GLB path."""
    args = {"projectId": project_id, "rootMotion": root_motion, "fps": fps}
    if rest_pose:
        args["restPose"] = True
        suffix = "rest"
    elif animation_id:
        args["animationId"] = animation_id
        suffix = "anim-" + animation_id
    elif preset_id:
        args["presetId"] = preset_id
        suffix = "preset-" + preset_id
    else:
        suffix = "initial"
    mcp = client()
    result = mcp.call("bake3d_export", **args)
    path = _unique_path(folder or content_folder(), project_id + "-" + suffix)
    size = mcp.download(result["downloadUrl"], path)
    _log("Downloaded " + str(size) + " bytes to " + path)
    return path


def interchange_available():
    return unreal is not None and hasattr(unreal, "InterchangeGenericAssetsPipeline") and hasattr(unreal, "InterchangeGLTFTranslator")


def _set(target, name, value):
    try:
        target.set_editor_property(name, value)
        return True
    except Exception as error:  # Property names vary slightly between engine versions.
        _warn("Interchange option " + name + " not applied: " + str(error))
        return False


def import_glb(path, destination=ASSET_PATH, scale=1.0):
    """Import a GLB through Interchange as a Skeletal Mesh with Skeleton, materials and Animation Sequences."""
    if not interchange_available():
        raise RuntimeError("The glTF Interchange importer is not available. Enable the Interchange Framework plugin under Edit > Plugins, restart the editor and import again. The file is kept at " + path)
    pipeline = unreal.InterchangeGenericAssetsPipeline()
    _set(pipeline, "import_offset_uniform_scale", scale)
    force = getattr(getattr(unreal, "InterchangeForceMeshType", None), "IFMT_SKELETAL_MESH", None)
    if force is not None:
        _set(pipeline.get_editor_property("common_meshes_properties"), "force_all_mesh_as_type", force)
    skeletal = pipeline.get_editor_property("common_skeletal_meshes_and_animations_properties")
    _set(skeletal, "import_only_animations", False)
    _set(skeletal, "import_meshes_in_bone_hierarchy", True)
    mesh = pipeline.get_editor_property("mesh_pipeline")
    _set(mesh, "import_skeletal_meshes", True)
    _set(mesh, "import_static_meshes", False)
    _set(mesh, "create_physics_asset", False)
    animation = pipeline.get_editor_property("animation_pipeline")
    _set(animation, "import_animations", True)
    _set(animation, "import_bone_tracks", True)
    material = pipeline.get_editor_property("material_pipeline")
    _set(material, "import_materials", True)
    task = unreal.AssetImportTask()
    task.set_editor_property("filename", path)
    task.set_editor_property("destination_path", destination)
    task.set_editor_property("automated", True)
    task.set_editor_property("save", True)
    task.set_editor_property("replace_existing", False)
    task.set_editor_property("options", pipeline)
    unreal.AssetToolsHelpers.get_asset_tools().import_asset_tasks([task])
    imported = [str(p) for p in task.get_editor_property("imported_object_paths")]
    if not imported:
        raise RuntimeError("Interchange imported nothing from " + path + ". Check the Output Log and Message Log for importer errors")
    _log("Imported " + str(len(imported)) + " assets into " + destination + ": " + ", ".join(imported))
    return imported


def import_character(project_id, animation_id=None, preset_id=None, rest_pose=False, root_motion=False, fps=30, destination=ASSET_PATH):
    """Export, download and import in one call. Pass nothing but the project ID for the included initial motion."""
    path = export_glb(project_id, animation_id, preset_id, rest_pose, root_motion, fps)
    return import_glb(path, destination)


def _latest_ready_project():
    projects = client().call("bake3d_list_projects", filter="ready", sort="updated", limit=1).get("projects", [])
    if not projects:
        raise RuntimeError("No rigged project found. Generate one first: bake3d_unreal.generate(prompt=...)")
    return projects[0]


def _guarded(action):
    try:
        action()
    except Exception as error:
        _warn(str(error))
        _dialog("Bake3D", str(error))


def menu_set_key():
    _dialog("Bake3D API key", "Run  bake3d_unreal.set_api_key(\"bk_live_...\")  in the Output Log's Python console, or start the editor with BAKE3D_API_KEY set. Keys stay in memory for this session. Create keys at bake3d.ai/account/api (Pro or Studio).")


def menu_account():
    _guarded(lambda: _dialog("Bake3D account", json.dumps(account(), indent=2)))


def menu_list_projects():
    _guarded(lambda: _dialog("Bake3D projects", "\n".join(str(p.get("id")) + "  " + str(p.get("title")) for p in list_projects(filter="ready").get("projects", [])) or "No rigged projects yet."))


def menu_import_latest():
    def run():
        project = _latest_ready_project()
        import_character(project["id"])
        _dialog("Bake3D", "Imported " + str(project.get("title")) + " with its included motion into " + ASSET_PATH)
    _guarded(run)


def menu_import_latest_rest():
    def run():
        project = _latest_ready_project()
        import_character(project["id"], rest_pose=True)
        _dialog("Bake3D", "Imported the rest pose of " + str(project.get("title")) + " into " + ASSET_PATH)
    _guarded(run)


MENU_ENTRIES = [
    ("SetKey", "Set API key…", "menu_set_key"),
    ("Account", "Account and credits", "menu_account"),
    ("ListProjects", "List rigged projects", "menu_list_projects"),
    ("ImportLatest", "Import latest project with included motion", "menu_import_latest"),
    ("ImportLatestRest", "Import latest project rest pose", "menu_import_latest_rest"),
]


def register_menu():
    """Tools > Bake3D. Console functions accept IDs and prompts; menu entries cover the common paths."""
    menus = unreal.ToolMenus.get()
    tools = menus.find_menu("LevelEditor.MainMenu.Tools")
    if tools is None:
        _warn("Tools menu not found; use the console functions instead")
        return False
    submenu = tools.add_sub_menu(tools.menu_name, "Bake3D", "Bake3D", "Bake3D", "Generate and import rigged, animated characters")
    for name, label, function in MENU_ENTRIES:
        entry = unreal.ToolMenuEntry(name=name, type=unreal.MultiBlockType.MENU_ENTRY)
        entry.set_label(label)
        entry.set_string_command(unreal.ToolMenuStringCommandType.PYTHON, "", "import bake3d_unreal; bake3d_unreal." + function + "()")
        submenu.add_menu_entry("Bake3D", entry)
    menus.refresh_all_widgets()
    return True


if unreal is not None:
    try:
        register_menu()
    except Exception as menu_error:
        _warn("Menu registration failed: " + str(menu_error))