Research-Stack/5-Applications/text-to-cad/skills/cad/scripts/snapshot/cli.py

from __future__ import annotations

import argparse
import math
import zlib
from dataclasses import dataclass
from pathlib import Path

import numpy as np

REPO_ROOT = Path.cwd().resolve()
CAD_ROOT = REPO_ROOT
DEFAULT_MODEL_COLOR = (0.80, 0.84, 0.90)
DEFAULT_BACKGROUND_COLOR = (0.98, 0.985, 0.99)
FALLBACK_COMPONENT_COLORS: tuple[tuple[float, float, float], ...] = (
    (0.82, 0.84, 0.88),
    (0.68, 0.77, 0.91),
    (0.70, 0.86, 0.79),
    (0.93, 0.79, 0.62),
    (0.88, 0.72, 0.78),
    (0.76, 0.72, 0.90),
    (0.85, 0.83, 0.62),
    (0.68, 0.86, 0.87),
)
MAX_RENDER_TRIANGLES_PER_MESH = 12000
FEATURE_EDGE_ANGLE_DEG = 32.0
BASE_MARGIN_PX = 12.0
CROP_PADDING_PX = 12
AXIS_BOX_SIZE_PX = 46
VIEW_OUTPUT_ORDER = ("isometric", "front", "back", "right", "left", "top", "bottom")


@dataclass(frozen=True)
class CameraView:
    name: str
    direction: tuple[float, float, float]
    up: tuple[float, float, float]


VIEW_PRESETS: dict[str, CameraView] = {
    "front": CameraView(name="front", direction=(0.0, 0.0, 1.0), up=(0.0, 1.0, 0.0)),
    "back": CameraView(name="back", direction=(0.0, 0.0, -1.0), up=(0.0, 1.0, 0.0)),
    "right": CameraView(name="right", direction=(1.0, 0.0, 0.0), up=(0.0, 1.0, 0.0)),
    "left": CameraView(name="left", direction=(-1.0, 0.0, 0.0), up=(0.0, 1.0, 0.0)),
    "top": CameraView(name="top", direction=(0.0, 1.0, 0.0), up=(0.0, 0.0, 1.0)),
    "bottom": CameraView(name="bottom", direction=(0.0, -1.0, 0.0), up=(0.0, 0.0, 1.0)),
    "isometric": CameraView(name="isometric", direction=(1.0, 1.0, 1.0), up=(0.0, 1.0, 0.0)),
}


@dataclass(frozen=True)
class MeshInstance:
    vertices: np.ndarray
    triangles: np.ndarray
    color_rgb: tuple[float, float, float] | None = None


@dataclass(frozen=True)
class ProjectedMeshInstance:
    screen_points: np.ndarray
    view_points: np.ndarray
    triangles: np.ndarray
    face_brightness: np.ndarray
    face_normals: np.ndarray
    feature_edges: tuple[tuple[int, int], ...]
    color_rgb: tuple[float, float, float]


def _rgb_default(rgb: tuple[float, float, float]) -> str:
    return ",".join(str(channel) for channel in rgb)


def parse_rgb(raw_value: str) -> tuple[float, float, float]:
    parts = [part.strip() for part in raw_value.split(",")]
    if len(parts) != 3:
        raise ValueError(f"Invalid RGB value: {raw_value}")
    rgb = tuple(float(part) for part in parts)
    if not all(0.0 <= channel <= 1.0 for channel in rgb):
        raise ValueError(f"RGB values must be in range [0, 1]: {raw_value}")
    return rgb  # type: ignore[return-value]


def resolve_view(view: str | CameraView) -> CameraView:
    return VIEW_PRESETS[view] if isinstance(view, str) else view


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Render CAD snapshot PNGs from GLB or STL mesh files. Prefer generated assembly GLBs."
    )
    parser.add_argument(
        "input",
        type=Path,
        help="Path to a part GLB, generated assembly GLB, or STL mesh.",
    )
    parser.add_argument("--out", type=Path, help="Write one PNG snapshot to this path.")
    parser.add_argument(
        "--out-dir",
        type=Path,
        help="Write one PNG per view into this directory. Required with --views.",
    )
    parser.add_argument(
        "--view",
        choices=sorted(VIEW_PRESETS),
        default=None,
        help="Camera preset. Defaults to isometric unless --align-ref is used.",
    )
    parser.add_argument(
        "--align-ref",
        help="Resolve an @cad[...] face or edge ref and choose the closest orthographic view automatically.",
    )
    parser.add_argument(
        "--views",
        help="Comma-separated camera presets, or 'all', for batched snapshots. Requires --out-dir.",
    )
    parser.add_argument("--width", type=int, default=1400, help="Maximum output width")
    parser.add_argument("--height", type=int, default=900, help="Maximum output height")
    parser.add_argument(
        "--color",
        default=_rgb_default(DEFAULT_MODEL_COLOR),
        help="Model RGB in 0..1, e.g. '0.80,0.84,0.90'",
    )
    parser.add_argument(
        "--background",
        default=_rgb_default(DEFAULT_BACKGROUND_COLOR),
        help="Background RGB in 0..1, e.g. '0.98,0.985,0.99'",
    )
    parser.add_argument(
        "--edges",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Overlay visible feature edges. Default: true",
    )
    parser.add_argument(
        "--axes",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Show orientation axes in a reserved inset. Default: true",
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    render_jobs = _resolve_render_jobs(args, parser)

    mesh_instances = load_mesh_instances(args.input)
    if not mesh_instances:
        raise RuntimeError(f"No mesh geometry found in {args.input}")

    model_color = parse_rgb(args.color)
    background_color = parse_rgb(args.background)
    for view_name, png_out in render_jobs:
        render_mesh_instances(
            mesh_instances,
            png_out=png_out,
            view=view_name,
            width=args.width,
            height=args.height,
            model_color=model_color,
            background_color=background_color,
            edges=bool(args.edges),
            axes=bool(args.axes),
        )
        label = f"{view_name} " if len(render_jobs) > 1 else ""
        print(f"saved {label}png: {png_out.resolve()}")
    return 0


def _resolve_render_jobs(args: argparse.Namespace, parser: argparse.ArgumentParser) -> list[tuple[str, Path]]:
    views = str(args.views or "").strip()
    if views:
        if args.view:
            parser.error("--views cannot be combined with --view")
        if args.align_ref:
            parser.error("--views cannot be combined with --align-ref")
        if args.out:
            parser.error("--views writes multiple files; use --out-dir instead of --out")
        if not args.out_dir:
            parser.error("--views requires --out-dir")
        try:
            view_names = _parse_views_arg(views)
        except ValueError as exc:
            parser.error(str(exc))
        output_stem = _snapshot_output_stem(args.input)
        return [(view_name, args.out_dir / f"{output_stem}-{view_name}.png") for view_name in view_names]

    if args.out_dir:
        parser.error("--out-dir requires --views")
    if not args.out:
        parser.error("--out is required unless --views and --out-dir are used")
    view_name = str(args.view or "").strip() or _resolve_aligned_view_name(args.align_ref) or "isometric"
    return [(view_name, args.out)]


def _parse_views_arg(raw_value: str) -> tuple[str, ...]:
    if raw_value.strip().lower() == "all":
        return VIEW_OUTPUT_ORDER
    view_names: list[str] = []
    for raw_part in raw_value.split(","):
        view_name = raw_part.strip()
        if not view_name:
            continue
        if view_name not in VIEW_PRESETS:
            allowed = ", ".join(VIEW_OUTPUT_ORDER)
            raise ValueError(f"Unknown snapshot view {view_name!r}; expected one of: {allowed}, or all")
        if view_name not in view_names:
            view_names.append(view_name)
    if not view_names:
        raise ValueError("--views must include at least one view")
    return tuple(view_names)


def _snapshot_output_stem(input_path: Path) -> str:
    stem = input_path.stem
    if input_path.name.lower() == "model.glb" and input_path.parent.name.startswith("."):
        stem = input_path.parent.name[1:]
        for suffix in (".step", ".stp"):
            if stem.lower().endswith(suffix):
                stem = stem[: -len(suffix)]
                break
    cleaned = "".join(character if character.isalnum() or character in {"-", "_"} else "_" for character in stem)
    cleaned = cleaned.strip("_")
    return cleaned or "snapshot"


def _resolve_aligned_view_name(cad_ref: str | None) -> str | None:
    if not cad_ref:
        return None
    from cadref import analysis as cadref_analysis
    from cadref.inspect import CadRefError, inspect_cad_refs

    try:
        result = inspect_cad_refs(cad_ref, facts=True)
    except CadRefError as exc:
        raise ValueError(str(exc)) from exc
    tokens = result.get("tokens")
    if not isinstance(tokens, list) or len(tokens) != 1:
        raise ValueError(f"Failed to resolve render alignment ref: {cad_ref}")
    token = tokens[0]
    if not isinstance(token, dict):
        raise ValueError(f"Failed to resolve render alignment ref: {cad_ref}")
    selections = token.get("selections")
    if not isinstance(selections, list) or len(selections) != 1:
        raise ValueError("snapshot --align-ref expects exactly one face or edge ref.")
    selection = selections[0]
    if not isinstance(selection, dict) or selection.get("status") != "resolved":
        raise ValueError(f"snapshot failed to resolve {cad_ref}")
    selector_type = str(selection.get("selectorType") or "")
    if selector_type not in {"face", "edge"}:
        raise ValueError("snapshot --align-ref only supports face or edge refs.")
    geometry_facts = selection.get("geometryFacts")
    if not isinstance(geometry_facts, dict):
        raise ValueError(f"snapshot could not derive geometry facts for {cad_ref}")
    view_name = cadref_analysis.aligned_view_name_for_facts(selector_type, geometry_facts)
    if not view_name:
        raise ValueError(f"snapshot could not determine an aligned view for {cad_ref}")
    return view_name


def load_mesh_instances(input_path: Path) -> list[MeshInstance]:
    resolved_input = _resolve_cad_path(input_path, kind="input")
    lowered = resolved_input.name.lower()
    if lowered.endswith(".py"):
        raise ValueError(
            "Python assembly inputs are no longer supported by snapshot; "
            "render the generated assembly GLB instead."
        )
    if lowered.endswith(".glb"):
        return [_read_glb_mesh(resolved_input)]
    if lowered.endswith(".stl"):
        return [_read_stl_mesh(resolved_input)]
    raise ValueError(f"Unsupported snapshot input: {input_path}")


def render_mesh_instances(
    mesh_instances: list[MeshInstance],
    *,
    png_out: Path,
    view: str,
    width: int,
    height: int,
    model_color: tuple[float, float, float],
    background_color: tuple[float, float, float],
    edges: bool,
    axes: bool,
) -> None:
    active_instances = [instance for instance in mesh_instances if _instance_has_geometry(instance)]
    if not active_instances:
        raise RuntimeError("No renderable mesh geometry found")

    projected_instances, right, true_up = _project_instances(
        active_instances,
        view=view,
        width=width,
        height=height,
        model_color=model_color,
        include_edges=edges,
    )
    image = _render_scene(
        projected_instances,
        width=width,
        height=height,
        background_color=background_color,
        edges=edges,
        axes=axes,
        right=right,
        true_up=true_up,
    )
    _write_png(image, png_out)


def _resolve_cad_path(path: Path, *, kind: str) -> Path:
    resolved = path.resolve() if path.is_absolute() else (Path.cwd() / path).resolve()
    if not resolved.exists():
        raise FileNotFoundError(f"snapshot {kind} not found: {path}")
    lowered = resolved.name.lower()
    if lowered.endswith((".glb", ".stl")):
        return resolved
    return resolved


def _read_glb_mesh(
    glb_path: Path,
    *,
    transform: object | None = None,
    mesh_cache: dict[Path, MeshInstance] | None = None,
) -> MeshInstance:
    import trimesh

    resolved_path = glb_path.resolve()
    base_mesh = mesh_cache.get(resolved_path) if mesh_cache is not None else None
    if base_mesh is None:
        loaded = trimesh.load(resolved_path, force="scene")
        if isinstance(loaded, trimesh.Scene):
            mesh = loaded.to_geometry()
        elif isinstance(loaded, trimesh.Trimesh):
            mesh = loaded
        else:
            raise RuntimeError(f"No GLB geometry loaded from {resolved_path}")
        if mesh.vertices.size <= 0 or mesh.faces.size <= 0:
            raise RuntimeError(f"No GLB geometry loaded from {resolved_path}")
        vertices = np.asarray(mesh.vertices, dtype=np.float64) * 1000.0
        triangles = np.asarray(mesh.faces, dtype=np.int64)
        base_mesh = MeshInstance(vertices=vertices, triangles=triangles)
        if mesh_cache is not None:
            mesh_cache[resolved_path] = base_mesh

    if transform is None:
        return base_mesh
    return MeshInstance(
        vertices=_apply_transform(base_mesh.vertices, transform),
        triangles=base_mesh.triangles,
        color_rgb=base_mesh.color_rgb,
    )


def _read_stl_mesh(
    stl_path: Path,
    *,
    transform: object | None = None,
    mesh_cache: dict[Path, MeshInstance] | None = None,
) -> MeshInstance:
    vtk, vtk_to_numpy = _vtk_modules()

    resolved_path = stl_path.resolve()
    base_mesh = mesh_cache.get(resolved_path) if mesh_cache is not None else None
    if base_mesh is None:
        reader = vtk.vtkSTLReader()
        reader.SetFileName(str(resolved_path))
        reader.Update()

        polydata = vtk.vtkPolyData()
        polydata.ShallowCopy(reader.GetOutput())
        if polydata.GetNumberOfPoints() <= 0:
            raise RuntimeError(f"No STL geometry loaded from {resolved_path}")

        preview_polydata = _prepare_preview_polydata(polydata)
        triangles = _triangle_indices(preview_polydata)
        points = preview_polydata.GetPoints()
        if points is None or points.GetNumberOfPoints() <= 0:
            raise RuntimeError(f"No STL point data loaded from {resolved_path}")
        vertices = np.asarray(vtk_to_numpy(points.GetData()), dtype=np.float64)
        base_mesh = MeshInstance(vertices=vertices, triangles=triangles)
        if mesh_cache is not None:
            mesh_cache[resolved_path] = base_mesh

    if transform is None:
        return base_mesh
    return MeshInstance(
        vertices=_apply_transform(base_mesh.vertices, transform),
        triangles=base_mesh.triangles,
        color_rgb=base_mesh.color_rgb,
    )


def _vtk_modules() -> tuple[object, object]:
    import vtk
    from vtk.util.numpy_support import vtk_to_numpy

    return vtk, vtk_to_numpy


def _prepare_preview_polydata(polydata: object) -> object:
    vtk, _vtk_to_numpy = _vtk_modules()

    triangle_filter = vtk.vtkTriangleFilter()
    triangle_filter.SetInputData(polydata)
    triangle_filter.Update()

    current = vtk.vtkPolyData()
    current.ShallowCopy(triangle_filter.GetOutput())
    triangle_count = max(current.GetNumberOfPolys(), current.GetNumberOfCells())
    if triangle_count > MAX_RENDER_TRIANGLES_PER_MESH:
        decimator = vtk.vtkQuadricDecimation()
        decimator.SetInputData(current)
        decimator.SetTargetReduction(
            max(0.0, min(0.99, 1.0 - (MAX_RENDER_TRIANGLES_PER_MESH / float(triangle_count))))
        )
        volume_preservation = getattr(decimator, "VolumePreservationOn", None)
        if callable(volume_preservation):
            volume_preservation()
        decimator.Update()
        current.ShallowCopy(decimator.GetOutput())

        retriangulate = vtk.vtkTriangleFilter()
        retriangulate.SetInputData(current)
        retriangulate.Update()
        current.ShallowCopy(retriangulate.GetOutput())

    return current


def _triangle_indices(polydata: object) -> np.ndarray:
    _vtk, vtk_to_numpy = _vtk_modules()

    polys = polydata.GetPolys()
    if polys is None or polys.GetNumberOfCells() <= 0 or polys.GetData() is None:
        raise RuntimeError("STL mesh does not contain polygon cells")

    raw = np.asarray(vtk_to_numpy(polys.GetData()), dtype=np.int64)
    if raw.size % 4 != 0 or np.any(raw[::4] != 3):
        raise RuntimeError("STL mesh contains non-triangular faces after preprocessing")
    return raw.reshape(-1, 4)[:, 1:]


def _apply_transform(vertices: np.ndarray, transform: object) -> np.ndarray:
    if not isinstance(transform, (list, tuple)) or len(transform) != 16:
        raise ValueError("manifest transform must be a 16-number array")
    matrix = np.asarray([float(value) for value in transform], dtype=np.float64).reshape(4, 4)
    homogeneous = np.concatenate([vertices, np.ones((vertices.shape[0], 1), dtype=np.float64)], axis=1)
    transformed = homogeneous @ matrix.T
    w = transformed[:, 3:4]
    safe_w = np.where(np.abs(w) > 1e-12, w, 1.0)
    return transformed[:, :3] / safe_w


def _instance_has_geometry(instance: MeshInstance) -> bool:
    return instance.vertices.size > 0 and instance.triangles.size > 0


def _project_instances(
    mesh_instances: list[MeshInstance],
    *,
    view: str,
    width: int,
    height: int,
    model_color: tuple[float, float, float],
    include_edges: bool,
) -> tuple[list[ProjectedMeshInstance], np.ndarray, np.ndarray]:
    resolved_view = resolve_view(view)
    right, true_up, view_normal = _camera_basis(resolved_view.direction, resolved_view.up)

    all_vertices = np.concatenate([instance.vertices for instance in mesh_instances], axis=0)
    center = 0.5 * (all_vertices.min(axis=0) + all_vertices.max(axis=0))

    projected_xyz: list[np.ndarray] = []
    x_min = math.inf
    x_max = -math.inf
    y_min = math.inf
    y_max = -math.inf
    for instance in mesh_instances:
        relative = instance.vertices - center
        view_points = np.column_stack(
            (
                relative @ right,
                relative @ true_up,
                relative @ view_normal,
            )
        )
        projected_xyz.append(view_points)
        x_min = min(x_min, float(view_points[:, 0].min()))
        x_max = max(x_max, float(view_points[:, 0].max()))
        y_min = min(y_min, float(view_points[:, 1].min()))
        y_max = max(y_max, float(view_points[:, 1].max()))

    available_width = max(1.0, width - (2.0 * BASE_MARGIN_PX))
    available_height = max(1.0, height - (2.0 * BASE_MARGIN_PX))
    span_x = max(x_max - x_min, 1e-6)
    span_y = max(y_max - y_min, 1e-6)
    scale = min(available_width / span_x, available_height / span_y)
    geometry_width = span_x * scale
    geometry_height = span_y * scale
    offset_x = BASE_MARGIN_PX + ((available_width - geometry_width) * 0.5)
    geometry_bottom = height - (BASE_MARGIN_PX + ((available_height - geometry_height) * 0.5))

    projected_instances: list[ProjectedMeshInstance] = []
    component_count = len(mesh_instances)
    for index, (instance, view_points) in enumerate(zip(mesh_instances, projected_xyz, strict=True)):
        screen_x = offset_x + ((view_points[:, 0] - x_min) * scale)
        screen_y = geometry_bottom - ((view_points[:, 1] - y_min) * scale)
        face_normals, face_brightness = _face_shading(view_points, instance.triangles)
        projected_instances.append(
            ProjectedMeshInstance(
                screen_points=np.column_stack((screen_x, screen_y)),
                view_points=view_points,
                triangles=instance.triangles,
                face_brightness=face_brightness,
                face_normals=face_normals,
                feature_edges=_feature_edges(instance.triangles, face_normals) if include_edges else (),
                color_rgb=_component_color(
                    index,
                    default_color=model_color,
                    count=component_count,
                    explicit_color=instance.color_rgb,
                ),
            )
        )
    return projected_instances, right, true_up


def _camera_basis(
    direction: tuple[float, float, float],
    up: tuple[float, float, float],
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    view_normal = _normalize(np.asarray(direction, dtype=np.float64))
    up_vector = np.asarray(up, dtype=np.float64)
    projected_up = up_vector - (view_normal * float(np.dot(up_vector, view_normal)))
    if np.linalg.norm(projected_up) <= 1e-9:
        fallback_up = np.asarray((0.0, 1.0, 0.0) if abs(view_normal[1]) < 0.9 else (0.0, 0.0, 1.0))
        projected_up = fallback_up - (view_normal * float(np.dot(fallback_up, view_normal)))
    true_up = _normalize(projected_up)
    right = _normalize(np.cross(true_up, view_normal))
    return right, true_up, view_normal


def _normalize(vector: np.ndarray) -> np.ndarray:
    length = float(np.linalg.norm(vector))
    if length <= 1e-9:
        raise ValueError(f"Cannot normalize near-zero vector: {vector}")
    return vector / length


def _component_color(
    index: int,
    *,
    default_color: tuple[float, float, float],
    count: int,
    explicit_color: tuple[float, float, float] | None,
) -> tuple[float, float, float]:
    if explicit_color is not None:
        return explicit_color
    if count <= 1:
        return default_color
    return FALLBACK_COMPONENT_COLORS[index % len(FALLBACK_COMPONENT_COLORS)]


def _face_shading(view_points: np.ndarray, triangles: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    tri_points = view_points[triangles]
    normals = np.cross(tri_points[:, 1] - tri_points[:, 0], tri_points[:, 2] - tri_points[:, 0])
    lengths = np.linalg.norm(normals, axis=1)
    safe_lengths = np.where(lengths > 1e-9, lengths, 1.0)
    normals = normals / safe_lengths[:, None]

    oriented = normals.copy()
    oriented[oriented[:, 2] < 0.0] *= -1.0
    light_dir = _normalize(np.asarray((-0.30, 0.45, 1.0), dtype=np.float64))
    brightness = 0.38 + (0.58 * np.clip(oriented @ light_dir, 0.0, 1.0)) + (0.08 * np.clip(oriented[:, 2], 0.0, 1.0))
    return normals, np.clip(brightness, 0.25, 1.0)


def _feature_edges(triangles: np.ndarray, face_normals: np.ndarray) -> tuple[tuple[int, int], ...]:
    edge_faces: dict[tuple[int, int], list[int]] = {}
    for face_index, triangle in enumerate(triangles):
        a, b, c = (int(triangle[0]), int(triangle[1]), int(triangle[2]))
        for start, end in ((a, b), (b, c), (c, a)):
            edge = (start, end) if start < end else (end, start)
            edge_faces.setdefault(edge, []).append(face_index)

    cosine_threshold = math.cos(math.radians(FEATURE_EDGE_ANGLE_DEG))
    feature_edges: list[tuple[int, int]] = []
    for edge, incident_faces in edge_faces.items():
        if len(incident_faces) == 1:
            feature_edges.append(edge)
            continue
        if len(incident_faces) > 2:
            feature_edges.append(edge)
            continue
        first, second = incident_faces
        normal_a = face_normals[first]
        normal_b = face_normals[second]
        if float(np.dot(normal_a, normal_b)) <= cosine_threshold:
            feature_edges.append(edge)
            continue
        if (normal_a[2] >= 0.0) != (normal_b[2] >= 0.0):
            feature_edges.append(edge)
    feature_edges.sort()
    return tuple(feature_edges)


def _render_scene(
    mesh_instances: list[ProjectedMeshInstance],
    *,
    width: int,
    height: int,
    background_color: tuple[float, float, float],
    edges: bool,
    axes: bool,
    right: np.ndarray,
    true_up: np.ndarray,
) -> np.ndarray:
    background_rgb = np.asarray(_rgb_u8(background_color), dtype=np.uint8)
    image = np.empty((height, width, 3), dtype=np.uint8)
    image[:, :, :] = background_rgb
    depth_buffer = np.full((height, width), -np.inf, dtype=np.float32)

    for instance in mesh_instances:
        _rasterize_faces(image, depth_buffer, instance)
    if edges:
        for instance in mesh_instances:
            _rasterize_feature_edges(image, depth_buffer, instance)
    cropped = _crop_to_content(image, background_rgb)
    if axes:
        _draw_axes_overlay(cropped, right=right, true_up=true_up, background_rgb=background_rgb)
    return cropped


def _rasterize_faces(image: np.ndarray, depth_buffer: np.ndarray, instance: ProjectedMeshInstance) -> None:
    base_color = np.asarray(_rgb_u8(instance.color_rgb), dtype=np.float32)
    triangle_count = instance.triangles.shape[0]
    for face_index in range(triangle_count):
        triangle = instance.triangles[face_index]
        indices = (int(triangle[0]), int(triangle[1]), int(triangle[2]))
        screen_triangle = instance.screen_points[list(indices)]
        view_triangle = instance.view_points[list(indices)]
        _rasterize_triangle(
            image=image,
            depth_buffer=depth_buffer,
            screen_triangle=screen_triangle,
            view_triangle=view_triangle,
            rgb=np.asarray(np.clip(np.rint(base_color * instance.face_brightness[face_index]), 0, 255), dtype=np.uint8),
        )


def _rasterize_triangle(
    *,
    image: np.ndarray,
    depth_buffer: np.ndarray,
    screen_triangle: np.ndarray,
    view_triangle: np.ndarray,
    rgb: np.ndarray,
) -> None:
    x_coords = screen_triangle[:, 0]
    y_coords = screen_triangle[:, 1]
    min_x = max(int(math.floor(float(x_coords.min()))), 0)
    max_x = min(int(math.ceil(float(x_coords.max()))), image.shape[1] - 1)
    min_y = max(int(math.floor(float(y_coords.min()))), 0)
    max_y = min(int(math.ceil(float(y_coords.max()))), image.shape[0] - 1)
    if min_x > max_x or min_y > max_y:
        return

    p0, p1, p2 = screen_triangle
    denominator = ((p1[1] - p2[1]) * (p0[0] - p2[0])) + ((p2[0] - p1[0]) * (p0[1] - p2[1]))
    if abs(float(denominator)) <= 1e-9:
        return

    x_range = np.arange(min_x, max_x + 1, dtype=np.float32) + 0.5
    y_range = np.arange(min_y, max_y + 1, dtype=np.float32) + 0.5
    grid_x, grid_y = np.meshgrid(x_range, y_range)

    w0 = (((p1[1] - p2[1]) * (grid_x - p2[0])) + ((p2[0] - p1[0]) * (grid_y - p2[1]))) / denominator
    w1 = (((p2[1] - p0[1]) * (grid_x - p2[0])) + ((p0[0] - p2[0]) * (grid_y - p2[1]))) / denominator
    w2 = 1.0 - w0 - w1

    epsilon = 1e-5
    inside = (w0 >= -epsilon) & (w1 >= -epsilon) & (w2 >= -epsilon)
    if not np.any(inside):
        return

    z0, z1, z2 = (float(view_triangle[0, 2]), float(view_triangle[1, 2]), float(view_triangle[2, 2]))
    interpolated_depth = (w0 * z0) + (w1 * z1) + (w2 * z2)
    depth_patch = depth_buffer[min_y : max_y + 1, min_x : max_x + 1]
    update_mask = inside & (interpolated_depth >= (depth_patch - 1e-4))
    if not np.any(update_mask):
        return

    image_patch = image[min_y : max_y + 1, min_x : max_x + 1]
    image_patch[update_mask] = rgb
    depth_patch[update_mask] = interpolated_depth[update_mask]


def _rasterize_feature_edges(
    image: np.ndarray,
    depth_buffer: np.ndarray,
    instance: ProjectedMeshInstance,
) -> None:
    edge_color = np.asarray(_edge_rgb(instance.color_rgb), dtype=np.uint8)
    for start, end in instance.feature_edges:
        screen_start = instance.screen_points[start]
        screen_end = instance.screen_points[end]
        depth_start = float(instance.view_points[start, 2])
        depth_end = float(instance.view_points[end, 2])
        _draw_depth_tested_segment(
            image=image,
            depth_buffer=depth_buffer,
            start=screen_start,
            end=screen_end,
            depth_start=depth_start,
            depth_end=depth_end,
            color=edge_color,
            radius_px=1,
        )


def _draw_depth_tested_segment(
    *,
    image: np.ndarray,
    depth_buffer: np.ndarray,
    start: np.ndarray,
    end: np.ndarray,
    depth_start: float,
    depth_end: float,
    color: np.ndarray,
    radius_px: int,
) -> None:
    x0 = float(start[0])
    y0 = float(start[1])
    x1 = float(end[0])
    y1 = float(end[1])
    dx = x1 - x0
    dy = y1 - y0
    steps = max(1, int(math.ceil(max(abs(dx), abs(dy)))))
    offsets = _brush_offsets(radius_px)
    for step in range(steps + 1):
        t = step / steps
        xi = int(round(x0 + (dx * t)))
        yi = int(round(y0 + (dy * t)))
        depth = depth_start + ((depth_end - depth_start) * t) + 5e-4
        for ox, oy in offsets:
            px = xi + ox
            py = yi + oy
            if 0 <= px < image.shape[1] and 0 <= py < image.shape[0]:
                if depth >= (float(depth_buffer[py, px]) - 1e-3):
                    image[py, px] = color


def _brush_offsets(radius_px: int) -> tuple[tuple[int, int], ...]:
    offsets: list[tuple[int, int]] = []
    radius_squared = radius_px * radius_px
    for dy in range(-radius_px, radius_px + 1):
        for dx in range(-radius_px, radius_px + 1):
            if (dx * dx) + (dy * dy) <= radius_squared:
                offsets.append((dx, dy))
    return tuple(offsets)


def _draw_axes_overlay(
    image: np.ndarray,
    *,
    right: np.ndarray,
    true_up: np.ndarray,
    background_rgb: np.ndarray,
) -> None:
    box_size = min(AXIS_BOX_SIZE_PX, max(24, min(image.shape[0], image.shape[1]) // 4))
    placement = _best_axis_corner(image, background_rgb, box_size=box_size)
    if placement is None:
        return
    x0, y0 = placement
    x1 = min(x0 + box_size, image.shape[1])
    y1 = min(y0 + box_size, image.shape[0])
    image[y0:y1, x0:x1] = background_rgb
    origin = np.asarray((x0 + 10.0, y1 - 10.0), dtype=np.float64)
    axis_length = float(min(x1 - x0, y1 - y0) - 18)
    if axis_length <= 6.0:
        return
    basis_by_axis = (
        (np.asarray((1.0, 0.0, 0.0)), np.asarray((214, 71, 71), dtype=np.uint8)),
        (np.asarray((0.0, 1.0, 0.0)), np.asarray((62, 165, 83), dtype=np.uint8)),
        (np.asarray((0.0, 0.0, 1.0)), np.asarray((65, 111, 219), dtype=np.uint8)),
    )
    for world_axis, color in basis_by_axis:
        projected = np.asarray((float(np.dot(world_axis, right)), -float(np.dot(world_axis, true_up))))
        magnitude = float(np.linalg.norm(projected))
        if magnitude <= 1e-9:
            continue
        end = origin + ((projected / magnitude) * axis_length)
        _draw_flat_segment(
            image=image,
            start=origin,
            end=end,
            color=color,
            radius_px=1,
        )


def _best_axis_corner(image: np.ndarray, background_rgb: np.ndarray, *, box_size: int) -> tuple[int, int] | None:
    corners = (
        (0, image.shape[0] - box_size),
        (image.shape[1] - box_size, image.shape[0] - box_size),
        (0, 0),
        (image.shape[1] - box_size, 0),
    )
    best_corner: tuple[int, int] | None = None
    best_score = -1
    for x0, y0 in corners:
        x0 = max(0, x0)
        y0 = max(0, y0)
        x1 = min(x0 + box_size, image.shape[1])
        y1 = min(y0 + box_size, image.shape[0])
        if x0 >= x1 or y0 >= y1:
            continue
        patch = image[y0:y1, x0:x1]
        background_count = int(np.count_nonzero(np.all(patch == background_rgb, axis=2)))
        if background_count > best_score:
            best_score = background_count
            best_corner = (x0, y0)
    return best_corner


def _draw_flat_segment(
    *,
    image: np.ndarray,
    start: np.ndarray,
    end: np.ndarray,
    color: np.ndarray,
    radius_px: int,
) -> None:
    x0 = float(start[0])
    y0 = float(start[1])
    x1 = float(end[0])
    y1 = float(end[1])
    dx = x1 - x0
    dy = y1 - y0
    steps = max(1, int(math.ceil(max(abs(dx), abs(dy)))))
    offsets = _brush_offsets(radius_px)
    for step in range(steps + 1):
        t = step / steps
        xi = int(round(x0 + (dx * t)))
        yi = int(round(y0 + (dy * t)))
        for ox, oy in offsets:
            px = xi + ox
            py = yi + oy
            if 0 <= px < image.shape[1] and 0 <= py < image.shape[0]:
                image[py, px] = color


def _crop_to_content(image: np.ndarray, background_rgb: np.ndarray) -> np.ndarray:
    content_mask = np.any(image != background_rgb, axis=2)
    if not np.any(content_mask):
        return image
    ys, xs = np.where(content_mask)
    min_x = max(int(xs.min()) - CROP_PADDING_PX, 0)
    max_x = min(int(xs.max()) + CROP_PADDING_PX + 1, image.shape[1])
    min_y = max(int(ys.min()) - CROP_PADDING_PX, 0)
    max_y = min(int(ys.max()) + CROP_PADDING_PX + 1, image.shape[0])
    return image[min_y:max_y, min_x:max_x].copy()


def _edge_rgb(color_rgb: tuple[float, float, float]) -> tuple[int, int, int]:
    base = np.asarray(_rgb_u8(color_rgb), dtype=np.float32)
    darkened = np.clip(np.rint((base * 0.38) - 6.0), 0, 255).astype(np.uint8)
    return (int(darkened[0]), int(darkened[1]), int(darkened[2]))


def _rgb_u8(rgb: tuple[float, float, float]) -> tuple[int, int, int]:
    return tuple(int(round(max(0.0, min(1.0, float(channel))) * 255.0)) for channel in rgb)


def _write_png(image: np.ndarray, png_path: Path) -> None:
    if image.ndim != 3 or image.shape[2] != 3:
        raise ValueError("PNG image must be an HxWx3 RGB array")
    png_path.parent.mkdir(parents=True, exist_ok=True)

    height, width, _channels = image.shape
    scanlines = bytearray()
    for row in image:
        scanlines.append(0)
        scanlines.extend(row.astype(np.uint8, copy=False).tobytes())

    def chunk(tag: bytes, data: bytes) -> bytes:
        payload = tag + data
        return (
            len(data).to_bytes(4, "big")
            + payload
            + zlib.crc32(payload).to_bytes(4, "big")
        )

    png_bytes = bytearray(b"\x89PNG\r\n\x1a\n")
    png_bytes.extend(
        chunk(
            b"IHDR",
            width.to_bytes(4, "big")
            + height.to_bytes(4, "big")
            + bytes((8, 2, 0, 0, 0)),
        )
    )
    png_bytes.extend(chunk(b"IDAT", zlib.compress(bytes(scanlines), level=9)))
    png_bytes.extend(chunk(b"IEND", b""))
    png_path.write_bytes(png_bytes)