# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import random

from cosmos_rl.utils.logging import logger


def save_rollout_video(
    rollout_images, rollout_dir: str, task_name: str, success: bool
) -> str:
    """
    Saves an MP4 replay of an episode.

    Args:
        rollout_images: List of images (numpy arrays) to save as video
        exp_name: Experiment name for organizing videos
        task_name: Task identifier
        step_idx: Current training step index
        success: Whether the episode was successful

    Returns:
        str: Path to the saved video file
    """

    try:
        import imageio
    except ImportError:
        logger.warning(
            "imageio not installed, cannot save rollout videos. Install with: pip install imageio imageio-ffmpeg"
        )
        return ""

    # Create rollout directory
    os.makedirs(rollout_dir, exist_ok=True)

    # Generate unique filename
    ran_id = random.randint(1, 10000)
    mp4_path = f"{rollout_dir}/task={task_name}--success={success}--ran={ran_id}.mp4"

    # Write video
    video_writer = imageio.get_writer(mp4_path, fps=30)
    for img in rollout_images:
        video_writer.append_data(img)
    video_writer.close()

    return mp4_path
