"""Bake3D for Blender: MCP client, generation, motion and native GLB import.
Install this file through Preferences > Add-ons > Install from Disk.
API keys are held in the current Blender session, never written to the .blend file.
"""
bl_info = {"name": "Bake3D", "author": "Bake3D", "version": (1, 2, 0),
           "blender": (4, 2, 0), "location": "View3D > Sidebar > Bake3D",
           "description": "Generate and import rigged, animated characters with Bake3D", "category": "Import-Export"}
import json
import math
import os
import tempfile
import urllib.error
import urllib.parse
import urllib.request


class Client:
    def __init__(self, api_key=None, origin="https://bake3d.ai"):
        self.origin = 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 os.environ.get("BAKE3D_API_KEY", "")
        if not self.api_key.startswith("bk_live_"):
            raise ValueError("Create an API key in Bake3D Account > API keys")
        self.sequence = 0
        self._rpc("initialize", {"protocolVersion": "2025-11-25", "capabilities": {},
                                  "clientInfo": {"name": "bake3d-blender", "version": "1.2.0"}})

    def _rpc(self, method, params):
        self.sequence += 1
        payload = json.dumps({"jsonrpc": "2.0", "id": self.sequence, "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": "2025-11-25"})
        try:
            with urllib.request.urlopen(req, timeout=180) as res:
                data = json.load(res)
        except urllib.error.HTTPError as error:
            raise RuntimeError("Bake3D HTTP " + str(error.code) + ": " + error.read(4096).decode(errors="replace")) from None
        if "error" in data:
            raise RuntimeError(data["error"]["message"])
        result = data["result"]
        if result.get("isError"):
            raise RuntimeError("; ".join(c.get("text", "") for c in result.get("content", [])))
        return result

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

    def import_character(self, project_id, animation_id=None, preset_id=None, root_motion=False, rest_pose=False):
        """Can also be called from a Blender MCP bridge's Python execution tool."""
        import bpy
        args = {"projectId": project_id, "rootMotion": root_motion}
        if rest_pose:
            args["restPose"] = True
        elif animation_id:
            args["animationId"] = animation_id
        elif preset_id:
            args["presetId"] = preset_id
        result = self.call("bake3d_export", **args)
        url = result["downloadUrl"]
        # Signed downloads never receive the account's API key.
        if urllib.parse.urlparse(url).scheme not in {"http", "https"}:
            raise ValueError("Invalid download URL")
        path = None
        try:
            with urllib.request.urlopen(url, timeout=60) as response, tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as output:
                path = output.name
                count = 0
                while True:
                    chunk = response.read(1024 * 1024)
                    if not chunk:
                        break
                    count += len(chunk)
                    if count > 64 * 1024 * 1024:
                        raise ValueError("GLB exceeds the 64 MiB editor download limit")
                    output.write(chunk)
            before = set(bpy.data.objects)
            bpy.ops.import_scene.gltf(filepath=path, import_pack_images=True)
            imported = [obj for obj in bpy.data.objects if obj not in before]
            for obj in imported:
                obj["bake3d_project_id"] = project_id
            # glTF times are seconds; Blender converts them using the scene's current FPS.
            # Changing FPS after import would change playback speed for every existing action.
            for obj in imported:
                if obj.animation_data and obj.animation_data.action:
                    bpy.context.scene.frame_start = min(bpy.context.scene.frame_start, math.floor(obj.animation_data.action.frame_range[0]))
                    bpy.context.scene.frame_end = max(1, math.ceil(obj.animation_data.action.frame_range[1]))
            return imported
        finally:
            if path and os.path.exists(path):
                os.unlink(path)


# UI is optional: importing Client from another Python application does not require Blender.
try:
    import bpy
except ImportError:
    bpy = None

if bpy:
    class Bake3DState(bpy.types.PropertyGroup):
        api_key: bpy.props.StringProperty(name="API key", subtype='PASSWORD', options={'SKIP_SAVE'})
        prompt: bpy.props.StringProperty(name="Character", default="A ceramic fox standing on four legs")
        project_id: bpy.props.StringProperty(name="Project ID")
        job_id: bpy.props.StringProperty(name="Job ID")
        motion: bpy.props.StringProperty(name="Motion", default="Look around curiously, then settle")
        animation_id: bpy.props.StringProperty(name="Animation ID")
        preset_id: bpy.props.StringProperty(name="Preset ID", description="Optional; blank motion fields use the included initial animation")
        rest_pose: bpy.props.BoolProperty(name="Rest pose only", description="Import a copy without animation clips")
        status: bpy.props.StringProperty(default="Connect your Bake3D account with an API key.")

    def client_for(context):
        return Client(context.window_manager.bake3d.api_key)

    class BAKE3D_OT_generate(bpy.types.Operator):
        bl_idname = "bake3d.generate"
        bl_label = "Generate character (35 credits)"
        def execute(self, context):
            state = context.window_manager.bake3d
            try:
                result = client_for(context).call("bake3d_generate", prompt=state.prompt, initialAnimation=not state.rest_pose)
                state.project_id = result["projectId"]
                state.job_id = result["job"]["id"]
                state.animation_id = ""
                state.status = "Generation started. Check the job for progress."
                return {'FINISHED'}
            except Exception as error:
                self.report({'ERROR'}, str(error))
                return {'CANCELLED'}

    class BAKE3D_OT_check(bpy.types.Operator):
        bl_idname = "bake3d.check_job"
        bl_label = "Check job"
        def execute(self, context):
            state = context.window_manager.bake3d
            try:
                job = client_for(context).call("bake3d_get_job", jobId=state.job_id)["job"]
                state.status = str(job["status"]) + " · " + str(job["progress"]) + "% · " + (job.get("error") or job.get("stage") or "")
                return {'FINISHED'}
            except Exception as error:
                self.report({'ERROR'}, str(error))
                return {'CANCELLED'}

    class BAKE3D_OT_animate(bpy.types.Operator):
        bl_idname = "bake3d.animate"
        bl_label = "Compose motion (5 credits)"
        def execute(self, context):
            state = context.window_manager.bake3d
            try:
                result = client_for(context).call("bake3d_animate", projectId=state.project_id, prompt=state.motion, durationSeconds=3, loop=True)
                state.animation_id = result["animation"]["id"]
                state.rest_pose = False
                state.status = "Motion saved. Import it to preview the animation."
                return {'FINISHED'}
            except Exception as error:
                self.report({'ERROR'}, str(error))
                return {'CANCELLED'}

    class BAKE3D_OT_initial(bpy.types.Operator):
        bl_idname = "bake3d.initial_animation"
        bl_label = "Create / retry included animation (free)"
        def execute(self, context):
            state = context.window_manager.bake3d
            try:
                result = client_for(context).call("bake3d_initial_animation", projectId=state.project_id)
                state.animation_id = result.get("animationId", "")
                state.preset_id = ""
                state.rest_pose = False
                if result.get("job"):
                    state.job_id = result["job"]["id"]
                    state.status = "Initial animation started. Check its job, then import."
                else:
                    state.status = "Initial animation is ready to import."
                return {'FINISHED'}
            except Exception as error:
                self.report({'ERROR'}, str(error))
                return {'CANCELLED'}

    class BAKE3D_OT_list_animations(bpy.types.Operator):
        bl_idname = "bake3d.list_animations"
        bl_label = "Use latest saved animation (free)"
        bl_description = "List this project's saved animations and select the newest one"
        def execute(self, context):
            state = context.window_manager.bake3d
            try:
                result = client_for(context).call("bake3d_list_animations", projectId=state.project_id)
                animations = result.get("animations", [])
                if not animations:
                    state.status = "No saved animations yet. Compose motion or create the included animation."
                    return {'FINISHED'}
                state.animation_id = animations[0]["id"]
                state.preset_id = ""
                state.rest_pose = False
                state.status = "Selected " + str(animations[0]["name"]) + ". Saved: " + "; ".join(
                    str(a["name"]) + " (" + str(a.get("duration")) + "s, " + str(a.get("kind")) + ")" for a in animations[:6])
                return {'FINISHED'}
            except Exception as error:
                self.report({'ERROR'}, str(error))
                return {'CANCELLED'}

    class BAKE3D_OT_import(bpy.types.Operator):
        bl_idname = "bake3d.import_character"
        bl_label = "Import character and motion"
        bl_options = {'REGISTER', 'UNDO'}
        project_id: bpy.props.StringProperty()
        animation_id: bpy.props.StringProperty()
        preset_id: bpy.props.StringProperty()
        def execute(self, context):
            state = context.window_manager.bake3d
            try:
                imported = client_for(context).import_character(self.project_id or state.project_id,
                    self.animation_id or state.animation_id or None, self.preset_id or state.preset_id or None, rest_pose=state.rest_pose)
                state.status = "Imported " + str(len(imported)) + " objects. Play the timeline to preview."
                return {'FINISHED'}
            except Exception as error:
                self.report({'ERROR'}, str(error))
                return {'CANCELLED'}

    class BAKE3D_PT_panel(bpy.types.Panel):
        bl_label = "Bake3D"
        bl_idname = "BAKE3D_PT_panel"
        bl_space_type = 'VIEW_3D'
        bl_region_type = 'UI'
        bl_category = 'Bake3D'
        def draw(self, context):
            layout = self.layout
            state = context.window_manager.bake3d
            layout.prop(state, "api_key")
            layout.operator("wm.url_open", text="Get an API key").url = "https://bake3d.ai/account/api"
            layout.separator()
            layout.prop(state, "prompt")
            layout.operator("bake3d.generate")
            layout.prop(state, "job_id")
            layout.operator("bake3d.check_job")
            layout.separator()
            layout.prop(state, "project_id")
            layout.operator("bake3d.initial_animation")
            layout.prop(state, "motion")
            layout.operator("bake3d.animate")
            layout.operator("bake3d.list_animations")
            layout.prop(state, "animation_id")
            layout.prop(state, "preset_id")
            layout.prop(state, "rest_pose")
            layout.operator("bake3d.import_character")
            layout.label(text=state.status)

    classes = (Bake3DState, BAKE3D_OT_generate, BAKE3D_OT_check, BAKE3D_OT_animate, BAKE3D_OT_initial, BAKE3D_OT_list_animations, BAKE3D_OT_import, BAKE3D_PT_panel)
    def register():
        for cls in classes:
            bpy.utils.register_class(cls)
        bpy.types.WindowManager.bake3d = bpy.props.PointerProperty(type=Bake3DState, options={'SKIP_SAVE'})

    def unregister():
        del bpy.types.WindowManager.bake3d
        for cls in reversed(classes):
            bpy.utils.unregister_class(cls)
