from math import floor
from sys import float_info, stdout
from internetarchive import download
from moviepy.editor import VideoFileClip, concatenate_videoclips, vfx
from moviepy.video.tools.cuts import find_video_period
from random import choice, sample
import time
from argparse import ArgumentParser
import os
import pytesseract
import pickle
import shutil

options = ArgumentParser(prog="Wingnuts Video Mixer")
options.add_argument("identifiers", nargs="+", help="Internet Archive identifier(s)")
options.add_argument("-l", "--length", default=30, type=int, help="Length of output file (seconds)")
options.add_argument("-m", "--min", default=1.5, type=float, help="Minimum length per clip (seconds)")
options.add_argument("-M", "--max", default=float_info.max, type=float, help="Maximum length per clip (seconds)")
options.add_argument("-w", "--window", default=5, type=int, help="Search window length (seconds)")
options.add_argument("-t", "--text", default=False, type=bool, help="Allow clips with text")
options.add_argument("-u", "--unbalanced", default=False, type=bool, help="Clip pool size proportional to video length")
options.add_argument("-W", "--width", default=1280, type=int, help="Output video width")
options.add_argument("-H", "--height", default=720, type=int, help="Output video height")
options.add_argument("-P", "--process_only", default=False, type=bool, help="Don't output a video, just download and process")
args = options.parse_args()

clip_settings = dict(min=args.min, max=args.max, window=args.window)

print("Starting Wingnuts Video Mixer.")

def process_clips(identifier):
  clips = []
  
  source_directory = "./source/{}".format(identifier)
  clips_directory = "./clips/{}".format(identifier)

  os.mkdir(clips_directory)

  settings_file = open("{}/.settings".format(clips_directory), "wb")
  pickle.dump(clip_settings, settings_file)
  settings_file.close()

  source = VideoFileClip("{}/{}".format(source_directory, os.listdir(source_directory)[0]), target_resolution=(args.height, args.width))

  sample_start = 0
  clip_index = 0

  print("Processing {} ({})".format(source.filename, source.duration))
  while sample_start + args.window < source.duration:
    subclip = source.subclip(sample_start, sample_start + args.window)
    video_period = find_video_period(subclip)
    if video_period >= args.min and video_period <= args.max:
      print("Found candidate at ({}, {}) with length {}, appending.".format(sample_start, sample_start + args.window, video_period))
      subclip.write_videofile("{}/{}.mp4".format(clips_directory, clip_index), audio=False)
      clips.append(subclip.subclip(0.3, video_period))
      sample_start += video_period
      clip_index += 1
    else:
      print("Skipping subclip at ({}, {})".format(sample_start, sample_start + args.window))
      sample_start += args.window / 2
  return clips

def download_and_process_clips(identifier):
  download(identifier, formats=("MPEG2","MPEG4"), verbose=True, destdir="source")
  return process_clips(identifier)

all_clips = []

for identifier in args.identifiers:
  clips_directory = "./clips/{}".format(identifier)
  if os.path.isdir(clips_directory):
    try:
      settings_file = open("{}/.settings".format(clips_directory), "rb")
      saved_settings = pickle.load(settings_file)
      settings_file.close()
      if not saved_settings == clip_settings:
        print("Clip settings mismatch, reprocessing.")
        shutil.rmtree(clips_directory)
        all_clips.append(download_and_process_clips(identifier))
      else:
        print("Clips exist and settings are unchanged.")
        clips = []
        for file in os.listdir(clips_directory):
          clips.append(VideoFileClip("{}/{}".format(clips_directory, file), target_resolution=(args.height, args.width)))
        all_clips.append(clips)
    except:
      print("Could not open saved settings. Reprocessing.")
      shutil.rmtree(clips_directory)
      all_clips.append(download_and_process_clips(identifier))
  else:
    print("No saved clips found.")
    all_clips.append(download_and_process_clips(identifier))

if (args.process_only):
  pass
else:
  final_clips = []
  sampled_clips = []
  total_duration = 0

  if not args.unbalanced:
    max_length = min(map(lambda x: len(x), all_clips))
    for clips in all_clips:
      sampled_clips = sampled_clips + sample(clips, max_length)
    all_clips = sampled_clips
  else:
    for clips in all_clips:
      final_clips = final_clips + clips

  print("Remixing clips...")

  while total_duration < args.length:
    clip = choice(all_clips)
    if not args.text and len(pytesseract.image_to_string(clip.get_frame(0)).strip()):
      print("Text detected, skipping clip.")
      continue
    final_clips.append(clip)
    total_duration += clip.duration
    print("Adding clip ({} seconds so far).".format(total_duration))

  concatenate_videoclips(final_clips).write_videofile("./output/{}_{}.mp4".format(identifier, floor(time.time())), audio=False)

print("Done.")