instruct_particulate/utils/auto_kinematics_utils.py

from __future__ import annotations

import base64
import json
import mimetypes
import os
import re
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Sequence

from google import genai
from google.genai import types
import numpy as np
from openai import OpenAI
import trimesh

from instruct_particulate.utils.postprocessing_utils import resolve_point_prompt_face_ids


DEFAULT_AUTO_KINEMATICS_AZIMUTHS = (-45.0, 45.0, 135.0, 215.0)
DEFAULT_RENDER_PITCH_DEG = 70.0
DEFAULT_RENDER_ELEVATION_DEG = 90.0 - DEFAULT_RENDER_PITCH_DEG
DEFAULT_RENDER_DISTANCE = 2.0
DEFAULT_RENDER_RESOLUTION = 512
DEFAULT_POINT_PROMPT_SNAP_RADIUS_PX = 24
DEFAULT_BLENDER_RENDER_ENGINE = "CYCLES"
DEFAULT_BLENDER_RENDER_SAMPLES = 256
DEFAULT_RENDER_BACKGROUND_RGB = (192, 192, 192)


@dataclass(frozen=True)
class RenderedView:
    image_id: int
    image_path: Path
    camera_path: Path
    azimuth_deg: float
    elevation_deg: float
    width: int
    height: int
    intrinsic: np.ndarray
    camera_to_world: np.ndarray
    world_to_camera: np.ndarray
    face_ids: np.ndarray
    hit_points: np.ndarray
    normals: np.ndarray
    depth: np.ndarray


def detect_provider(model_id: str) -> str:
    lowered = str(model_id).strip().lower()
    if "gemini" in lowered:
        return "gemini"
    return "openai"


def image_detail_for_model(model_id: str) -> str:
    normalized = str(model_id).strip().lower()
    match = re.match(r"^gpt-(\d+)\.(\d+)", normalized)
    if match is not None:
        major = int(match.group(1))
        minor = int(match.group(2))
        if major > 5 or (major == 5 and minor >= 4):
            return "original"
    return "high"


def _normalize_reasoning_effort(reasoning_effort: str | None) -> str:
    normalized = str(reasoning_effort or "").strip().lower()
    if normalized == "med":
        normalized = "medium"
    if normalized not in {"low", "medium", "high"}:
        return "medium"
    return normalized


def _extract_openai_response_text(response: Any) -> str:
    output_text = getattr(response, "output_text", None)
    if isinstance(output_text, str) and output_text.strip():
        return output_text.strip()

    output = getattr(response, "output", None)
    text_fragments: list[str] = []
    if isinstance(output, list):
        for item in output:
            if item is None:
                continue
            item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None)
            if item_type != "message":
                continue
            content = item.get("content") if isinstance(item, dict) else getattr(item, "content", None)
            if not isinstance(content, list):
                continue
            for part in content:
                if part is None:
                    continue
                part_type = part.get("type") if isinstance(part, dict) else getattr(part, "type", None)
                if part_type not in {"output_text", "text", "input_text"}:
                    continue
                text = part.get("text") if isinstance(part, dict) else getattr(part, "text", None)
                if isinstance(text, str) and text.strip():
                    text_fragments.append(text.strip())
    if text_fragments:
        return "\n".join(text_fragments).strip()

    if isinstance(response, dict):
        fallback = response.get("output_text")
        return fallback.strip() if isinstance(fallback, str) else ""

    model_dump = getattr(response, "model_dump", None)
    if callable(model_dump):
        try:
            dumped = model_dump()
        except Exception:
            dumped = None
        if isinstance(dumped, dict):
            fallback = dumped.get("output_text")
            if isinstance(fallback, str) and fallback.strip():
                return fallback.strip()
    return ""


def extract_json_object(raw_text: str) -> dict[str, Any]:
    cleaned = raw_text.strip()
    if cleaned.startswith("```"):
        lines = cleaned.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].startswith("```"):
            lines = lines[:-1]
        cleaned = "\n".join(lines).strip()

    candidates = [cleaned]
    start = cleaned.find("{")
    end = cleaned.rfind("}")
    if start != -1 and end != -1 and end > start:
        candidates.append(cleaned[start : end + 1])

    for candidate in candidates:
        try:
            parsed = json.loads(candidate)
        except json.JSONDecodeError:
            continue
        if not isinstance(parsed, dict):
            raise ValueError("Model response was valid JSON but not a JSON object.")
        return parsed

    raise ValueError("Model response was not parseable as a JSON object.")


def _load_text(path: Path) -> str:
    if not path.exists():
        raise FileNotFoundError(f"Prompt file does not exist: {path}")
    return path.read_text(encoding="utf-8").strip()


def _guess_mime_type(path: Path) -> str:
    mime_type, _ = mimetypes.guess_type(path.name)
    return mime_type or "application/octet-stream"


def _encode_image_data(path: Path) -> tuple[bytes, str]:
    return path.read_bytes(), _guess_mime_type(path)


def _save_render_with_gray_background(
    source_path: Path,
    output_path: Path,
    *,
    background_rgb: tuple[int, int, int] = DEFAULT_RENDER_BACKGROUND_RGB,
) -> None:
    from PIL import Image

    with Image.open(source_path) as image:
        rgba_image = image.convert("RGBA")
        background = Image.new("RGBA", rgba_image.size, (*background_rgb, 255))
        background.alpha_composite(rgba_image)
        background.convert("RGB").save(output_path)


def resolve_blender_binary(explicit_path: str | os.PathLike[str] | None = None) -> Path:
    if explicit_path is not None:
        candidate = Path(explicit_path).expanduser().resolve()
        if not candidate.is_file():
            raise FileNotFoundError(f"Blender binary does not exist: {candidate}")
        return candidate

    home = Path.home()
    candidates = [
        home / "blender" / "blender",
        home / "blender/blender",
    ]
    for directory in sorted(home.glob("blender*-linux-x64"), reverse=True):
        candidates.append(directory / "blender")

    which_blender = shutil.which("blender")
    if which_blender:
        candidates.append(Path(which_blender))

    for candidate in candidates:
        candidate = candidate.expanduser()
        if candidate.is_file():
            return candidate.resolve()

    raise FileNotFoundError(
        "Could not find Blender. Pass --blender-bin or install Blender under your home directory."
    )


def render_mesh_auto_kinematics_views(
    mesh: trimesh.Trimesh,
    *,
    output_dir: Path,
    azimuths_deg: Sequence[float],
    mesh_path: str | os.PathLike[str] | None = None,
    up_dir: str | None = None,
    pitch_deg: float = DEFAULT_RENDER_PITCH_DEG,
    camera_distance: float = DEFAULT_RENDER_DISTANCE,
    resolution: int = DEFAULT_RENDER_RESOLUTION,
    blender_bin: str | os.PathLike[str] | None = None,
) -> list[RenderedView]:
    output_dir.mkdir(parents=True, exist_ok=True)
    rendered_views: list[RenderedView] = []
    azimuths = [float(azimuth_deg) for azimuth_deg in azimuths_deg]
    blender_path = resolve_blender_binary(blender_bin)
    repo_root = Path(__file__).resolve().parents[2]
    blender_script_path = repo_root / "scripts" / "render_auto_kinematics_blender.py"
    if not blender_script_path.exists():
        raise FileNotFoundError(f"Auto-kinematics Blender script does not exist: {blender_script_path}")

    with tempfile.TemporaryDirectory(prefix="auto_kinematics_renders_") as temp_dir:
        render_output_dir = Path(temp_dir)
        if mesh_path is None:
            render_mesh_path = render_output_dir / "_auto_kinematics_render_mesh.glb"
            mesh.export(render_mesh_path)
        else:
            render_mesh_path = Path(mesh_path).expanduser().resolve()
            if not render_mesh_path.exists():
                raise FileNotFoundError(f"Auto-kinematics render mesh path does not exist: {render_mesh_path}")
        command = [
            str(blender_path),
            "-b",
            "-P",
            str(blender_script_path),
            "--",
            "--mesh-path",
            str(render_mesh_path),
            "--output-dir",
            str(render_output_dir),
            "--resolution",
            str(int(resolution)),
            "--camera-distance",
            repr(float(camera_distance)),
            "--pitch-deg",
            repr(float(pitch_deg)),
            "--engine",
            DEFAULT_BLENDER_RENDER_ENGINE,
            "--samples",
            str(int(DEFAULT_BLENDER_RENDER_SAMPLES)),
            "--azimuths",
            *[repr(float(azimuth_deg)) for azimuth_deg in azimuths],
        ]
        if up_dir is not None and str(up_dir).strip():
            command.append(f"--up-dir={str(up_dir).strip()}")
        print(
            f"Rendering {len(azimuths)} auto-kinematics views with Blender at {blender_path}"
        )
        completed_process = subprocess.run(
            command,
            check=False,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
        )
        blender_output = completed_process.stdout or ""
        if blender_output:
            print(blender_output, end="" if blender_output.endswith("\n") else "\n")
        if completed_process.returncode != 0:
            raise subprocess.CalledProcessError(
                completed_process.returncode,
                command,
                output=blender_output,
            )

        total_views = len(azimuths)
        for image_id, azimuth_deg in enumerate(azimuths):
            temp_image_path = render_output_dir / f"view_{image_id:03d}.png"
            temp_camera_path = render_output_dir / f"view_{image_id:03d}_camera.npz"
            if not temp_image_path.exists():
                raise FileNotFoundError(
                    "Blender render output is missing image: "
                    f"{temp_image_path}\nBlender output:\n{blender_output}"
                )
            if not temp_camera_path.exists():
                raise FileNotFoundError(
                    "Blender render output is missing camera data: "
                    f"{temp_camera_path}\nBlender output:\n{blender_output}"
                )

            image_path = output_dir / temp_image_path.name
            camera_path = output_dir / temp_camera_path.name
            _save_render_with_gray_background(temp_image_path, image_path)
            shutil.copy2(temp_camera_path, camera_path)

            print(
                f"Loading auto-kinematics view {image_id + 1}/{total_views} "
                f"(azimuth={float(azimuth_deg):.1f} deg)"
            )
            camera_payload = np.load(temp_camera_path)
            intrinsic = np.asarray(camera_payload["intrinsic"], dtype=np.float64)
            camera_to_world = np.asarray(camera_payload["camera_to_world"], dtype=np.float64)
            world_to_camera = np.asarray(camera_payload["world_to_camera"], dtype=np.float64)
            face_ids = np.asarray(camera_payload["face_ids"], dtype=np.int32)
            hit_points = np.asarray(camera_payload["hit_points"], dtype=np.float32)
            normals = np.asarray(camera_payload["normals"], dtype=np.float32)
            depth = np.asarray(camera_payload["depth"], dtype=np.float32)
            elevation_deg = (
                float(camera_payload["elevation_deg"])
                if "elevation_deg" in camera_payload.files
                else float(DEFAULT_RENDER_ELEVATION_DEG)
            )
            rendered_views.append(
                RenderedView(
                    image_id=int(image_id),
                    image_path=image_path,
                    camera_path=camera_path,
                    azimuth_deg=float(azimuth_deg),
                    elevation_deg=elevation_deg,
                    width=int(resolution),
                    height=int(resolution),
                    intrinsic=intrinsic.astype(np.float32),
                    camera_to_world=camera_to_world.astype(np.float32),
                    world_to_camera=world_to_camera.astype(np.float32),
                    face_ids=face_ids.astype(np.int32, copy=False),
                    hit_points=hit_points.astype(np.float32, copy=False),
                    normals=normals.astype(np.float32, copy=False),
                    depth=depth.astype(np.float32, copy=False),
                )
            )
    return rendered_views


def build_auto_kinematics_user_prompt(rendered_views: Sequence[RenderedView]) -> str:
    lines = [
        "Analyze the rendered object views and infer the rigid-part kinematic structure.",
        "Return only a valid JSON object matching the schema in the system instruction.",
        "Use only rigid parts. Do not include non-rigid or decorative subparts.",
        "Do not omit small independently moving rigid parts; controls or attachments such as knobs, buttons, switches, sliders, levers, handles, trays, lids, doors, or latches must each be their own link when they move independently.",
        "For every link, choose one representative visible surface point.",
        "For every point_prompt, return integer x and y bins in the inclusive range [0, 1000].",
        "These bins represent normalized image coordinates with top-left origin:",
        "x_normalized = x / 1000 and y_normalized = y / 1000.",
        "So (0, 0) is the top-left corner and (1000, 1000) is the bottom-right corner.",
        "Each image is identified by the explicit image ID text that appears immediately before it.",
        "Point prompts for different links do not need to come from the same image.",
        "For each link independently, choose the image where that link is most visible and easiest to localize.",
        f"Number of images: {len(rendered_views)}.",
    ]
    for view in rendered_views:
        lines.append(
            f"Image ID {view.image_id}: square render with size {view.width}x{view.height} pixels."
        )
    return "\n".join(lines)


def _get_openai_api_key(explicit_api_key: str | None) -> str | None:
    return explicit_api_key or os.getenv("OPENAI_API_KEY")


def _get_gemini_api_key(explicit_api_key: str | None) -> str | None:
    return explicit_api_key or os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")


def call_auto_kinematics_model(
    *,
    model_id: str,
    system_prompt_path: Path,
    rendered_views: Sequence[RenderedView],
    provider: str | None = None,
    api_key: str | None = None,
    reasoning_effort: str = "medium",
) -> dict[str, Any]:
    resolved_provider = detect_provider(model_id) if provider is None else provider
    system_prompt = _load_text(system_prompt_path)
    user_prompt = build_auto_kinematics_user_prompt(rendered_views)
    if resolved_provider == "openai":
        return _call_openai(
            model_id=model_id,
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            rendered_views=rendered_views,
            api_key=api_key,
            reasoning_effort=reasoning_effort,
        )
    if resolved_provider == "gemini":
        return _call_gemini(
            model_id=model_id,
            system_prompt=system_prompt,
            user_prompt=user_prompt,
            rendered_views=rendered_views,
            api_key=api_key,
            reasoning_effort=reasoning_effort,
        )
    raise ValueError(f"Unsupported provider: {resolved_provider}")


def _call_openai(
    *,
    model_id: str,
    system_prompt: str,
    user_prompt: str,
    rendered_views: Sequence[RenderedView],
    api_key: str | None,
    reasoning_effort: str,
) -> dict[str, Any]:
    resolved_api_key = _get_openai_api_key(api_key)
    if not resolved_api_key:
        raise EnvironmentError("Missing OpenAI API key. Set OPENAI_API_KEY.")

    content: list[dict[str, Any]] = [{"type": "input_text", "text": user_prompt}]
    image_detail = image_detail_for_model(model_id)
    for view in rendered_views:
        content.append(
            {
                "type": "input_text",
                "text": (
                    f"Image ID {view.image_id}: the next image corresponds to this ID. "
                    "Use this integer in point_prompt.image_id."
                ),
            }
        )
        image_bytes, mime_type = _encode_image_data(view.image_path)
        encoded = base64.b64encode(image_bytes).decode("ascii")
        content.append(
            {
                "type": "input_image",
                "image_url": f"data:{mime_type};base64,{encoded}",
                "detail": image_detail,
            }
        )

    client = OpenAI(api_key=resolved_api_key)
    request_payload: dict[str, Any] = {
        "model": model_id,
        "reasoning": {"effort": _normalize_reasoning_effort(reasoning_effort)},
        "text": {"format": {"type": "json_object"}},
        "input": [
            {
                "role": "system",
                "content": [{"type": "input_text", "text": system_prompt}],
            },
            {
                "role": "user",
                "content": content,
            },
        ],
        "store": False,
    }
    response = client.responses.create(**request_payload)
    content_text = _extract_openai_response_text(response)
    if not content_text:
        raise ValueError("Unexpected OpenAI response format: response output text was empty.")
    return extract_json_object(content_text)


def _call_gemini(
    *,
    model_id: str,
    system_prompt: str,
    user_prompt: str,
    rendered_views: Sequence[RenderedView],
    api_key: str | None,
    reasoning_effort: str,
) -> dict[str, Any]:
    resolved_api_key = _get_gemini_api_key(api_key)
    if not resolved_api_key:
        raise EnvironmentError("Missing Gemini API key. Set GEMINI_API_KEY or GOOGLE_API_KEY.")

    client = genai.Client(api_key=resolved_api_key)
    parts: list[Any] = [types.Part.from_text(text=user_prompt)]
    for view in rendered_views:
        parts.append(
            types.Part.from_text(
                text=(
                    f"Image ID {view.image_id}: the next image corresponds to this ID. "
                    "Use this integer in point_prompt.image_id."
                )
            )
        )
        image_bytes, mime_type = _encode_image_data(view.image_path)
        parts.append(types.Part.from_bytes(data=image_bytes, mime_type=mime_type))

    contents = [types.Content(role="user", parts=parts)]
    config_kwargs: dict[str, Any] = {
        "system_instruction": system_prompt,
        "response_mime_type": "application/json",
        "temperature": 0.1,
    }
    config_kwargs["thinking_config"] = types.ThinkingConfig(
        thinking_level=_normalize_reasoning_effort(reasoning_effort)
    )
    response = client.models.generate_content(
        model=model_id,
        contents=contents,
        config=types.GenerateContentConfig(**config_kwargs),
    )
    response_text = getattr(response, "text", None)
    if not response_text:
        raise ValueError("Gemini returned an empty response.")
    return extract_json_object(response_text)


def parse_auto_kinematics_response(response: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(response, dict):
        raise TypeError(f"Expected a JSON object, got {type(response).__name__}")
    total_links = int(response.get("total_links", 0))
    raw_links = response.get("links")
    raw_joints = response.get("joints")
    if not isinstance(raw_links, list):
        raise ValueError("Response must contain a 'links' list.")
    if not isinstance(raw_joints, list):
        raise ValueError("Response must contain a 'joints' list.")
    if total_links != len(raw_links):
        raise ValueError(
            f"total_links={total_links} does not match len(links)={len(raw_links)}"
        )

    links: list[dict[str, Any]] = []
    seen_link_ids: set[int] = set()
    for raw_link in raw_links:
        if not isinstance(raw_link, dict):
            raise ValueError("Each entry in 'links' must be an object.")
        link_id = int(raw_link.get("link_id", -1))
        if link_id in seen_link_ids:
            raise ValueError(f"Duplicate link_id {link_id} in response.")
        seen_link_ids.add(link_id)
        link_name = str(raw_link.get("name", "")).strip()
        if not link_name:
            raise ValueError(f"Link {link_id} is missing a non-empty name.")
        point_prompt = raw_link.get("point_prompt")
        if not isinstance(point_prompt, dict):
            raise ValueError(f"Link {link_id} is missing a point_prompt object.")
        image_id = int(point_prompt.get("image_id", -1))
        x_bin = point_prompt.get("x")
        y_bin = point_prompt.get("y")
        if x_bin is None or y_bin is None:
            coordinate = point_prompt.get("coordinate")
            if (
                not isinstance(coordinate, (list, tuple))
                or len(coordinate) != 2
            ):
                raise ValueError(
                    "Link "
                    f"{link_id} point_prompt must contain integer x/y bins or a length-2 coordinate array, "
                    f"got {point_prompt!r}"
                )
            x_bin = coordinate[0]
            y_bin = coordinate[1]
        x_bin = int(round(float(x_bin)))
        y_bin = int(round(float(y_bin)))
        if not (0 <= x_bin <= 1000 and 0 <= y_bin <= 1000):
            raise ValueError(
                f"Link {link_id} point_prompt bins must lie in [0, 1000], got x={x_bin}, y={y_bin}"
            )
        links.append(
            {
                "link_id": link_id,
                "name": link_name,
                "point_prompt": {
                    "image_id": image_id,
                    "x": x_bin,
                    "y": y_bin,
                    "normalized_coordinate": [
                        float(x_bin) / 1000.0,
                        float(y_bin) / 1000.0,
                    ],
                },
            }
        )
    links.sort(key=lambda link: int(link["link_id"]))
    expected_link_ids = list(range(len(links)))
    observed_link_ids = [int(link["link_id"]) for link in links]
    if observed_link_ids != expected_link_ids:
        raise ValueError(
            f"Link IDs must be continuous from 0, got {observed_link_ids}"
        )

    joints: list[dict[str, Any]] = []
    seen_joint_ids: set[int] = set()
    for raw_joint in raw_joints:
        if not isinstance(raw_joint, dict):
            raise ValueError("Each entry in 'joints' must be an object.")
        joint_id = int(raw_joint.get("joint_id", -1))
        if joint_id in seen_joint_ids:
            raise ValueError(f"Duplicate joint_id {joint_id} in response.")
        seen_joint_ids.add(joint_id)
        parent_link_id = int(raw_joint.get("parent_link_id", -1))
        child_link_id = int(raw_joint.get("child_link_id", -1))
        joint_type = str(raw_joint.get("joint_type", "")).strip().lower()
        if joint_type not in {"revolute", "prismatic"}:
            raise ValueError(
                f"Joint {joint_id} must have joint_type 'revolute' or 'prismatic', got {joint_type!r}"
            )
        joints.append(
            {
                "joint_id": joint_id,
                "parent_link_id": parent_link_id,
                "child_link_id": child_link_id,
                "joint_type": joint_type,
            }
        )
    joints.sort(key=lambda joint: int(joint["joint_id"]))
    expected_joint_ids = list(range(len(joints)))
    observed_joint_ids = [int(joint["joint_id"]) for joint in joints]
    if observed_joint_ids != expected_joint_ids:
        raise ValueError(
            f"Joint IDs must be continuous from 0, got {observed_joint_ids}"
        )

    link_names = [str(link["name"]) for link in links]
    joint_specs = [
        (
            int(joint["parent_link_id"]),
            int(joint["child_link_id"]),
            str(joint["joint_type"]),
        )
        for joint in joints
    ]
    return {
        "object_name": str(response.get("object_name", "")).strip(),
        "total_links": total_links,
        "links": links,
        "joints": joints,
        "link_names": link_names,
        "joint_specs": joint_specs,
    }


def _resolve_prompt_pixel(
    view: RenderedView,
    *,
    point_prompt: dict[str, Any],
    max_snap_radius_px: int,
) -> tuple[int, int, float, float, float]:
    raw_x = float(np.clip(float(point_prompt["x"]) / 1000.0, 0.0, 1.0)) * float(view.width - 1)
    raw_y = float(np.clip(float(point_prompt["y"]) / 1000.0, 0.0, 1.0)) * float(view.height - 1)
    requested_x = int(np.floor(raw_x + 0.5))
    requested_y = int(np.floor(raw_y + 0.5))
    raw_x = float(np.clip(raw_x, 0.0, float(view.width - 1)))
    raw_y = float(np.clip(raw_y, 0.0, float(view.height - 1)))
    clipped_x = int(np.clip(requested_x, 0, view.width - 1))
    clipped_y = int(np.clip(requested_y, 0, view.height - 1))
    if view.face_ids[clipped_y, clipped_x] >= 0:
        return clipped_x, clipped_y, float(np.hypot(clipped_x - raw_x, clipped_y - raw_y)), raw_x, raw_y

    valid_pixels = np.argwhere(view.face_ids >= 0)
    if valid_pixels.size == 0:
        raise ValueError(f"Rendered view {view.image_id} does not contain any visible mesh pixels.")
    deltas = valid_pixels - np.asarray([clipped_y, clipped_x], dtype=np.int32)
    squared_distances = (deltas.astype(np.float64) ** 2).sum(axis=1)
    best_index = int(np.argmin(squared_distances))
    best_distance = float(np.sqrt(squared_distances[best_index]))
    if best_distance > float(max_snap_radius_px):
        raise ValueError(
            "Point prompt lies too far from the rendered object surface: "
            f"view={view.image_id}, requested=({raw_x:.2f}, {raw_y:.2f}), nearest_distance={best_distance:.2f}px"
        )
    best_y, best_x = valid_pixels[best_index]
    return int(best_x), int(best_y), best_distance, raw_x, raw_y


def lift_point_prompts_from_rendered_views(
    *,
    links: Sequence[dict[str, Any]],
    rendered_views: Sequence[RenderedView],
    max_snap_radius_px: int = DEFAULT_POINT_PROMPT_SNAP_RADIUS_PX,
) -> dict[str, Any]:
    views_by_id = {int(view.image_id): view for view in rendered_views}
    lifted_points: list[np.ndarray] = []
    lifted_normals: list[np.ndarray] = []
    debug_records: list[dict[str, Any]] = []
    for link in links:
        point_prompt = link["point_prompt"]
        image_id = int(point_prompt["image_id"])
        if image_id not in views_by_id:
            raise ValueError(
                f"Link {link['link_id']} references image_id={image_id}, which is not among the rendered views."
            )
        view = views_by_id[image_id]
        resolved_x, resolved_y, snap_distance_px, requested_x_px, requested_y_px = _resolve_prompt_pixel(
            view,
            point_prompt=point_prompt,
            max_snap_radius_px=max_snap_radius_px,
        )
        point = np.asarray(view.hit_points[resolved_y, resolved_x], dtype=np.float32)
        normal = np.asarray(view.normals[resolved_y, resolved_x], dtype=np.float32)
        if not np.isfinite(point).all():
            raise ValueError(
                f"Resolved point prompt for link {link['link_id']} did not map to a finite 3D point."
            )
        normal_norm = float(np.linalg.norm(normal))
        if not np.isfinite(normal).all() or normal_norm <= 1e-8:
            raise ValueError(
                f"Resolved point prompt for link {link['link_id']} did not map to a valid normal."
            )
        normal = normal / normal_norm
        lifted_points.append(point)
        lifted_normals.append(normal.astype(np.float32, copy=False))
        debug_records.append(
            {
                "link_id": int(link["link_id"]),
                "name": str(link["name"]),
                "image_id": image_id,
                "requested_coordinate_bins": [
                    int(point_prompt["x"]),
                    int(point_prompt["y"]),
                ],
                "requested_coordinate_normalized": [
                    float(point_prompt["x"]) / 1000.0,
                    float(point_prompt["y"]) / 1000.0,
                ],
                "requested_coordinate_pixels": [
                    float(requested_x_px),
                    float(requested_y_px),
                ],
                "resolved_coordinate": [int(resolved_x), int(resolved_y)],
                "snap_distance_px": float(snap_distance_px),
                "point_prompt_render_normalized": point.astype(np.float32).tolist(),
                "point_prompt_normalized": point.astype(np.float32).tolist(),
                "normal": normal.astype(np.float32).tolist(),
            }
        )
    return {
        "points": np.stack(lifted_points, axis=0).astype(np.float32),
        "normals": np.stack(lifted_normals, axis=0).astype(np.float32),
        "records": debug_records,
    }


def _denormalize_points(
    points: np.ndarray,
    *,
    center: np.ndarray,
    scale: float,
) -> np.ndarray:
    return (np.asarray(points, dtype=np.float32) / np.float32(scale) + center).astype(
        np.float32,
        copy=False,
    )


def prepare_lifted_prompt_records_for_saving(
    *,
    normalized_mesh: trimesh.Trimesh,
    lifted: dict[str, Any],
    center: np.ndarray,
    scale: float,
    render_to_model_rotation: np.ndarray | None = None,
) -> dict[str, Any]:
    """Transforms lifted prompt records into model space and annotates face IDs."""
    prompt_points = np.asarray(lifted["points"], dtype=np.float32)
    prompt_normals = np.asarray(lifted["normals"], dtype=np.float32)
    if render_to_model_rotation is not None:
        rotation = np.asarray(render_to_model_rotation, dtype=np.float32)
        prompt_points = prompt_points @ rotation.T
        prompt_normals = prompt_normals @ rotation.T
    normal_norms = np.linalg.norm(prompt_normals, axis=1, keepdims=True)
    normal_norms = np.clip(normal_norms, a_min=1e-8, a_max=None)
    prompt_normals = prompt_normals / normal_norms

    records = list(lifted["records"])
    for record, point, normal in zip(records, prompt_points, prompt_normals, strict=True):
        record["point_prompt_normalized"] = point.astype(np.float32).tolist()
        record["normal"] = normal.astype(np.float32).tolist()

    prompt_face_ids = resolve_point_prompt_face_ids(normalized_mesh, prompt_points)
    for record, face_id in zip(records, prompt_face_ids.tolist(), strict=True):
        record["point_prompt_face_id"] = int(face_id)

    point_prompts_world = _denormalize_points(
        prompt_points,
        center=np.asarray(center, dtype=np.float32),
        scale=float(scale),
    )
    return {
        "point_prompts": prompt_points.astype(np.float32, copy=False),
        "point_prompt_normals": prompt_normals.astype(np.float32, copy=False),
        "point_prompt_face_ids": prompt_face_ids.astype(np.int64, copy=False),
        "point_prompts_world": point_prompts_world.astype(np.float32, copy=False),
        "records": records,
    }


def save_auto_kinematics_artifacts(
    *,
    output_dir: Path,
    raw_response: dict[str, Any],
    parsed_response: dict[str, Any],
    lifted_prompt_records: Sequence[dict[str, Any]],
    center: np.ndarray,
    scale: float,
    visualization_mesh_relative_path: str = "_auto_kinematics_visualization_mesh.glb",
) -> dict[str, Path]:
    output_dir.mkdir(parents=True, exist_ok=True)
    raw_response_path = output_dir / "auto_kinematics_raw_response.json"
    parsed_response_path = output_dir / "auto_kinematics_parsed.json"
    lifted_prompts_path = output_dir / "auto_kinematics_lifted_point_prompts.json"
    raw_response_path.write_text(json.dumps(raw_response, indent=2) + "\n", encoding="utf-8")
    parsed_response_path.write_text(json.dumps(parsed_response, indent=2) + "\n", encoding="utf-8")

    mesh_path_pointer = str(visualization_mesh_relative_path)
    lifted_payload = {
        "normalization": {
            "center": np.asarray(center, dtype=np.float32).tolist(),
            "scale": float(scale),
        },
        "visualization_mesh_path": mesh_path_pointer,
        "point_prompt_face_mesh_path": mesh_path_pointer,
        "links": list(lifted_prompt_records),
    }
    for link_record in lifted_payload["links"]:
        if "point_prompt_face_id" in link_record and "point_prompt_face_mesh_path" not in link_record:
            link_record["point_prompt_face_mesh_path"] = mesh_path_pointer
        if "point_prompt_world" not in link_record:
            point_prompt_normalized = np.asarray(
                link_record["point_prompt_normalized"],
                dtype=np.float32,
            )
            point_prompt_world = _denormalize_points(
                point_prompt_normalized[None, :],
                center=np.asarray(center, dtype=np.float32),
                scale=float(scale),
            )[0]
            link_record["point_prompt_world"] = point_prompt_world.tolist()
    lifted_prompts_path.write_text(
        json.dumps(lifted_payload, indent=2) + "\n",
        encoding="utf-8",
    )
    return {
        "raw_response_path": raw_response_path,
        "parsed_response_path": parsed_response_path,
        "lifted_prompts_path": lifted_prompts_path,
    }