zcbot/windows-node/adapters/blender.scene.author@v3/blender_worker.py

"""Blender-hosted executor for managed Scene Recipe jobs."""

from __future__ import annotations

import hashlib
import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import bpy

sys.path.insert(0, str(Path(__file__).resolve().parent))

from scene_recipe import SceneBuilder

ADAPTER_VERSION = "1.0.0"
PREVIEW_FILENAMES = {
    f"preview_{index:02d}": f"preview-{index:02d}.png" for index in range(1, 9)
}


def atomic_json(path: Path, value: Any) -> None:
    temporary = path.with_name(path.name + ".tmp-" + os.urandom(8).hex())
    try:
        with temporary.open("w", encoding="utf-8", newline="\n") as handle:
            json.dump(value, handle, ensure_ascii=False, indent=2)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
    finally:
        temporary.unlink(missing_ok=True)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def terminal(
    record: dict[str, Any],
    status: str,
    artifacts: list[dict[str, Any]],
    error: dict[str, str],
) -> dict[str, Any]:
    return {
        "job_id": record.get("job_id", ""),
        "lease_id": record.get("lease_id", ""),
        "request_digest": record.get("request_digest", ""),
        "status": status,
        "error": error,
        "artifact_manifest": artifacts,
        "terminal_at": datetime.now(timezone.utc).isoformat(),
    }


def artifact(path: Path, artifact_id: str, media_type: str) -> dict[str, Any]:
    if not path.is_file() or path.stat().st_size == 0:
        raise RuntimeError(f"OUTPUT_EXPORT_EMPTY:{path.name}")
    if path.suffix.lower() == ".png" and path.read_bytes()[:8] != b"\x89PNG\r\n\x1a\n":
        raise RuntimeError("PNG_EXPORT_INVALID")
    return {
        "artifact_id": artifact_id,
        "filename": path.name,
        "media_type": media_type,
        "size_bytes": path.stat().st_size,
        "sha256": sha256(path),
    }


def validate_semantics(request: dict[str, Any], spec: dict[str, Any]) -> list[dict[str, Any]]:
    if request.get("inputs") or spec.get("type") != "recipe" or spec.get("recipe_version") != 1:
        raise ValueError("BLENDER_REQUEST_SEMANTICS_INVALID")
    views = list(spec["render"]["views"])
    output_keys = [item["output_key"] for item in views]
    if len(output_keys) != len(set(output_keys)) or "preview_01" not in output_keys:
        raise ValueError("BLENDER_RENDER_VIEWS_INVALID")
    requested = {
        item.get("key")
        for item in request["outputs"]
        if isinstance(item, dict) and isinstance(item.get("key"), str)
    }
    expected = set(output_keys) - {"preview_01"}
    if requested - {"scene_glb"} != expected:
        raise ValueError("BLENDER_RENDER_OUTPUTS_MISMATCH")
    return views


def execute(job_dir: Path, record: dict[str, Any]) -> list[dict[str, Any]]:
    request = record["request"]
    spec = request["operation"]["scene"]
    views = validate_semantics(request, spec)
    builder = SceneBuilder(spec)
    builder.compile()

    output = job_dir / "output"
    output.mkdir(exist_ok=True)
    project = output / "scene.blend"
    manifest = output / "scene-manifest.json"
    provenance = output / "provenance.json"
    bpy.ops.wm.save_as_mainfile(filepath=str(project), check_existing=False)

    rendered: list[tuple[Path, str]] = []
    for view in views:
        builder.configure_camera(view)
        output_id = view["output_key"]
        preview = output / PREVIEW_FILENAMES[output_id]
        bpy.context.scene.render.filepath = str(preview)
        bpy.ops.render.render(write_still=True)
        rendered.append((preview, output_id))

    atomic_json(manifest, builder.manifest(views))
    workspace = json.loads((job_dir / "workspace.json").read_text(encoding="utf-8"))
    atomic_json(
        provenance,
        {
            "adapter_version": ADAPTER_VERSION,
            "blender_version": bpy.app.version_string,
            "request_digest": record["request_digest"],
            "workspace_mode": workspace.get("mode"),
            "generator": "scene_recipe@1",
        },
    )

    artifacts = [
        artifact(project, "project", "application/x-blender"),
        *[artifact(path, output_id, "image/png") for path, output_id in rendered],
        artifact(manifest, "scene_manifest", "application/json"),
    ]
    if any(item.get("key") == "scene_glb" for item in request["outputs"]):
        glb = output / "scene.glb"
        bpy.ops.export_scene.gltf(
            filepath=str(glb),
            export_format="GLB",
            export_apply=True,
            use_renderable=True,
        )
        artifacts.append(artifact(glb, "scene_glb", "model/gltf-binary"))
    artifacts.append(artifact(provenance, "provenance", "application/json"))
    return artifacts


def job_directory() -> Path:
    arguments = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    if len(arguments) != 1:
        raise ValueError("BLENDER_JOB_DIRECTORY_MISSING")
    return Path(arguments[0]).resolve(strict=True)


def main() -> int:
    job_dir = job_directory()
    record: dict[str, Any] = {}
    try:
        record = json.loads((job_dir / "request" / "request.json").read_text(encoding="utf-8"))
        artifacts = execute(job_dir, record)
        atomic_json(job_dir / "artifacts.json", artifacts)
        atomic_json(job_dir / "terminal.json", terminal(record, "succeeded", artifacts, {}))
        print("[OK] Blender scene recipe completed")
        return 0
    except Exception as exc:
        code = str(exc).split(":", 1)[0] if str(exc) else type(exc).__name__
        atomic_json(
            job_dir / "terminal.json",
            terminal(record, "failed", [], {"code": code[:80], "detail": str(exc)[:500]}),
        )
        print(f"[ERR] {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())