"""
Video generation utilities for visualizing agent behavior.

Wraps Gymnasium's video recording functionality to generate MP4 videos
of agent rollouts in the environment.
"""

from pathlib import Path
import gymnasium as gym

def generate_video(
    agent,
    env_id: str,
    output_dir: Path,
    num_episodes: int = 3,
    video_prefix: str = "video",
    deterministic: bool = True,
) -> None:
    """
    Generate video recordings of agent episodes.
    
    Args:
        agent: Agent to record
        env_id: Gymnasium environment ID
        output_dir: Directory to save videos
        num_episodes: Number of episodes to record
        video_prefix: Prefix for video filenames
        deterministic: If True, use deterministic policy
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Create environment with video recording
    env = gym.make(env_id, render_mode="rgb_array")
    env = gym.wrappers.RecordVideo(
        env,
        video_folder=str(output_dir),
        name_prefix=video_prefix,
        episode_trigger=lambda ep: True,  # Record all episodes
    )
    
    print(f"Recording {num_episodes} episodes to: {output_dir}")
    
    for ep in range(num_episodes):
        obs, _ = env.reset()
        done = False
        episode_return = 0.0
        
        while not done:
            action = agent.get_action(obs, deterministic=deterministic)
            obs, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            episode_return += reward
        
        print(f"  Episode {ep + 1}: Return = {episode_return:.2f}")
    
    env.close()
    print(f" Videos saved to: {output_dir}")