HartOMat/backend/app/services/render_blender.py

"""Direct Blender rendering service — runs Blender as a subprocess.

Used by the render-worker Celery container (which has BLENDER_BIN set and
cadquery installed). The backend and standard workers fall back to the Pillow
placeholder when this service is unavailable.
"""
import json
import logging
import os
import shutil
import signal
import subprocess
from pathlib import Path

logger = logging.getLogger(__name__)


def _glb_from_step(step_path: Path, glb_path: Path) -> None:
    """Convert STEP → GLB via OCC (export_step_to_gltf.py, no Blender needed)."""
    import subprocess
    import sys as _sys

    linear_deflection  = 0.3
    angular_deflection = 0.5

    scripts_dir = Path(os.environ.get("RENDER_SCRIPTS_DIR", "/render-scripts"))
    script_path = scripts_dir / "export_step_to_gltf.py"

    cmd = [
        _sys.executable, str(script_path),
        "--step_path", str(step_path),
        "--output_path", str(glb_path),
        "--linear_deflection", str(linear_deflection),
        "--angular_deflection", str(angular_deflection),
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    for line in result.stdout.splitlines():
        logger.info("[occ-gltf] %s", line)
    for line in result.stderr.splitlines():
        logger.warning("[occ-gltf stderr] %s", line)
    if result.returncode != 0 or not glb_path.exists() or glb_path.stat().st_size == 0:
        raise RuntimeError(
            f"export_step_to_gltf.py failed (exit {result.returncode}).\n"
            f"STDERR: {result.stderr[-1000:]}"
        )
    logger.info("GLB converted: %s (%d KB)", glb_path.name, glb_path.stat().st_size // 1024)


def find_blender() -> str:
    """Locate the Blender binary via $BLENDER_BIN or PATH."""
    env_bin = os.environ.get("BLENDER_BIN", "")
    if env_bin and Path(env_bin).exists():
        return env_bin
    found = shutil.which("blender")
    return found or ""


def is_blender_available() -> bool:
    """Return True if a Blender binary is reachable from this process."""
    return bool(find_blender())


def render_still(
    step_path: Path,
    output_path: Path,
    width: int = 512,
    height: int = 512,
    engine: str = "cycles",
    samples: int = 256,
    smooth_angle: int = 30,
    cycles_device: str = "auto",
    transparent_bg: bool = False,
    part_colors: dict | None = None,
    template_path: str | None = None,
    target_collection: str = "Product",
    material_library_path: str | None = None,
    material_map: dict | None = None,
    part_names_ordered: list | None = None,
    lighting_only: bool = False,
    shadow_catcher: bool = False,
    rotation_x: float = 0.0,
    rotation_y: float = 0.0,
    rotation_z: float = 0.0,
    noise_threshold: str = "",
    denoiser: str = "",
    denoising_input_passes: str = "",
    denoising_prefilter: str = "",
    denoising_quality: str = "",
    denoising_use_gpu: str = "",
    mesh_attributes: dict | None = None,
    log_callback: "Callable[[str], None] | None" = None,
    usd_path: "Path | None" = None,
) -> dict:
    """Convert STEP → GLB (OCC) → PNG (Blender subprocess).

    When usd_path is provided and the file exists, the GLB conversion step is
    skipped and Blender imports the USD stage directly (--usd-path flag).

    Returns a dict with timing, sizes, engine_used, and log_lines.
    Raises RuntimeError on failure.
    """
    import time

    blender_bin = find_blender()
    if not blender_bin:
        raise RuntimeError("Blender binary not found — check BLENDER_BIN env or PATH")

    script_path = Path(os.environ.get("RENDER_SCRIPTS_DIR", "/render-scripts")) / "blender_render.py"
    if not script_path.exists():
        alt = Path(__file__).parent.parent.parent.parent / "render-worker" / "scripts" / "blender_render.py"
        if alt.exists():
            script_path = alt
        else:
            raise RuntimeError(f"blender_render.py not found at {script_path}")

    t0 = time.monotonic()

    # 1. GLB conversion (OCC) — skipped when usd_path is provided
    glb_path = step_path.parent / f"{step_path.stem}_thumbnail.glb"
    use_usd = bool(usd_path and usd_path.exists())

    t_glb = time.monotonic()
    if use_usd:
        logger.info("[render_blender] using USD path: %s", usd_path)
        glb_size_bytes = 0
    else:
        if not glb_path.exists() or glb_path.stat().st_size == 0:
            _glb_from_step(step_path, glb_path)
        else:
            logger.info("GLB local hit: %s (%d KB)", glb_path.name, glb_path.stat().st_size // 1024)
        glb_size_bytes = glb_path.stat().st_size if glb_path.exists() else 0
    glb_duration_s = round(time.monotonic() - t_glb, 2)

    # 2. Blender render
    output_path.parent.mkdir(parents=True, exist_ok=True)

    env = dict(os.environ)
    if engine == "eevee":
        env.update({
            "VK_ICD_FILENAMES":         "/usr/share/vulkan/icd.d/lvp_icd.x86_64.json",
            "LIBGL_ALWAYS_SOFTWARE":    "1",
            "MESA_GL_VERSION_OVERRIDE": "4.5",
            "EGL_PLATFORM":             "surfaceless",
        })
    else:
        env["EGL_PLATFORM"] = "surfaceless"

    def _build_cmd(eng: str) -> list:
        # Pass "" as glb_path when using USD — blender_render.py reads --usd-path instead
        glb_arg = "" if use_usd else str(glb_path)
        cmd = [
            blender_bin,
            "--background",
            "--python", str(script_path),
            "--",
            glb_arg,
            str(output_path),
            str(width), str(height),
            eng, str(samples), str(smooth_angle),
            cycles_device,
            "1" if transparent_bg else "0",
            template_path or "",
            target_collection,
            material_library_path or "",
            json.dumps(material_map) if material_map else "{}",
            json.dumps(part_names_ordered) if part_names_ordered else "[]",
            "1" if lighting_only else "0",
            "1" if shadow_catcher else "0",
            str(rotation_x), str(rotation_y), str(rotation_z),
            noise_threshold or "", denoiser or "",
            denoising_input_passes or "", denoising_prefilter or "",
            denoising_quality or "", denoising_use_gpu or "",
        ]
        if use_usd:
            cmd += ["--usd-path", str(usd_path)]
            if mesh_attributes:
                logger.debug("[render_blender] usd_path active — mesh_attributes ignored")
        elif mesh_attributes:
            cmd += ["--mesh-attributes", json.dumps(mesh_attributes)]
        return cmd

    def _run(eng: str) -> tuple[int, list[str], list[str]]:
        """Run Blender subprocess, streaming stdout line-by-line.

        Returns (returncode, stdout_lines, stderr_lines).
        """
        import selectors
        proc = subprocess.Popen(
            _build_cmd(eng),
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            text=True, env=env, start_new_session=True,
        )
        stdout_lines: list[str] = []
        stderr_lines: list[str] = []
        deadline = time.monotonic() + 600

        sel = selectors.DefaultSelector()
        sel.register(proc.stdout, selectors.EVENT_READ, "stdout")
        sel.register(proc.stderr, selectors.EVENT_READ, "stderr")

        try:
            while sel.get_map():
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    try:
                        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
                    except (ProcessLookupError, OSError):
                        pass
                    break
                events = sel.select(timeout=min(remaining, 2.0))
                for key, _ in events:
                    line = key.fileobj.readline()
                    if not line:
                        sel.unregister(key.fileobj)
                        continue
                    line = line.rstrip("\n")
                    if key.data == "stdout":
                        stdout_lines.append(line)
                        logger.info("[blender] %s", line)
                        if log_callback and "[blender_render]" in line:
                            log_callback(line)
                    else:
                        stderr_lines.append(line)
                        logger.warning("[blender stderr] %s", line)
        finally:
            sel.close()

        proc.wait(timeout=10)
        return proc.returncode, stdout_lines, stderr_lines

    t_render = time.monotonic()
    returncode, stdout_lines, stderr_lines = _run(engine)
    engine_used = engine

    log_lines = [l for l in stdout_lines if "[blender_render]" in l]

    # Parse RENDER_DEVICE_USED token from stdout
    device_used = "unknown"
    compute_type = "unknown"
    gpu_fallback = False
    for line in stdout_lines:
        if line.startswith("RENDER_DEVICE_USED:"):
            parts = line.split()
            for part in parts:
                if part.startswith("device="):
                    device_used = part.split("=", 1)[1]
                elif part.startswith("compute_type="):
                    compute_type = part.split("=", 1)[1]
            gpu_fallback = (device_used == "CPU")
            break

    # EEVEE fallback removed (Phase 5.2): EEVEE Next in Blender 5.0+ is stable.
    # If EEVEE fails, it is a hard failure — no silent retry.

    if returncode == 2:
        raise RuntimeError(
            "GPU required but render used CPU — strict mode (CYCLES_DEVICE=gpu). "
            "Check that the render-worker has a visible NVIDIA GPU."
        )

    if returncode != 0:
        stdout_tail = "\n".join(stdout_lines[-50:]) if stdout_lines else ""
        stderr_tail = "\n".join(stderr_lines[-20:]) if stderr_lines else ""
        raise RuntimeError(
            f"Blender exited with code {returncode}.\n"
            f"stdout: {stdout_tail[-2000:]}\n"
            f"stderr: {stderr_tail[-500:]}"
        )

    render_duration_s = round(time.monotonic() - t_render, 2)

    return {
        "total_duration_s": round(time.monotonic() - t0, 2),
        "stl_duration_s": glb_duration_s,   # key kept for backward compat with DB render_log
        "render_duration_s": render_duration_s,
        "stl_size_bytes": glb_size_bytes,
        "output_size_bytes": output_path.stat().st_size if output_path.exists() else 0,
        "parts_count": 0,
        "engine_used": engine_used,
        "device_used": device_used,
        "compute_type": compute_type,
        "gpu_fallback": gpu_fallback,
        "log_lines": log_lines,
    }


def render_turntable_to_file(
    step_path: Path,
    output_path: Path,
    frame_count: int = 24,
    fps: int = 25,
    width: int = 1920,
    height: int = 1920,
    engine: str = "cycles",
    samples: int = 128,
    smooth_angle: int = 30,
    cycles_device: str = "auto",
    transparent_bg: bool = False,
    bg_color: str = "",
    turntable_axis: str = "world_z",
    part_colors: dict | None = None,
    template_path: str | None = None,
    target_collection: str = "Product",
    material_library_path: str | None = None,
    material_map: dict | None = None,
    part_names_ordered: list | None = None,
    lighting_only: bool = False,
    shadow_catcher: bool = False,
    rotation_x: float = 0.0,
    rotation_y: float = 0.0,
    rotation_z: float = 0.0,
    usd_path: "Path | None" = None,
) -> dict:
    """Render a turntable animation: STEP → STL → N frames (Blender) → mp4 (ffmpeg).

    When usd_path is provided and exists, the GLB conversion step is skipped.

    Returns a dict with timing, frame count, engine_used, log_lines.
    Raises RuntimeError on failure.
    """
    import shutil as _shutil
    import tempfile
    import time

    blender_bin = find_blender()
    if not blender_bin:
        raise RuntimeError("Blender binary not found — check BLENDER_BIN env or PATH")

    script_path = Path(os.environ.get("RENDER_SCRIPTS_DIR", "/render-scripts")) / "turntable_render.py"
    if not script_path.exists():
        alt = Path(__file__).parent.parent.parent.parent / "render-worker" / "scripts" / "turntable_render.py"
        if alt.exists():
            script_path = alt
        else:
            raise RuntimeError(f"turntable_render.py not found at {script_path}")

    ffmpeg_bin = _shutil.which("ffmpeg")
    if not ffmpeg_bin:
        raise RuntimeError("ffmpeg not found — install ffmpeg in the render-worker container")

    t0 = time.monotonic()

    # 1. GLB conversion (OCC) — skipped when usd_path is provided
    glb_path = step_path.parent / f"{step_path.stem}_thumbnail.glb"
    use_usd = bool(usd_path and usd_path.exists())

    t_glb = time.monotonic()
    if use_usd:
        logger.info("[render_blender] turntable using USD path: %s", usd_path)
    else:
        if not glb_path.exists() or glb_path.stat().st_size == 0:
            _glb_from_step(step_path, glb_path)
        else:
            logger.info("GLB local hit: %s (%d KB)", glb_path.name, glb_path.stat().st_size // 1024)
    glb_duration_s = round(time.monotonic() - t_glb, 2)

    # 2. Render frames with Blender
    frames_dir = output_path.parent / f"_frames_{output_path.stem}"
    frames_dir.mkdir(parents=True, exist_ok=True)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    env = dict(os.environ)
    env["EGL_PLATFORM"] = "surfaceless"

    glb_arg = "" if use_usd else str(glb_path)
    cmd = [
        blender_bin,
        "--background",
        "--python", str(script_path),
        "--",
        glb_arg,
        str(frames_dir),
        str(frame_count),
        "360",  # degrees
        str(width), str(height),
        engine, str(samples),
        json.dumps(part_colors or {}),
        template_path or "",
        target_collection,
        material_library_path or "",
        json.dumps(material_map) if material_map else "{}",
        json.dumps(part_names_ordered) if part_names_ordered else "[]",
        "1" if lighting_only else "0",
        cycles_device,
        "1" if shadow_catcher else "0",
        str(rotation_x), str(rotation_y), str(rotation_z),
        turntable_axis,
        bg_color or "",
        "1" if transparent_bg else "0",
    ]
    if use_usd:
        cmd += ["--usd-path", str(usd_path)]

    log_lines: list[str] = []

    t_render = time.monotonic()
    proc = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        text=True, env=env, start_new_session=True,
    )
    try:
        stdout, stderr = proc.communicate(timeout=3600)  # 1hr max for full animation
    except subprocess.TimeoutExpired:
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        except (ProcessLookupError, OSError):
            pass
        stdout, stderr = proc.communicate()

    for line in (stdout or "").splitlines():
        logger.info("[turntable] %s", line)
        if "[turntable_render]" in line:
            log_lines.append(line)
    for line in (stderr or "").splitlines():
        logger.warning("[turntable stderr] %s", line)

    if proc.returncode != 0:
        raise RuntimeError(
            f"turntable_render.py exited with code {proc.returncode}.\n"
            f"stdout: {(stdout or '')[-2000:]}\n"
            f"stderr: {(stderr or '')[-500:]}"
        )

    render_duration_s = round(time.monotonic() - t_render, 2)

    # Check frames were written
    frame_files = sorted(frames_dir.glob("frame_*.png"))
    if not frame_files:
        raise RuntimeError(f"No frames rendered in {frames_dir}")
    logger.info("Rendered %d frames in %.1fs", len(frame_files), render_duration_s)

    # 3. Compose frames → mp4 with ffmpeg
    t_ffmpeg = time.monotonic()

    ffmpeg_cmd = [
        ffmpeg_bin,
        "-y",
        "-framerate", str(fps),
        "-i", str(frames_dir / "frame_%04d.png"),
        "-vcodec", "libx264",
        "-pix_fmt", "yuv420p",
        "-crf", "18",
        "-movflags", "+faststart",
        str(output_path),
    ]

    # If bg_color is set and transparent_bg is True, overlay frames on solid bg
    if bg_color and transparent_bg:
        hex_color = bg_color.lstrip("#")
        r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
        ffmpeg_cmd = [
            ffmpeg_bin, "-y",
            "-framerate", str(fps),
            "-i", str(frames_dir / "frame_%04d.png"),
            "-f", "lavfi", "-i", f"color=c=0x{hex_color}:size={width}x{height}:rate={fps}",
            "-filter_complex", "[1:v][0:v]overlay=0:0:shortest=1",
            "-vcodec", "libx264",
            "-pix_fmt", "yuv420p",
            "-crf", "18",
            "-movflags", "+faststart",
            str(output_path),
        ]

    ffmpeg_proc = subprocess.run(
        ffmpeg_cmd, capture_output=True, text=True, timeout=300
    )
    ffmpeg_duration_s = round(time.monotonic() - t_ffmpeg, 2)

    for line in (ffmpeg_proc.stdout or "").splitlines():
        logger.info("[ffmpeg] %s", line)
    for line in (ffmpeg_proc.stderr or "").splitlines():
        logger.debug("[ffmpeg stderr] %s", line)

    if ffmpeg_proc.returncode != 0:
        raise RuntimeError(
            f"ffmpeg exited with code {ffmpeg_proc.returncode}.\n"
            f"stderr: {(ffmpeg_proc.stderr or '')[-1000:]}"
        )

    # Clean up frames directory
    try:
        _shutil.rmtree(frames_dir)
    except Exception:
        pass

    return {
        "total_duration_s": round(time.monotonic() - t0, 2),
        "stl_duration_s": glb_duration_s,   # key kept for backward compat with DB render_log
        "render_duration_s": render_duration_s,
        "ffmpeg_duration_s": ffmpeg_duration_s,
        "stl_size_bytes": 0,
        "output_size_bytes": output_path.stat().st_size if output_path.exists() else 0,
        "frame_count": len(frame_files),
        "engine_used": engine,
        "log_lines": log_lines,
    }