renderer.py

"""
Manim rendering engine - handles subprocess execution and video generation
"""

import asyncio
import os
import shutil
import logging
from datetime import datetime
from pathlib import Path
from typing import Tuple, Optional, List
from config import QUALITY_PRESETS, RENDER_TIMEOUT, BASE_RENDER_DIR

logger = logging.getLogger("manim_studio.renderer")


async def render_animation(
    request_id: str,
    code: str,
    quality: str,
    fps: Optional[int],
    scene_name: str
) -> Tuple[bool, str, Optional[Path]]:
    """
    Execute Manim rendering in subprocess

    Args:
        request_id: Unique identifier for this render
        code: Python code containing Manim Scene
        quality: Quality preset (4k, 2k, 1080p, 720p, 480p)
        fps: Custom FPS (overrides quality preset)
        scene_name: Name of Scene class to render

    Returns:
        Tuple of (success, message, output_path)
        - success: True if render completed successfully
        - message: Status message or error description
        - output_path: Path to rendered MP4 file (None if failed)
    """
    # Create request directory
    request_dir = BASE_RENDER_DIR / request_id
    request_dir.mkdir(parents=True, exist_ok=True)

    source_file = request_dir / "source.py"
    media_dir = request_dir / "media"

    try:
        # Write code to source file
        source_file.write_text(code, encoding="utf-8")
        logger.info(f"[{request_id}] Source file written: {source_file}")

        # Build manim command
        command = build_manim_command(
            source_file=source_file,
            quality=quality,
            fps=fps,
            scene_name=scene_name,
            media_dir=media_dir
        )

        logger.info(f"[{request_id}] Executing command: {' '.join(command)}")

        # Execute manim with timeout
        success, message, stdout, stderr = await execute_manim_command(
            command=command,
            timeout=RENDER_TIMEOUT,
            request_id=request_id
        )

        if not success:
            logger.error(f"[{request_id}] Render failed: {message}")
            # Log stdout and stderr for debugging
            if stdout:
                logger.error(f"[{request_id}] Manim stdout: {stdout[:1000]}")
            if stderr:
                logger.error(f"[{request_id}] Manim stderr: {stderr[:1000]}")

            # Return detailed error message
            error_details = f"{message}\n\nStdout: {stdout[:500]}\n\nStderr: {stderr[:500]}"
            return False, error_details, None

        # Find the output video file
        logger.info(f"[{request_id}] Searching for output video in: {media_dir}")

        # Debug: List all files created
        if media_dir.exists():
            logger.info(f"[{request_id}] Media directory exists, listing contents...")
            for root, dirs, files in os.walk(media_dir):
                logger.info(f"[{request_id}]   Dir: {root}")
                for f in files:
                    logger.info(f"[{request_id}]     File: {f}")
        else:
            logger.error(f"[{request_id}] Media directory does NOT exist: {media_dir}")

        output_video = find_output_video(media_dir, quality, fps)

        if not output_video or not output_video.exists():
            # Try to find ANY mp4 file in the media directory
            mp4_files = list(media_dir.rglob("*.mp4")) if media_dir.exists() else []

            if mp4_files:
                logger.info(f"[{request_id}] Found MP4 files via fallback search: {mp4_files}")
                output_video = mp4_files[0]  # Use the first one found
            else:
                error_msg = f"Render completed but output video not found. Expected in {media_dir}"
                logger.error(f"[{request_id}] {error_msg}")
                return False, error_msg, None

        # Copy to standardized location
        standardized_output = request_dir / "output.mp4"
        shutil.copy2(output_video, standardized_output)

        logger.info(f"[{request_id}] Render completed successfully: {standardized_output}")
        return True, "Render completed successfully", standardized_output

    except Exception as e:
        error_msg = f"Unexpected error during rendering: {str(e)}"
        logger.exception(f"[{request_id}] {error_msg}")
        return False, error_msg, None


def build_manim_command(
    source_file: Path,
    quality: str,
    fps: Optional[int],
    scene_name: str,
    media_dir: Path
) -> List[str]:
    """
    Build Manim CLI command with appropriate flags

    Args:
        source_file: Path to Python file with Manim code
        quality: Quality preset key
        fps: Custom FPS (None to use quality preset default)
        scene_name: Scene class name to render
        media_dir: Output directory for media files

    Returns:
        List of command arguments for subprocess execution
    """
    # Get quality preset
    preset = QUALITY_PRESETS.get(quality, QUALITY_PRESETS["720p"])
    quality_flag = preset["flag"]

    # Base command
    command = [
        "manim",
        quality_flag,
        "--media_dir", str(media_dir),
        "--disable_caching",  # Disable caching for clean renders
    ]

    # Add custom FPS if specified
    if fps is not None:
        command.extend(["--fps", str(fps)])

    # Add source file and scene name
    command.append(str(source_file))
    command.append(scene_name)

    return command


async def execute_manim_command(
    command: List[str],
    timeout: int,
    request_id: str
) -> Tuple[bool, str, str, str]:
    """
    Execute Manim command as subprocess with timeout

    Args:
        command: Command arguments list
        timeout: Maximum execution time in seconds
        request_id: Request ID for logging

    Returns:
        Tuple of (success, message, stdout, stderr)
    """
    async def _read_stream(stream, buffer, log_file, prefix):
        while True:
            chunk = await stream.read(1024)
            if not chunk:
                break
            text = chunk.decode("utf-8", errors="replace")
            buffer.append(text)
            if log_file:
                log_file.write(f"{prefix}{text}")
                log_file.flush()

    log_file = None
    try:
        # Set up environment for non-interactive Manim
        import os
        env = os.environ.copy()
        env["TERM"] = "xterm-256color"
        env["PYTHONUNBUFFERED"] = "1"

        # Create subprocess
        process = await asyncio.create_subprocess_exec(
            *command,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env
        )

        log_path = BASE_RENDER_DIR / request_id / "render.log"
        log_path.parent.mkdir(parents=True, exist_ok=True)
        log_file = log_path.open("a", encoding="utf-8")
        log_file.write(f"[{datetime.utcnow().isoformat()}] Starting manim render\n")
        log_file.flush()

        stdout_chunks: List[str] = []
        stderr_chunks: List[str] = []

        read_stdout = asyncio.create_task(
            _read_stream(process.stdout, stdout_chunks, log_file, "")
        )
        read_stderr = asyncio.create_task(
            _read_stream(process.stderr, stderr_chunks, log_file, "ERR: ")
        )

        # Wait for completion with timeout
        try:
            await asyncio.wait_for(process.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            # Kill process on timeout
            process.kill()
            await process.wait()
            error_msg = f"Render exceeded timeout of {timeout} seconds"
            logger.error(f"[{request_id}] {error_msg}")
            await asyncio.gather(read_stdout, read_stderr, return_exceptions=True)
            if log_file:
                log_file.write(f"[{datetime.utcnow().isoformat()}] {error_msg}\n")
                log_file.flush()
            return False, error_msg, "".join(stdout_chunks), "".join(stderr_chunks)

        await asyncio.gather(read_stdout, read_stderr, return_exceptions=True)

        stdout = "".join(stdout_chunks)
        stderr = "".join(stderr_chunks)

        # Check return code
        if process.returncode == 0:
            if log_file:
                log_file.write(f"[{datetime.utcnow().isoformat()}] Manim finished successfully\n")
                log_file.flush()
            return True, "Command executed successfully", stdout, stderr
        else:
            error_msg = f"Manim command failed with exit code {process.returncode}"
            if stderr:
                error_msg += f"\n\nError output:\n{stderr}"
            if log_file:
                log_file.write(f"[{datetime.utcnow().isoformat()}] {error_msg}\n")
                log_file.flush()
            return False, error_msg, stdout, stderr

    except Exception as e:
        error_msg = f"Failed to execute command: {str(e)}"
        logger.exception(f"[{request_id}] {error_msg}")
        return False, error_msg, "", ""
    finally:
        if log_file:
            log_file.close()


def find_output_video(media_dir: Path, quality: str, fps: Optional[int]) -> Optional[Path]:
    """
    Locate the rendered video file in Manim's output directory structure

    Manim creates videos in: media/videos/{source_file_name}/{quality_fps}/{scene_name}.mp4

    Args:
        media_dir: Base media directory
        quality: Quality preset
        fps: Custom FPS (None if using preset default)

    Returns:
        Path to output video file, or None if not found
    """
    videos_dir = media_dir / "videos" / "source"

    if not videos_dir.exists():
        logger.warning(f"Videos directory not found: {videos_dir}")
        return None

    # Determine quality directory name
    preset = QUALITY_PRESETS.get(quality, QUALITY_PRESETS["720p"])
    if fps is not None:
        # Custom FPS: e.g., "720p30"
        resolution_height = preset["resolution"].split("x")[1]
        quality_dir_name = f"{resolution_height}p{fps}"
    else:
        # Use preset FPS
        resolution_height = preset["resolution"].split("x")[1]
        preset_fps = preset["fps"]
        quality_dir_name = f"{resolution_height}p{preset_fps}"

    quality_dir = videos_dir / quality_dir_name

    if not quality_dir.exists():
        # Try to find any subdirectory (fallback)
        logger.warning(f"Expected quality directory not found: {quality_dir}")
        subdirs = list(videos_dir.iterdir())
        if subdirs:
            quality_dir = subdirs[0]
            logger.info(f"Using alternative directory: {quality_dir}")
        else:
            return None

    # Find .mp4 files in the quality directory
    mp4_files = list(quality_dir.glob("*.mp4"))

    if mp4_files:
        # Return the first (and typically only) MP4 file
        return mp4_files[0]
    else:
        logger.warning(f"No MP4 files found in: {quality_dir}")
        return None


def cleanup_temp_files(request_dir: Path) -> None:
    """
    Clean up temporary files created during rendering (keep only output.mp4)

    Args:
        request_dir: Directory containing render files
    """
    try:
        # Remove source.py
        source_file = request_dir / "source.py"
        if source_file.exists():
            source_file.unlink()

        # Remove media directory (contains intermediate files)
        media_dir = request_dir / "media"
        if media_dir.exists():
            shutil.rmtree(media_dir)

        logger.info(f"Cleaned up temporary files in {request_dir}")
    except Exception as e:
        logger.warning(f"Failed to cleanup temp files in {request_dir}: {e}")