from pathlib import Path
from argparse import ArgumentParser

import json

from typing import Union


def generate_vae_annotation(data_root: str,
                            data_type: str,
                            save_path: str):
    assert data_type in ['image', 'video']
    
    data_root = Path(data_root)
    
    if data_type == "video":
        data_path_list = [*data_root.glob("*.mp4")]
    elif data_type == "image":
        data_path_list = [*data_root.glob("*.png"), *data_root.glob("*.jpeg"),
                          *data_root.glob("*.jpg"), *data_root.glob("*.JPEG")]
    else:
        raise NotImplemented
    
    with open(save_path, "w") as f:
        for data_path in data_path_list:
            f.write(json.dumps({data_type: str(data_path.resolve())}, ensure_ascii=False) + '\n')


if __name__ == '__main__':
    parser = ArgumentParser()
    
    parser.add_argument("--data_root", type=str)

    parser.add_argument("--data_type", type=str)

    parser.add_argument("--save_path", type=str)
    
    args = parser.parse_args()
    
    generate_vae_annotation(args.data_root,
                            args.data_type,
                            args.save_path)