from src.DataManager.base import BaseDataManager
from src.DataManager.utils import imwrite_rgb

import cv2
import numpy as np
from pathlib import Path
import shutil
from typing import Optional, Union

from moviepy.editor import AudioFileClip, VideoFileClip
from moviepy.video.io.ImageSequenceClip import ImageSequenceClip


class VideoDataManager(BaseDataManager):
    def __init__(self, src_data: Path, output_dir: Path, clean_work_dir: bool = False):
        self.video_handle: Optional[cv2.VideoCapture] = None
        self.audio_handle: Optional[AudioFileClip] = None

        self.output_dir = output_dir
        self.output_img_dir = output_dir / "img"
        self.output_dir.mkdir(exist_ok=True)
        self.output_img_dir.mkdir(exist_ok=True)
        self.video_name = None
        self.clean_work_dir = clean_work_dir

        if src_data.is_file():
            self.video_name = "swap_" + src_data.name

            if VideoFileClip(str(src_data)).audio is not None:
                self.audio_handle = AudioFileClip(str(src_data))

            self.video_handle = cv2.VideoCapture(str(src_data))
            self.video_handle.set(cv2.CAP_PROP_POS_FRAMES, 0)

            self.frame_count = int(self.video_handle.get(cv2.CAP_PROP_FRAME_COUNT))
            self.fps = self.video_handle.get(cv2.CAP_PROP_FPS)

        self.last_idx = -1

        assert self.video_handle, "Video file must be specified!"

    def __len__(self):
        return self.frame_count

    def get(self) -> np.ndarray:
        img: Union[None, np.ndarray] = None

        while img is None and self.last_idx < self.frame_count:
            status, img = self.video_handle.read()
            self.last_idx += 1

        if img is not None:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img

    def save(self, img: np.ndarray):
        filename = "frame_{:0>7d}.jpg".format(self.last_idx)
        imwrite_rgb(self.output_img_dir / filename, img)

        if (self.frame_count - 1) == self.last_idx:
            self._close()

    def _close(self):
        image_filenames = [str(x) for x in sorted(self.output_img_dir.glob("*.jpg"))]
        clip = ImageSequenceClip(image_filenames, fps=self.fps)

        if self.audio_handle is not None:
            clip = clip.set_audio(self.audio_handle)

        clip.write_videofile(str(self.output_dir / self.video_name))

        if self.clean_work_dir:
            shutil.rmtree(self.output_img_dir, ignore_errors=True)