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

import argparse
from pathlib import Path

import arrow
import cv2
from arrow import Arrow
from numpy import ndarray

from jcx.sys.fs import time_to_file, make_parents
from jcx.ui.key import Key
from jvi.geo.point2d import Point
from jvi.geo.size2d import size_parse, Size
from jvi.gui.video_viewer import VideoViewer

TITLE = "video2image"


class ImageCapture(VideoViewer):
    """图片抓取"""

    def __init__(
        self, url, dst_dir: Path, size: Size = Size(), start_time: Arrow = arrow.now()
    ):
        super().__init__(url, TITLE, size)
        self.__dst_dir = dst_dir
        self.__start_time = start_time

    def on_key(self, key: int):
        """按键响应"""
        if key == Key.BLANK:
            self.__save_image()
        else:
            return super().on_key(key)
        return 0

    def on_draw(self, canvas: ndarray, pos: Point):
        """设置窗口重绘事件响应"""
        # self.record().draw_on(canvas, pos)
        pass

    def __save_image(self):
        """保存图片"""
        image = self.image_at(self._index)
        if image is None:
            print("[ERROR] image is none")
            return

        interval = self._index * 1000000 / 25
        time = self.__start_time.shift(microseconds=interval)
        file = self.__dst_dir / time_to_file(time, ".jpg")
        make_parents(file)
        cv2.imwrite(str(file), image)
        print("save to:", file)


# file = 'file:///home/jiang/ws/lift/2021-08-19/1.mp4'


def main():
    parser = argparse.ArgumentParser(description="视频转图片工具")
    parser.add_argument("src_url", type=Path, help="来源视频URL")
    parser.add_argument("dst_dir", type=Path, help="目的图片所在目录")
    parser.add_argument("-e", "--ext", type=str, default=".jpg", help="图片文件扩展名")
    parser.add_argument("-s", "--size", type=str, default="HD", help="输出视频尺寸(HD)")
    parser.add_argument("-i", "--interval", type=float, default=1, help="时间间隔")
    parser.add_argument("-t", "--start_time", type=str, help="图片开始时间")
    parser.add_argument("-v", "--verbose", action="store_true", help="显示详细信息")
    opt = parser.parse_args()

    size = size_parse(opt.size)
    # TODO: start_time解析
    win = ImageCapture(opt.src_url, opt.dst_dir, size=Size(1280, 720))
    if not win.is_opened():
        print("[ERROR] cannot open url:", opt.src_url)
    print("frame count:", win.image_count())
    print("video size:", win.image_size())
    # print(opt.start_time)

    win.run()


if __name__ == "__main__":
    main()
