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__)

MIN_BLENDER_VERSION = (5, 0, 1)


def _glb_from_step(step_path: Path, glb_path: Path, quality: str = "low") -> None:
    """Convert STEP → GLB via OCC (export_step_to_gltf.py, no Blender needed).

    quality: "low" → coarser mesh (~0.3 mm deflection, fast)
             "high" → finer mesh (~0.05 mm deflection, slower)
    """
    import subprocess
    import sys as _sys

    linear_deflection  = 0.3  if quality == "low" else 0.05
    angular_deflection = 0.5  if quality == "low" else 0.2

    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,
    stl_quality: str = "low",
    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,
) -> dict:
    """Convert STEP → GLB (OCC) → PNG (Blender subprocess).

    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 — replaces cadquery STL)
    glb_path = step_path.parent / f"{step_path.stem}_thumbnail.glb"

    t_glb = time.monotonic()
    if not glb_path.exists() or glb_path.stat().st_size == 0:
        _glb_from_step(step_path, glb_path, quality=stl_quality)
    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:
        cmd = [
            blender_bin,
            "--background",
            "--python", str(script_path),
            "--",
            str(glb_path),
            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 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]

    # EEVEE fallback to Cycles on non-signal error
    if returncode > 0 and engine == "eevee":
        logger.warning("EEVEE failed (exit %d) — retrying with Cycles", returncode)
        returncode, stdout_lines2, stderr_lines2 = _run("cycles")
        engine_used = "cycles (eevee fallback)"
        log_lines.extend(l for l in stdout_lines2 if "[blender_render]" in l)

    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,
        "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,
    stl_quality: str = "low",
    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,
) -> dict:
    """Render a turntable animation: STEP → STL → N frames (Blender) → mp4 (ffmpeg).

    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 — replaces cadquery STL)
    glb_path = step_path.parent / f"{step_path.stem}_thumbnail.glb"

    t_glb = time.monotonic()
    if not glb_path.exists() or glb_path.stat().st_size == 0:
        _glb_from_step(step_path, glb_path, quality=stl_quality)
    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"

    cmd = [
        blender_bin,
        "--background",
        "--python", str(script_path),
        "--",
        str(glb_path),
        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",
    ]

    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,
    }