import imageio
import numpy as np
from tqdm import tqdm


def save_video(frames, save_path, fps, quality=9, ffmpeg_params=None):

    if len(frames) == 1 :
        frames[0].save(save_path.replace('.mp4', '.png'))
        return

    # Ensure browser-compatible pixel format (yuv420p) to avoid Gradio conversion warning
    if ffmpeg_params is None:
        ffmpeg_params = ['-pix_fmt', 'yuv420p']
    elif '-pix_fmt' not in ffmpeg_params:
        ffmpeg_params = list(ffmpeg_params) + ['-pix_fmt', 'yuv420p']
    
    writer = imageio.get_writer(save_path, fps=fps, quality=quality, ffmpeg_params=ffmpeg_params)
    for frame in tqdm(frames, desc="Saving video"):
        frame = np.array(frame)
        writer.append_data(frame)
    writer.close()