import multiprocessing
import subprocess
import argparse
import datetime
import requests
import warnings
import hashlib
import typing
import random
import glob
import json
import time
import os

from transformers import GPT2Tokenizer
from google.cloud import storage
import tensorflow as tf
import numpy as np
import youtube_dl
import cv2


def _int64_feature(value):
    """Returns an int64_list."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


def _bytes_feature(value):
    """Returns an bytes_list."""
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def diveision_zero(x, y):
    '''
    x / y
    Helper function to divide number by zero.
    IF divided by zero zero will be return.
    '''

    try:
        return x / y
    except ZeroDivisionError:
        return 0


def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


class Downloader:

    def __init__(self, max_try: int = 3, webshare_io_key: str = None):
        self.max_try = max_try
        self.webshare_io_key = webshare_io_key
        self.proxies = None

        self.update_proxy()

    def download(self, url: str, filename: str, use_proxy: bool):
        try_count = 0

        while try_count < self.max_try:
            try:
                if use_proxy:
                    r = requests.get(url, stream=True, proxies=self.proxies)
                else:
                    r = requests.get(url, stream=True)

                with open(filename, 'wb') as f:
                    for chunk in r:
                        f.write(chunk)

            except Exception as e:
                print(f'Download error with proxy: {use_proxy} \n error:', e)
                try_count += 1

                if use_proxy:
                    self.update_proxy()

            else:
                return True

        print('Retry exceeded for URL:', url)

        if os.path.exists(filename):
            os.remove(filename)

        return False

    def update_proxy(self):

        if self.webshare_io_key is not None:

            proxies = []
            _next = "/api/proxy/list/?page=1"

            while _next is not None:

                r = requests.get('https://proxy.webshare.io' + _next,
                                 headers={"Authorization": f"Token {self.webshare_io_key}"})
                dump = r.json()

                _next = None
                if dump is not None:
                    if 'next' in dump:
                        _next = dump['next']

                    if 'results' in dump:
                        proxies = proxies + [res for res in dump['results'] if res['valid']]

            random.shuffle(proxies)
            random.shuffle(proxies)

            if len(proxies) > 0:
                username = proxies[0]['username']
                password = proxies[0]['password']
                proxy_address = proxies[0]['proxy_address']
                ports = proxies[0]['ports']['http']

                self.proxies = {"http": f"http://{username}:{password}@{proxy_address}:{ports}",
                                "https": f"http://{username}:{password}@{proxy_address}:{ports}"}
            else:
                self.proxies = None


def frame_encoder(frame,
                  text_tokens: typing.List[int] = None,
                  skip_frame: typing.List[bool] = [False],
                  mask: typing.List[int] = None,
                  concat: typing.List[bool] = [False]):
    '''
    :param frame: A byte String containing a jpg encoded image.
    :param text_tokens: A list containing int ped tokens.
    :param skip_frame: A list containing a single bool that
    determines if this frame include an image or just text.
    :param mask: A int that determines when the padding tokens start.
    :param concat: A list containing a single bool that determents if tow videos are concat at that point.

    This Function will encode frame to proto buffer.
    '''

    # Encoding Key.
    feature = {
        'frame': _bytes_feature(frame),
        'concat': _int64_feature(concat)
    }

    if text_tokens is not None:
        feature.update({'tokens': _int64_feature(text_tokens),
                        'skip_frame': _int64_feature(skip_frame),
                        'mask': _int64_feature(mask)})

    # Encode
    proto = tf.train.Example(features=tf.train.Features(feature=feature))

    # Serialize proto buffer to string.
    proto = proto.SerializeToString()

    return proto


def split_equal(ids: list, duration: list, num: int, min_duration: int = 256):

    sort = sorted(zip(duration, ids))[::-1]

    ids_split = [[] for i in range(num)]
    duration_spit = [[] for i in range(num)]
    duration_sum = [0] * num

    for d, i in sort:
        if d > min_duration or min_duration <= 0:
            pos = np.argmin(duration_sum)

            ids_split[pos].append(i)
            duration_spit[pos].append(d)
            duration_sum[pos] = duration_sum[pos] + d

    return ids_split, duration_spit


def decode_vtt(content: str):
    '''
    :param content: String with the of the .vtt file.
    :return: String with combined text, List with Strings split at time stamps and List with float second time stamps.

    This Function decodes a vtt to get the contend with  time stamps.
    '''

    if '</c><' in content and '><c>' in content:

        # Split the content at line brake and check if word level time stamps are in the line.
        content = [l for l in content.split('\n') if '<c>' in l]

        # Connect list of strings back together.
        content = "".join(content)

        # Split String at time stamp headers.
        content = content.split('><c>')

        # Create output lists.
        words = []
        stamps = []

        # Loop word and time stamp string list.
        for c in content:

            # Split word and time stamp part.
            word = c[:-12]
            stam = c[-12:]

            # Clean word string.
            if not '</c><' in word:
                word = word.replace('</c>', ' ')

            word = word.replace('</c>', '')
            word = word.replace('<', '')
            word = word.lstrip().rstrip()

            # Check if time stamp is in stamp string.
            if ':' in stam and '.' in stam:

                # Split time stamp string at time punctuation marks.
                stam = stam.split(':')
                stam = stam[:-1] + stam[-1].split('.')

                # Converting time stamp string in to second based float.
                stam = datetime.timedelta(hours=int(stam[0]), minutes=int(stam[1]), seconds=int(stam[2]), milliseconds=int(stam[3]))
                stam = stam.total_seconds()

                # add word string and second based float to output list.
                words.append(' ' + word)
                stamps.append(stam)
            else:
                # If no time stamp contain in time stamp part we assume that it is a another word.
                # If it as a another word we add it to the previous word string.
                if len(words) > 0:
                    words[-1] = words[-1] + " " + c.replace('</c>', '').replace('<', '').lstrip().rstrip()

        return ''.join(words), words, stamps

    else:

        # Split the content at line brake.
        content = content.split('\n')

        # Create output lists.
        words_buffer = []
        stamps_buffer = []
        words = []
        stamps = []

        # Loop word and time stamp string list.
        for idx in range(len(content)):
            if ' --> ' in content[idx]:
                stamps_buffer.append(content[idx])

                word_buffer = []
                idx += 1
                while idx + 1 < len(content) and ' --> ' not in content[idx + 1]:
                    word_buffer.append(content[idx])
                    idx += 1

                words_buffer.append(" ".join(word_buffer))

        for idx in range(len(stamps_buffer)):
            s = stamps_buffer[idx].split(' --> ')

            s_1 = s[0]
            s_1 = s_1.split(':')
            s_1 = s_1[:-1] + s_1[-1].split('.')

            s_2 = s[1]
            s_2 = s_2.split(':')
            s_2 = s_2[:-1] + s_2[-1].split('.')

            s_1 = datetime.timedelta(hours=int(s_1[0]), minutes=int(s_1[1]), seconds=int(s_1[2]),
                                     milliseconds=int(s_1[3]))
            s_1 = s_1.total_seconds()

            s_2 = datetime.timedelta(hours=int(s_2[0]), minutes=int(s_2[1]), seconds=int(s_2[2]),
                                     milliseconds=int(s_2[3]))
            s_2 = s_2.total_seconds()

            stamps_buffer[idx] = [s_1, s_2]

        for idx in range(len(words_buffer)):
            word = words_buffer[idx].lstrip().rstrip()
            wor = [' ' + w for w in word.split(' ')]

            stamp = stamps_buffer[idx]

            time_snip = (stamp[1] - stamp[0]) / len(wor)

            stamps += [stamp[0] + i * time_snip for i in range(len(wor))]
            words += wor

        return ''.join(words), words, stamps


def bpe_with_word_split(enc: GPT2Tokenizer, words: list, text: str):
    '''
    :param enc: BPE encoder to be used.
    :param words: List with word strings split at time stamps.
    :param text: String with total text. the sum of text in words list should be the same as the content of the text string.
    :return: List with list containing BPE tokens.

    This function will encode a text using a given tokenizer and
    that split it back up according to the string split provided by word list.
    '''

    # Encode text with tokenizer.
    tokens = enc.encode(text, max_length=None)

    # Decode single token.
    pair_split_word = [enc.decode(token).replace(' ', '') for token in tokens]

    # Creat output list.
    bpe_list = []

    # Creat token split index.
    idx = 0

    # Loop through word list.
    for word in words:

        # Create buffer list to holt the tokens for one time stamp interval.
        buffer = []

        # Buffer word string and remove empty space between single words.
        word_buffer = word.replace(' ', '')

        # Continue looping es long  idx is lower as token count and
        # the current token string is in string from time stamp interval.
        while (idx < len(pair_split_word)) and (pair_split_word[idx] in word):

            # Get the length of the token string.
            token_word_len = len(pair_split_word[idx])

            # Check if the token string is at the end of the time stamp interval string.
            # This is important for tokens like 'a' and 'i'.
            if word_buffer[:token_word_len] == pair_split_word[idx]:

                # Add token to buffer and remove the string length of the token from the time stamp interval string.
                buffer.append(tokens[idx])
                word_buffer = word_buffer[token_word_len:]
                idx += 1
            else:
                break

        # Append buffer list to output list.
        bpe_list.append(buffer)

    return bpe_list


def char_level_encoder(words: list):

    chars = []

    for word in words:
        w = word
        chars.append([ord(c) for c in w])

    return chars


def worker(work: list,
           worker_id: int,
           save_dir: str,
           target_fps: int = 0,
           target_resolution: tuple = None,
           download: bool = True,
           lock = None,
           keep_buffer_download: bool = False,
           download_buffer_dir: str = '',
           use_subtitles: bool = False,
           tokenizer=None,
           language_tokens_per_frame: int = 4,
           padding_token: int = 50257,
           concat_token: int = 50256,
           skip_if_no_subtitles: bool = True,
           service_account_json_path: str = '',
           webshare_io_key: str = None,
           bucket_name: str = '',
           youtube_base: str = 'https://www.youtube.com/watch?v='):
    '''
    :param work: List with path to existing videos (if so download need to be True (default))
    or list with youtube video IDs (if so download need to be False).
    :param worker_id: The ID of the worker.
    :param save_dir: Directory where the finished TF.record's are saved or google blob to upload TF.record.
    :param target_fps: The fps from the TF.record, if 0 original download fps will be kept (default).
    :param target_resolution: Tuple with (width, height) resolution.
    If None original download resolution will be kept (default).
    :param download: Bool if it needs to download the video or just proses it. If download=True (default) youtube ID's
    needs to be given and if download=False path to videos needs to be given.
    :param lock: The multiprocessing lock used to control the youtube web site scraping.
    :param keep_buffer_download: If True the downloaded files will not be deleted after the tfrecord got created.
    :param download_buffer_dir: Directory where YoutubDL will download the videos (only if download is True (default)).
    It is recommended to use a RAM Disk as buffer directory.
    :param use_subtitles: Bool, if true Text will be used to.
    :param tokenizer: A lambda that contains the tokenizer.
    :param language_tokens_per_frame: The number of language tokens encoded per single frame.
    :param padding_token: The token id that is used as padding token.
    :param skip_if_no_subtitles: If True the video will be skipped if no subtitles are available.
    (only if use_subtitles is True)
    :param service_account_json_path: The path to the json containing the service account informations.
    :param webshare_io_key: API key for webshare.io.
    :param bucket_name: The Name of the google cloud storage bucket the TFrecords are should to be uploaded to.
    :param youtube_base: Youtube base string https://www.youtube.com/watch?v=.

    This function will download youtube videos and proses them and save than as TF.record files.
    '''

    # Check if TF.record are uploaded to google cloud storage.
    if service_account_json_path:
        cloud_storage = True
        buffer_save_dir = download_buffer_dir

        cloud_storage_client = storage.Client.from_service_account_json(service_account_json_path)
        cloud_storage_bucket = cloud_storage_client.get_bucket(bucket_name)
    else:
        cloud_storage = False
        buffer_save_dir = save_dir

    # shuffle the work for better thread utilization.
    random.shuffle(work)

    # really small padding frame.
    pading_frame = cv2.imencode('.jpg', np.array([[[255]]]))[1].tobytes()

    # Check if video needs to be downloaded.
    if download:
        # Creat Youtube Downloader.
        youtube_getter = youtube_dl.YoutubeDL({'writeautomaticsub': True, 'ignore-errors': True, 'socket-timeout': 600})
        youtube_getter.add_default_info_extractors()

        downloader = Downloader(webshare_io_key=webshare_io_key)

    # Loop to list of work.
    for chunk_idx, wor in enumerate(work):

        _tfrecord_name = "|".join(wor)
        _tfrecord_name = hashlib.sha3_256(_tfrecord_name.encode('utf-8')).hexdigest()
        _save_name = os.path.join(buffer_save_dir, _tfrecord_name + '.tfrecord')

        contains_video_already = False

        # Create TF Record Writer
        with tf.io.TFRecordWriter(_save_name) as tf_writer:

            for wor_idx, _wor in enumerate(wor):

                print(f"worker: {worker_id} chunk: {chunk_idx} video: {wor_idx}")

                # Assume by default the download was successful.
                download_success = True

                # Assume by default the subtitles are available.
                subtitles_available = True

                video_buffer_path = ""
                video_urls = ""
                vtt_url = ""

                # Check if video needs to be downloaded.
                if download:

                    # Execute all in try except to so the script is not crashing from a single fails video.
                    try:

                        print(youtube_base + _wor)

                        # We have to lock this part because it can lead to errors if multiple thread try to
                        # scrap video Information at the same time.
                        with lock:
                            # Get video info.
                            info = youtube_getter.extract_info(youtube_base + _wor, download=False)

                        # Go and find all video links that is as small as possible
                        # and is still larger that the target resolution.
                        if 'formats' in info:
                            current_res = (9999999, 9999999)

                            for f in info['formats']:

                                if 'format_note' in f:
                                    if f['format_note'] != "tiny":
                                        if 'width' in f and 'height' in f:
                                            width = f['width']
                                            height = f['height']

                                            if width is not None and height is not None:
                                                if width > target_resolution[0] and height > target_resolution[1]:
                                                    if current_res[0] > width and current_res[1] > height:
                                                        video_urls = []
                                                        current_res = (width, height)

                                                    if current_res[0] == width and current_res[1] == height:
                                                        if 'ext' in f and 'url' in f:
                                                            video_urls.append(
                                                                {'width': width, 'height': height,
                                                                 'ext': f['ext'], 'url': f['url']})

                        # Go and find the english vtt subtitle link.
                        if 'automatic_captions' in info:
                            automatic_captions = info['automatic_captions']
                            if 'en' in automatic_captions:
                                en_automatic_captions = automatic_captions['en']

                                for en in en_automatic_captions:
                                    if 'ext' in en:
                                        if en['ext'] == 'vtt':
                                            if 'url' in en:
                                                vtt_url = en['url']

                                            break

                    except:
                        download_success = False
                        warnings.warn("worker: " + str(worker_id) + " failed to get info: " + _wor)

                    # Do some checking to ensure that the information get extracted successfully.
                    if video_urls is None:
                        download_success = False

                    if not download_success and len(video_urls) == 0:
                        download_success = False

                    # Download the video and subtitle.
                    if download_success:

                        # Put .webm at the bottom at the list.
                        for idx in range(len(video_urls)):
                            if video_urls[idx]['ext'] == 'webm':
                                video_urls[-1], video_urls[idx] = video_urls[idx], video_urls[-1]

                        for video_url_idx, video_url in enumerate(video_urls):
                            url = video_url['url']
                            ext = video_url['ext']

                            if url is not None and ext is not None:
                                if url != "" and ext != "":
                                    video_buffer_path = os.path.join(download_buffer_dir, _wor) + '.' + ext
                                    download_success = downloader.download(url, video_buffer_path, False)

                                    if download_success:

                                        # If no mp4 god downloaded yous ffmpag to converted it to mp4.
                                        if ext != 'mp4':
                                            new_video_buffer_path = os.path.join(download_buffer_dir, _wor) + '.mp4'

                                            subprocess.run(['ffmpeg', '-i', video_buffer_path, '-c',
                                                            'copy', new_video_buffer_path, '-y'],
                                                           capture_output=False, stdout=subprocess.DEVNULL,
                                                           stderr=subprocess.STDOUT)

                                            if os.path.exists(video_buffer_path):
                                                os.remove(video_buffer_path)

                                            video_buffer_path = new_video_buffer_path

                                        # Check if the file can be opened.
                                        try:
                                            video_cap = cv2.VideoCapture(video_buffer_path)
                                            success, _ = video_cap.read()
                                            video_cap.release()
                                        except:
                                            success = False

                                        if success:
                                            break
                                        else:
                                            warnings.warn("worker: " + str(worker_id) + " cv2 failed to open:" +
                                                          video_buffer_path + " " +
                                                          str([__ff["ext"] for __ff in video_urls]))

                                            if os.path.exists(video_buffer_path):
                                                os.remove(video_buffer_path)

                                    print(f"Video download error video_url_idx: {video_url_idx} type: {ext}")



                else:
                    vtt = os.path.splitext(_wor)[0]
                    vtt = os.path.join(download_buffer_dir, (vtt + '.*'))
                    vtt = [v for v in glob.glob(vtt) if '.vtt' in v]

                    if len(vtt) > 0:
                        bpe_list = tokenizer(words, text)
                    else:
                        subtitles_available = False
                        bpe_list = []


                # Download and tokenize teh vtt.
                if use_subtitles and download_success and download:

                    vtt_download_success = False
                    if vtt_url is not None:
                        if vtt_url != "":
                            vtt = os.path.join(download_buffer_dir, _wor) + '.vtt'
                            vtt_download_success = downloader.download(vtt_url, vtt, True)

                    if vtt_download_success:
                        try:
                            with open(vtt, "r") as r:
                                text, words, stamp = decode_vtt(r.read())
                        except:
                            warnings.warn("worker: " + str(worker_id) + " Failed to open and decode: " + vtt)
                            vtt_download_success = False

                    if vtt_download_success:
                        bpe_list = tokenizer(words, text)

                    else:
                        subtitles_available = False
                        bpe_list = []

                _wor = video_buffer_path
                print(_wor, 'subtitles_available:', subtitles_available)


                # Check if this video need to be skipt. This will happen if subtitles are required but are not available.
                if (not use_subtitles or \
                   (subtitles_available or not skip_if_no_subtitles) or \
                   (subtitles_available and skip_if_no_subtitles)) and download_success:

                    try:
                        # Setup CV2 video reader.
                        video_cap = cv2.VideoCapture(_wor)
                        success, frame = video_cap.read()
                    except:
                        success = False

                    if success:

                        # If TFrecord already contains a video, add a separator between it and the new video.
                        if contains_video_already:
                            if use_subtitles:
                                tf_writer.write(frame_encoder(frame=pading_frame,
                                                              text_tokens=[concat_token] * language_tokens_per_frame,
                                                              skip_frame=[False],
                                                              mask=[language_tokens_per_frame],
                                                              concat=[True]))
                            else:
                                tf_writer.write(frame_encoder(frame=pading_frame, concat=[True]))


                        # Set frame counter zo zero.
                        frame_count = 0

                        # Get base video FPS, this is needed to align language and video.
                        fps_raw = video_cap.get(cv2.CAP_PROP_FPS)

                        # Calculate Modulo number.
                        fps_split = diveision_zero(round(fps_raw), target_fps)


                        while success:

                            # Check if frame needs to be saved.
                            if frame_count % fps_split == 0:

                                # Do resizing if required.
                                if target_resolution is not None:
                                    frame = cv2.resize(frame, target_resolution)

                                # convert BGR to RGB frames.
                                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

                                # Encode frame to image string.
                                frame = cv2.imencode('.jpg', frame)[1].tobytes()

                                if use_subtitles:
                                    token_buffer = []
                                    proto = []

                                    while len(bpe_list) > 0 and stamp[0] < (frame_count + fps_split) * (1 / fps_raw):
                                        token_buffer += bpe_list[0]
                                        bpe_list.pop(0)
                                        stamp.pop(0)

                                    for i in range(0, len(token_buffer), (language_tokens_per_frame - 1)):
                                        buffer = token_buffer[i:i + (language_tokens_per_frame - 1)]
                                        mask = len(buffer)
                                        buffer += [padding_token] * (language_tokens_per_frame - mask)
                                        skip_buffer = i > 0

                                        proto.append(frame_encoder(pading_frame if skip_buffer else frame,
                                                                   buffer,
                                                                   [skip_buffer], [mask]))

                                    if len(proto) <= 0:
                                        proto.append(frame_encoder(frame,
                                                                   [padding_token] * language_tokens_per_frame,
                                                                   [False], [0]))

                                else:
                                    # Encode frame to proto buffer.
                                    proto = [frame_encoder(frame)]

                                # print(frame_count, len(proto))

                                # Write proto buffer to TF.record file.
                                for p in proto:
                                    tf_writer.write(p)

                            # Get next frame.
                            success, frame = video_cap.read()
                            frame_count += 1

                    # close video file.
                    video_cap.release()
                    success = True
                    contains_video_already = True

                # Remove video file if it was downloaded.
                if not keep_buffer_download:
                    if os.path.exists(_wor):
                        os.remove(_wor)

                    if use_subtitles and subtitles_available:
                        if os.path.exists(vtt):
                            os.remove(vtt)

        if not contains_video_already:
            if os.path.exists(_save_name):
                os.remove(_save_name)

        if cloud_storage and success and contains_video_already:
            blob = cloud_storage_bucket.blob(os.path.join(save_dir, _tfrecord_name + '.tfrecord'))
            if os.path.exists(_save_name):
                blob.upload_from_filename(_save_name)

        if cloud_storage and not keep_buffer_download:
            if os.path.exists(_save_name):
                os.remove(_save_name)

    print('DONE worker:', worker_id)



if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument('load_path', type=str,
                        help='The path to a json file containing video information, or a path to a folder containing '
                             'json files with video information.')
    parser.add_argument('save_path', type=str,
                        help='The path where the final TFrecords get saved.')
    parser.add_argument('download_buffer_path', type=str,
                        help="A Folder that gets used to save downloads. IF 'keep_buffer_download' is False (default) "
                             "the content in the folder will get automatically deleted when it is no longer needed.")

    parser.add_argument('-num_worker', type=int, default=1,
                        help='The number of parallel workers.')
    parser.add_argument('-download', type=str, default='True',
                        help='If the videos are supposed to be downloaded from youtube or if there should '
                             'be loaded from the buffer folder.')
    parser.add_argument('-use_subtitles', type=str, default='False',
                        help='If subtitles are supposed the be downloaded and encoded.')
    parser.add_argument('-skip_if_no_subtitles', type=str, default='True',
                        help='Skip the Video if no subtitles are available. If False the ')
    parser.add_argument('-keep_buffer_download', type=str, default='False',
                        help='If True the videos how god downloaded will be capt on disk.')

    parser.add_argument('-target_fps', type=int, default=1,
                        help='The FPS of the output TFrecord.')
    parser.add_argument('-target_width', type=int, default=320,
                        help='The frame width of the output TFrecord.')
    parser.add_argument('-target_height', type=int, default=176,
                        help='The frame height of the output TFrecord.')
    parser.add_argument('-language_tokens_per_frame', type=int, default=64,
                        help='The number of tokens aligned with a single frame.')
    parser.add_argument('-padding_token', type=int, default=0,
                        help="The token ID that gets usd to pad to the 'language_tokens_per_frame' number"
                             "if not enough tokens are present in the frame.")
    parser.add_argument('-concat_token', type=int, default=4,
                        help="The token ID that gets usd to concat multibel videos")
    parser.add_argument('-duration_need_larger', type=int, default=256,
                        help='A single Video needs to be LONGER than this variable in seconds.')

    parser.add_argument('--service_account_json_path', type=str,
                        help="The path to the google service account information. If this is sat the TFrecords "
                             "will get automatically uploaded to a google cloud storage bucket."
                             " The bucket name needs to be defined with '--bucket_name'")
    parser.add_argument('--bucket_name', type=str,
                        help="The Name of the google cloud storage bucket the TFrecords are should to be uploaded to."
                             " This only as an affect if 'google service account information' are present.")
    parser.add_argument('--webshare_io_key', type=str,
                        help="The API key for webshare.io if this param is not set no proxy will be used.")
    parser.add_argument('--start_delay', type=int, default=0,
                        help="The delay in second between the start of the next worker.")


    args = parser.parse_args()

    load_path = args.load_path
    save_path = args.save_path
    download_buffer_path = args.download_buffer_path

    num_worker = args.num_worker
    download = str2bool(args.download)
    use_subtitles = str2bool(args.use_subtitles)
    skip_if_no_subtitles = str2bool(args.skip_if_no_subtitles)
    keep_buffer_download = str2bool(args.keep_buffer_download)

    target_fps = args.target_fps
    target_width = args.target_width
    target_height = args.target_height
    language_tokens_per_frame = args.language_tokens_per_frame
    padding_token = args.padding_token
    concat_token = args.concat_token
    duration_need_larger = args.duration_need_larger

    webshare_io_key = args.webshare_io_key
    service_account_json_path = args.service_account_json_path
    bucket_name = args.bucket_name
    start_delay = args.start_delay

    assert os.path.exists(load_path), 'The load path is invalid.'
    assert download, 'Only download is supported at the moment.'

    if not os.path.exists(download_buffer_path):
        os.makedirs(download_buffer_path)

    if not os.path.exists(save_path) and service_account_json_path == '':
        os.makedirs(save_path)

    if os.path.isdir(load_path):
        load_path = [os.path.join(load_path, p) for p in os.listdir(load_path)]
    else:
        load_path = [load_path]

    ids = []
    duration = []

    for l in load_path:
        json_load = json.load(open(l))

        ids = ids + json_load['id']
        duration = duration + json_load['duration']

    if type(ids[0]) is not list:
        ids = [[id] for id in ids]
    else:
        duration = [np.sum(_duration) for _duration in duration]

    ids, duration = split_equal(ids, duration, num_worker, duration_need_larger)

    split_chunk_count = 0
    split_video_count = 0
    split_video_duration = 0

    for i in range(len(ids)):
        buffer_chunk_count = len(ids[i])
        buffer_video_count = sum([len(__ids) for __ids in ids[i]])
        buffer_video_duration = sum(duration[i])

        print('split:', i, 'chunks:', buffer_chunk_count, 'videos:',
              buffer_video_count, 'duration:', buffer_video_duration)

        split_chunk_count += buffer_chunk_count
        split_video_count += buffer_video_count
        split_video_duration += buffer_video_duration

    print('')
    print('total num of chunks:', split_chunk_count, 'total num of videos:',
          split_video_count, 'total video duration:', split_video_duration)
    print('')


    youtube_getter = youtube_dl.YoutubeDL({'writeautomaticsub': True, 'ignore-errors': True, 'socket-timeout': 600})
    youtube_getter.add_default_info_extractors()

    lock = multiprocessing.Lock()

    worker_list = []

    for i, work in enumerate(ids):

        p = multiprocessing.Process(target=worker, args=(work,
                                                         i,
                                                         save_path,
                                                         target_fps,
                                                         (target_width, target_height),
                                                         download,
                                                         lock,
                                                         keep_buffer_download,
                                                         download_buffer_path,
                                                         use_subtitles,
                                                         lambda words, text: char_level_encoder(words),
                                                         language_tokens_per_frame,
                                                         padding_token,
                                                         concat_token,
                                                         skip_if_no_subtitles,
                                                         service_account_json_path,
                                                         webshare_io_key,
                                                         bucket_name))

        p.start()
        print('START WORKER', i)
        worker_list.append(p)

        time.sleep(start_delay)

    print(f"worker_list len: {len(worker_list)} ids len: {len(ids)}")

    for w in worker_list:
        w.join()

    #if download and not keep_buffer_download:
    #    os.remove(download_buffer_path)
