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 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 convert_step_to_stl(step_path: Path, stl_path: Path, quality: str = "low") -> None:
    """Convert a STEP file to STL using cadquery.

    Raises ImportError if cadquery is not installed (not available in backend
    container — only in render-worker container).
    """
    import cadquery as cq  # only available in render-worker

    if quality == "high":
        shape = cq.importers.importStep(str(step_path))
        cq.exporters.export(shape, str(stl_path), tolerance=0.01, angularTolerance=0.02)
    else:
        shape = cq.importers.importStep(str(step_path))
        cq.exporters.export(shape, str(stl_path), tolerance=0.3, angularTolerance=0.3)

    if not stl_path.exists() or stl_path.stat().st_size == 0:
        raise RuntimeError("cadquery produced empty STL")


def export_per_part_stls(step_path: Path, parts_dir: Path, quality: str = "low") -> list:
    """Export one STL per named STEP leaf shape using OCP XCAF.

    Returns the manifest list (may be empty on failure — non-fatal).
    """
    tol = 0.01 if quality == "high" else 0.3
    angular_tol = 0.05 if quality == "high" else 0.3

    try:
        from OCP.STEPCAFControl import STEPCAFControl_Reader
        from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ShapeTool
        from OCP.TDataStd import TDataStd_Name
        from OCP.TDF import TDF_Label as TDF_Label_cls, TDF_LabelSequence
        from OCP.XCAFApp import XCAFApp_Application
        from OCP.TDocStd import TDocStd_Document
        from OCP.TCollection import TCollection_ExtendedString
        from OCP.IFSelect import IFSelect_RetDone
        import cadquery as cq
    except ImportError as e:
        logger.warning("per-part export skipped (import error): %s", e)
        return []

    app = XCAFApp_Application.GetApplication_s()
    doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf"))
    app.InitDocument(doc)

    reader = STEPCAFControl_Reader()
    reader.SetNameMode(True)
    status = reader.ReadFile(str(step_path))
    if status != IFSelect_RetDone:
        logger.warning("XCAF reader failed with status %s", status)
        return []

    if not reader.Transfer(doc):
        logger.warning("XCAF transfer failed")
        return []

    shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
    name_id = TDataStd_Name.GetID_s()

    leaves = []

    def _get_label_name(label):
        name_attr = TDataStd_Name()
        if label.FindAttribute(name_id, name_attr):
            return name_attr.Get().ToExtString()
        return ""

    def _collect_leaves(label):
        if XCAFDoc_ShapeTool.IsAssembly_s(label):
            components = TDF_LabelSequence()
            XCAFDoc_ShapeTool.GetComponents_s(label, components)
            for i in range(1, components.Length() + 1):
                comp_label = components.Value(i)
                if XCAFDoc_ShapeTool.IsReference_s(comp_label):
                    ref_label = TDF_Label_cls()
                    XCAFDoc_ShapeTool.GetReferredShape_s(comp_label, ref_label)
                    comp_name = _get_label_name(comp_label)
                    ref_name = _get_label_name(ref_label)
                    name = ref_name or comp_name
                    if XCAFDoc_ShapeTool.IsAssembly_s(ref_label):
                        _collect_leaves(ref_label)
                    elif XCAFDoc_ShapeTool.IsSimpleShape_s(ref_label):
                        shape = XCAFDoc_ShapeTool.GetShape_s(comp_label)
                        leaves.append((name or f"unnamed_{len(leaves)}", shape))
                else:
                    _collect_leaves(comp_label)
        elif XCAFDoc_ShapeTool.IsSimpleShape_s(label):
            name = _get_label_name(label)
            shape = XCAFDoc_ShapeTool.GetShape_s(label)
            leaves.append((name or f"unnamed_{len(leaves)}", shape))

    top_labels = TDF_LabelSequence()
    shape_tool.GetFreeShapes(top_labels)
    for i in range(1, top_labels.Length() + 1):
        _collect_leaves(top_labels.Value(i))

    if not leaves:
        logger.warning("no leaf shapes found via XCAF")
        return []

    parts_dir.mkdir(parents=True, exist_ok=True)
    manifest = []

    for idx, (name, shape) in enumerate(leaves):
        safe_name = name.replace("/", "_").replace("\\", "_").replace(" ", "_")
        filename = f"{idx:02d}_{safe_name}.stl"
        filepath = str(parts_dir / filename)
        try:
            cq_shape = cq.Shape(shape)
            cq_shape.exportStl(filepath, tolerance=tol, angularTolerance=angular_tol)
            manifest.append({"index": idx, "name": name, "file": filename})
        except Exception as e:
            logger.warning("failed to export part '%s': %s", name, e)

    manifest_path = parts_dir / "manifest.json"
    with open(manifest_path, "w") as f:
        json.dump({"parts": manifest}, f, indent=2)

    return manifest


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,
) -> dict:
    """Convert STEP → STL (cadquery) → 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():
        # Fallback: look next to this file (development mode)
        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. STL conversion (cadquery)
    stl_path = step_path.parent / f"{step_path.stem}_{stl_quality}.stl"
    parts_dir = step_path.parent / f"{step_path.stem}_{stl_quality}_parts"

    t_stl = time.monotonic()
    if not stl_path.exists() or stl_path.stat().st_size == 0:
        logger.info("STL cache miss — converting: %s", step_path.name)
        convert_step_to_stl(step_path, stl_path, stl_quality)
    else:
        logger.info("STL cache hit: %s (%d KB)", stl_path.name, stl_path.stat().st_size // 1024)
    stl_size_bytes = stl_path.stat().st_size if stl_path.exists() else 0

    if not (parts_dir / "manifest.json").exists():
        try:
            export_per_part_stls(step_path, parts_dir, stl_quality)
        except Exception as exc:
            logger.warning("per-part STL export failed (non-fatal): %s", exc)

    stl_duration_s = round(time.monotonic() - t_stl, 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(stl_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) -> subprocess.CompletedProcess:
        proc = subprocess.Popen(
            _build_cmd(eng),
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            text=True, env=env, start_new_session=True,
        )
        try:
            stdout, stderr = proc.communicate(timeout=600)
        except subprocess.TimeoutExpired:
            try:
                os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
            except (ProcessLookupError, OSError):
                pass
            stdout, stderr = proc.communicate()
        return subprocess.CompletedProcess(_build_cmd(eng), proc.returncode, stdout, stderr)

    t_render = time.monotonic()
    result = _run(engine)
    engine_used = engine

    log_lines = []
    for line in (result.stdout or "").splitlines():
        logger.info("[blender] %s", line)
        if "[blender_render]" in line:
            log_lines.append(line)
    for line in (result.stderr or "").splitlines():
        logger.warning("[blender stderr] %s", line)

    # EEVEE fallback to Cycles on non-signal error
    if result.returncode > 0 and engine == "eevee":
        logger.warning("EEVEE failed (exit %d) — retrying with Cycles", result.returncode)
        result = _run("cycles")
        engine_used = "cycles (eevee fallback)"
        for line in (result.stdout or "").splitlines():
            logger.info("[blender-fallback] %s", line)
            if "[blender_render]" in line:
                log_lines.append(line)

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

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

    parts_count = 0
    manifest_file = parts_dir / "manifest.json"
    if manifest_file.exists():
        try:
            data = json.loads(manifest_file.read_text())
            parts_count = len(data.get("parts", []))
        except Exception:
            pass

    return {
        "total_duration_s": round(time.monotonic() - t0, 2),
        "stl_duration_s": stl_duration_s,
        "render_duration_s": render_duration_s,
        "stl_size_bytes": stl_size_bytes,
        "output_size_bytes": output_path.stat().st_size if output_path.exists() else 0,
        "parts_count": parts_count,
        "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. STL conversion
    stl_path = step_path.parent / f"{step_path.stem}_{stl_quality}.stl"
    parts_dir = step_path.parent / f"{step_path.stem}_{stl_quality}_parts"

    t_stl = time.monotonic()
    if not stl_path.exists() or stl_path.stat().st_size == 0:
        logger.info("STL cache miss — converting: %s", step_path.name)
        convert_step_to_stl(step_path, stl_path, stl_quality)
    else:
        logger.info("STL cache hit: %s (%d KB)", stl_path.name, stl_path.stat().st_size // 1024)
    stl_size_bytes = stl_path.stat().st_size if stl_path.exists() else 0

    if not (parts_dir / "manifest.json").exists():
        try:
            export_per_part_stls(step_path, parts_dir, stl_quality)
        except Exception as exc:
            logger.warning("per-part STL export failed (non-fatal): %s", exc)

    stl_duration_s = round(time.monotonic() - t_stl, 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(stl_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": stl_duration_s,
        "render_duration_s": render_duration_s,
        "ffmpeg_duration_s": ffmpeg_duration_s,
        "stl_size_bytes": stl_size_bytes,
        "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,
    }