#!/opt/ias/env/bin/python

import argparse
from pathlib import Path

import cv2
from jvi.geo.size2d import size_parse


def main():
    parser = argparse.ArgumentParser(description="图片转视频工具")
    parser.add_argument(
        "src_dir", metavar="SRC_DIR", type=Path, help="来源图片所在目录"
    )
    parser.add_argument("dst_file", metavar="DST_FILE", type=Path, help="目的视频文件")
    parser.add_argument(
        "-p", "--pattern", type=str, default="*.jpg", help="图片文件模式(*.jpg)"
    )
    parser.add_argument("-s", "--size", type=str, default="HD", help="输出视频尺寸(HD)")
    parser.add_argument("-f", "--fps", type=float, default=5, help="视频帧率(5)")
    parser.add_argument(
        "-r", "--repeat", type=int, default=1, help="视频帧的重复次数(1)"
    )
    parser.add_argument("-v", "--verbose", action="store_true", help="显示详细信息")
    opt = parser.parse_args()

    size = size_parse(opt.size)

    fourcc_tab = {
        ".mp4": "mp4v",
        ".mjpg": "mjpg",
        ".webp": "webp",
    }  # require libwebp-dev
    # TODO: avc1/MJPG报错，https://stackoverflow.com/questions/52932157/ \
    #  opencv-ffmpeg-tag-0x34363268-h264-is-not-supported-with-codec/56723380

    fourcc = fourcc_tab.get(opt.dst_file.suffix)
    if not fourcc:
        print("[ERROR] 不支持的文件格式:", opt.dst_file.suffix)
        return -1
    fourcc = cv2.VideoWriter_fourcc(*fourcc)
    writer = cv2.VideoWriter(str(opt.dst_file), fourcc, opt.fps, size.to_tuple_int())

    files = opt.src_dir.rglob(opt.pattern)
    files = sorted(files)
    print("源图片目录：%s" % opt.src_dir)

    dst_img = None
    ok = 0
    err = 0
    for i, f in enumerate(files):
        src_img = ImageNda.load(str(f))  # BGR
        name = f.name
        if src_img is None:
            err += 1
            print(i, name, "\tERROR")
            continue
        else:
            ok += 1
            if opt.verbose:
                print(i, name, "\tOK")

        dst_img = cv2.resize(src_img, size.to_tuple_int(), dst_img)
        for _i in range(opt.repeat):
            writer.write(dst_img)

    print("\n压缩帧数：(%d/%d)\n" % (ok, ok + err))


if __name__ == "__main__":
    main()
