hkuds--deeptutor/deeptutor/agents/math_animator/renderer.py

"""Manim rendering service for the math animator capability."""

from __future__ import annotations

import asyncio
from pathlib import Path
import re
import subprocess
import sys
import threading
from typing import Awaitable, Callable

from deeptutor.services.path_service import get_path_service

from .models import RenderedArtifact, RenderResult
from .utils import slugify_filename, trim_error_message

YON_IMAGE_PATTERN = re.compile(
    r"###\s*YON_IMAGE_(\d+)_START\s*###\s*(.*?)\s*###\s*YON_IMAGE_\1_END\s*###",
    re.DOTALL | re.IGNORECASE,
)
SCENE_PATTERN = re.compile(r"class\s+([A-Za-z_][A-Za-z0-9_]*)\s*\(\s*.*?Scene.*?\)\s*:")

QUALITY_FLAG_MAP = {
    "low": "-ql",
    "medium": "-qm",
    "high": "-qh",
}


class ManimRenderError(RuntimeError):
    """Raised when Manim rendering fails."""


class ManimRenderService:
    def __init__(
        self,
        turn_id: str,
        progress_callback: Callable[[str, bool], Awaitable[None]] | None = None,
    ) -> None:
        self.turn_id = turn_id
        self.path_service = get_path_service()
        self.progress_callback = progress_callback
        self.base_dir = self.path_service.get_agent_dir("math_animator") / turn_id
        self.source_dir = self.base_dir / "source"
        self.artifacts_dir = self.base_dir / "artifacts"
        self.media_dir = self.base_dir / "media"
        self.meta_dir = self.base_dir / "meta"
        for path in (self.source_dir, self.artifacts_dir, self.media_dir, self.meta_dir):
            path.mkdir(parents=True, exist_ok=True)

    async def render(self, *, code: str, output_mode: str, quality: str) -> RenderResult:
        await self._emit_progress(f"Preparing {output_mode} render workspace (quality={quality}).")
        source_name = "scene.py" if output_mode == "video" else "scene_image.py"
        source_path = self.source_dir / source_name
        source_path.write_text(code, encoding="utf-8")
        await self._emit_progress(f"Saved generated code to {source_name}.", raw=True)

        if output_mode == "image":
            artifacts = await self._render_image_blocks(code=code, quality=quality)
        else:
            artifacts = [await self._render_video(code_path=source_path, quality=quality)]

        return RenderResult(
            output_mode=output_mode,
            artifacts=artifacts,
            source_code_path=str(source_path),
            quality=quality,
        )

    async def _render_video(self, *, code_path: Path, quality: str) -> RenderedArtifact:
        scene_name = self._extract_scene_name(code_path.read_text(encoding="utf-8"))
        await self._emit_progress(f"Launching Manim scene `{scene_name}`.")
        await self._run_manim(
            code_path=code_path, scene_name=scene_name, quality=quality, save_last_frame=False
        )
        video_file = self._find_rendered_file(".mp4")
        target_name = slugify_filename(f"{self.turn_id}-{scene_name}.mp4", f"{self.turn_id}.mp4")
        artifact_path = self.artifacts_dir / target_name
        artifact_path.write_bytes(video_file.read_bytes())
        await self._emit_progress(f"Saved rendered video as {artifact_path.name}.")
        return self._build_artifact(artifact_path, "video", "video/mp4", "Animation video")

    async def _render_image_blocks(self, *, code: str, quality: str) -> list[RenderedArtifact]:
        matches = list(YON_IMAGE_PATTERN.finditer(code))
        if not matches:
            raise ManimRenderError(
                "Image mode requires code blocks wrapped in ### YON_IMAGE_n_START ### / END ###."
            )

        residual = YON_IMAGE_PATTERN.sub("", code).strip()
        if residual:
            raise ManimRenderError("Image mode code must only contain YON_IMAGE anchor blocks.")

        artifacts: list[RenderedArtifact] = []
        for idx, match in enumerate(matches, start=1):
            block_code = match.group(2).strip()
            block_path = self.source_dir / f"image_block_{idx:02d}.py"
            block_path.write_text(block_code, encoding="utf-8")
            scene_name = self._extract_scene_name(block_code)
            await self._emit_progress(
                f"Rendering image block {idx}/{len(matches)} with scene `{scene_name}`."
            )
            await self._run_manim(
                code_path=block_path,
                scene_name=scene_name,
                quality=quality,
                save_last_frame=True,
            )
            image_file = self._find_rendered_file(".png")
            artifact_path = self.artifacts_dir / f"image-{idx:02d}.png"
            artifact_path.write_bytes(image_file.read_bytes())
            await self._emit_progress(f"Saved image artifact {artifact_path.name}.")
            artifacts.append(
                self._build_artifact(
                    artifact_path,
                    "image",
                    "image/png",
                    f"Image {idx}",
                )
            )
        return artifacts

    async def _run_manim(
        self,
        *,
        code_path: Path,
        scene_name: str,
        quality: str,
        save_last_frame: bool,
    ) -> None:
        quality_flag = QUALITY_FLAG_MAP.get(quality, "-qm")
        command = [
            sys.executable,
            "-m",
            "manim",
            quality_flag,
            str(code_path),
            scene_name,
            "--media_dir",
            str(self.media_dir),
            "--progress_bar",
            "none",
        ]
        if save_last_frame:
            command.append("-s")
        else:
            command.extend(["--format", "mp4"])

        await self._emit_progress(
            f"Started Manim process for `{scene_name}` with command: {' '.join(command)}",
            raw=True,
        )

        # Use subprocess.Popen instead of asyncio.create_subprocess_exec
        # for Windows compatibility (SelectorEventLoop doesn't support
        # asyncio subprocesses). Reader threads + asyncio.Queue preserve
        # real-time streaming output.
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        _SENTINEL = None
        queue: asyncio.Queue[tuple[str, str] | None] = asyncio.Queue()
        loop = asyncio.get_running_loop()

        def _reader(stream, prefix: str) -> None:
            assert stream is not None
            for raw_line in stream:
                line = raw_line.decode(errors="ignore").strip()
                if line:
                    loop.call_soon_threadsafe(queue.put_nowait, (prefix, line))
            loop.call_soon_threadsafe(queue.put_nowait, _SENTINEL)

        threading.Thread(target=_reader, args=(process.stdout, "stdout"), daemon=True).start()
        threading.Thread(target=_reader, args=(process.stderr, "stderr"), daemon=True).start()

        stdout_lines: list[str] = []
        stderr_lines: list[str] = []
        streams_open = 2
        while streams_open > 0:
            item = await queue.get()
            if item is _SENTINEL:
                streams_open -= 1
                continue
            prefix, line = item
            (stdout_lines if prefix == "stdout" else stderr_lines).append(line)
            await self._emit_progress(f"[{prefix}] {line}", raw=True)

        return_code = process.wait()
        await self._emit_progress(f"Manim process finished with exit code {return_code}.", raw=True)
        if return_code != 0:
            raise ManimRenderError(
                trim_error_message(
                    "\n".join(
                        part for part in ["\n".join(stdout_lines), "\n".join(stderr_lines)] if part
                    )
                )
            )

    async def _emit_progress(self, message: str, raw: bool = False) -> None:
        if self.progress_callback is None:
            return
        await self.progress_callback(message, raw)

    def _find_rendered_file(self, suffix: str) -> Path:
        # Manim stores many transient chunks under ``partial_movie_files``.
        # We only want the final exported artifact for the scene.
        matches = [
            path
            for path in self.media_dir.rglob(f"*{suffix}")
            if "partial_movie_files" not in path.parts
        ]
        if not matches:
            matches = list(self.media_dir.rglob(f"*{suffix}"))
        if not matches:
            raise ManimRenderError(f"Rendered {suffix} artifact not found.")
        return max(matches, key=lambda path: path.stat().st_mtime)

    @staticmethod
    def _extract_scene_name(code: str) -> str:
        match = SCENE_PATTERN.search(code)
        if not match:
            raise ManimRenderError("Generated code does not define a renderable Manim Scene class.")
        return match.group(1)

    def _build_artifact(
        self,
        artifact_path: Path,
        artifact_type: str,
        content_type: str,
        label: str,
    ) -> RenderedArtifact:
        rel_path = artifact_path.resolve().relative_to(self.path_service.user_data_dir.resolve())
        return RenderedArtifact(
            type=artifact_type,
            filename=artifact_path.name,
            url=f"/api/outputs/{rel_path.as_posix()}",
            content_type=content_type,
            label=label,
        )


__all__ = [
    "ManimRenderError",
    "ManimRenderService",
    "YON_IMAGE_PATTERN",
]