import numpy as np
import imageio
from PIL import Image
from typing import List


def video_frame_picker(
    video_path: str,
    offscreen: bool = True,
    vis_widown_path: str = 'test.png'   # if offscreen, save to this path
) -> List[int]:
    # return the frame indices user picked
    video = imageio.get_reader(video_path)
    len_video = video.count_frames()
    print(f"video has {len_video} frames")

    indices_picked = []
    for i in range(len_video):
        frame = video.get_data(i)
        img = Image.fromarray(frame)
        if offscreen:
            img.save(vis_widown_path)
        else:
            img.show()
        while (
            user_input := input(f"Frame {i}: 'y' for yes, 'n' for no, 'q' to quit: ")
        ) not in ['y', 'n', 'q']:
            print("Invalid input")
        if user_input == 'q':
            return indices_picked
        if user_input == 'y':
            indices_picked.append(i)
        if user_input == 'n':
            continue
    return indices_picked


if __name__ == '__main__':
    video_frame_picker(
        video_path="data/policy_eval_videos/policy_1.00/00.mp4",
    )