HartOMat/backend/app/domains/rendering/tasks.py

"""Rendering domain tasks — Celery tasks for Blender-based rendering.

These tasks run on the `thumbnail_rendering` queue in the render-worker
container, which has Blender and cadquery available.

Phase A2: Initial implementation replacing the blender-renderer HTTP service.
Phase B: This module will be expanded as part of the Domain-Driven restructure.
"""
import logging
from pathlib import Path

from app.tasks.celery_app import celery_app

logger = logging.getLogger(__name__)


@celery_app.task(
    bind=True,
    name="app.domains.rendering.tasks.render_still_task",
    queue="thumbnail_rendering",
    max_retries=2,
)
def render_still_task(
    self,
    step_path: str,
    output_path: str,
    engine: str = "cycles",
    samples: int = 256,
    stl_quality: str = "low",
    smooth_angle: int = 30,
    cycles_device: str = "auto",
    width: int = 512,
    height: int = 512,
    transparent_bg: bool = False,
    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 = "",
) -> dict:
    """Render a STEP file to a still PNG via Blender subprocess.

    Returns render metadata dict on success.
    Retries up to 2 times on failure (30s countdown).
    """
    try:
        from app.services.render_blender import render_still
        result = render_still(
            step_path=Path(step_path),
            output_path=Path(output_path),
            engine=engine,
            samples=samples,
            stl_quality=stl_quality,
            smooth_angle=smooth_angle,
            cycles_device=cycles_device,
            width=width,
            height=height,
            transparent_bg=transparent_bg,
            template_path=template_path,
            target_collection=target_collection,
            material_library_path=material_library_path,
            material_map=material_map,
            part_names_ordered=part_names_ordered,
            lighting_only=lighting_only,
            shadow_catcher=shadow_catcher,
            rotation_x=rotation_x,
            rotation_y=rotation_y,
            rotation_z=rotation_z,
            noise_threshold=noise_threshold,
            denoiser=denoiser,
            denoising_input_passes=denoising_input_passes,
            denoising_prefilter=denoising_prefilter,
            denoising_quality=denoising_quality,
            denoising_use_gpu=denoising_use_gpu,
        )
        logger.info(
            "render_still_task completed: %s → %s in %.1fs",
            Path(step_path).name, Path(output_path).name,
            result.get("total_duration_s", 0),
        )
        return result
    except Exception as exc:
        logger.error("render_still_task failed for %s: %s", step_path, exc)
        raise self.retry(exc=exc, countdown=30)


@celery_app.task(
    bind=True,
    name="app.domains.rendering.tasks.render_turntable_task",
    queue="thumbnail_rendering",
    max_retries=2,
)
def render_turntable_task(
    self,
    step_path: str,
    output_dir: str,
    output_name: str = "turntable",
    engine: str = "cycles",
    samples: int = 64,
    stl_quality: str = "low",
    smooth_angle: int = 30,
    cycles_device: str = "auto",
    width: int = 1920,
    height: int = 1080,
    frame_count: int = 120,
    fps: int = 30,
    turntable_degrees: float = 360.0,
    turntable_axis: str = "world_z",
    bg_color: str = "",
    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,
    camera_orbit: bool = True,
    rotation_x: float = 0.0,
    rotation_y: float = 0.0,
    rotation_z: float = 0.0,
) -> dict:
    """Render a STEP file as a turntable animation (frames + FFmpeg composite).

    Returns render metadata dict on success.
    """
    import json
    import os
    import shutil
    import subprocess
    from app.services.render_blender import (
        find_blender, convert_step_to_stl, export_per_part_stls
    )

    blender_bin = find_blender()
    if not blender_bin:
        raise RuntimeError("Blender binary not found in render-worker container")

    step = Path(step_path)
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

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

    # STL conversion
    stl_path = step.parent / f"{step.stem}_{stl_quality}.stl"
    if not stl_path.exists() or stl_path.stat().st_size == 0:
        convert_step_to_stl(step, stl_path, stl_quality)
    parts_dir = step.parent / f"{step.stem}_{stl_quality}_parts"
    if not (parts_dir / "manifest.json").exists():
        try:
            export_per_part_stls(step, parts_dir, stl_quality)
        except Exception as exc:
            logger.warning("per-part export non-fatal: %s", exc)

    # Build turntable render arguments
    frames_dir = out_dir / "frames"
    frames_dir.mkdir(exist_ok=True)

    cmd = [
        blender_bin, "--background",
        "--python", str(turntable_script),
        "--",
        str(stl_path),
        str(frames_dir),
        output_name,
        str(width), str(height),
        engine, str(samples), str(smooth_angle), cycles_device,
        str(frame_count), str(fps), str(turntable_degrees), turntable_axis,
        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",
        "1" if camera_orbit else "0",
        str(rotation_x), str(rotation_y), str(rotation_z),
    ]

    try:
        result = subprocess.run(
            cmd, capture_output=True, text=True, timeout=3600
        )
        if result.returncode != 0:
            raise RuntimeError(
                f"Blender turntable exited {result.returncode}:\n{result.stdout[-2000:]}"
            )
    except Exception as exc:
        logger.error("render_turntable_task failed: %s", exc)
        raise self.retry(exc=exc, countdown=60)

    # FFmpeg composite: frames → MP4 with optional background
    output_mp4 = out_dir / f"{output_name}.mp4"
    ffmpeg_cmd = _build_ffmpeg_cmd(
        frames_dir, output_mp4, fps=fps, bg_color=bg_color
    )
    try:
        subprocess.run(ffmpeg_cmd, check=True, capture_output=True, text=True, timeout=300)
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(f"FFmpeg composite failed: {exc.stderr[-500:]}")

    return {
        "output_mp4": str(output_mp4),
        "frame_count": frame_count,
        "fps": fps,
    }


def _build_ffmpeg_cmd(
    frames_dir: Path, output_mp4: Path, fps: int = 30, bg_color: str = ""
) -> list:
    """Build FFmpeg command for compositing turntable frames to MP4."""
    import shutil as _shutil
    ffmpeg = _shutil.which("ffmpeg") or "ffmpeg"
    frame_pattern = str(frames_dir / "%04d.png")

    if bg_color:
        # Overlay transparent frames onto solid color background
        r = int(bg_color[1:3], 16) if bg_color.startswith("#") else 255
        g = int(bg_color[3:5], 16) if bg_color.startswith("#") else 255
        b = int(bg_color[5:7], 16) if bg_color.startswith("#") else 255
        color_str = f"color=c=0x{r:02x}{g:02x}{b:02x}:s=1920x1080:r={fps}"
        return [
            ffmpeg, "-y",
            "-f", "lavfi", "-i", color_str,
            "-framerate", str(fps), "-i", frame_pattern,
            "-filter_complex", "[0:v][1:v]overlay=0:0",
            "-c:v", "libx264", "-pix_fmt", "yuv420p",
            "-movflags", "+faststart",
            str(output_mp4),
        ]
    else:
        return [
            ffmpeg, "-y",
            "-framerate", str(fps), "-i", frame_pattern,
            "-c:v", "libx264", "-pix_fmt", "yuv420p",
            "-movflags", "+faststart",
            str(output_mp4),
        ]