import os
import sys
import time
import subprocess
import argparse
import shutil
import re
from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv())

clean_re = re.compile('[^a-zA-Z0-9_\- ]')

start_time = time.time()

parser = argparse.ArgumentParser()
parser.add_argument('--folder', required=True)
parser.add_argument('--folder_id', type=int, default=0)
parser.add_argument('--out_fn', required=True)
parser.add_argument('--dataset', required=True)
parser.add_argument('--end', default="99999")
opt = parser.parse_args()

dataset = clean_re.sub('', opt.dataset)
step = '256'

os.makedirs('./uprez', exist_ok=True)

print(">>> {}".format(dataset))

in_dir = os.path.abspath(opt.folder)
vsr_duf_inputs_dir = os.path.abspath('../vsr-duf/inputs')
vsr_duf_results_dir = os.path.abspath('../vsr-duf/results/{}'.format(dataset))
out_dir = os.path.abspath('./uprez')

print("{}: {}".format('in_dir', in_dir))
print("{}: {}".format('vsr_duf_inputs', vsr_duf_inputs_dir))
print("{}: {}".format('vsr_duf_results', vsr_duf_results_dir))
print("{}: {}".format('out_dir', out_dir))

subprocess.call([ 'rm', '-rf', os.path.join(in_dir, '*') ])
subprocess.call([ 'rm', '-rf', os.path.join(vsr_duf_results_dir, '*') ])

print(">>>> crop frames")
subprocess.call([
  sys.executable,
  '../vsr-duf/crop_frames.py',
  '--step', step,
  '--in_dir', in_dir,
  '--out_dir', vsr_duf_inputs_dir,
  '--end', opt.end,
  '--ratio', '6.0'
])

print(">>>> uprez")
subprocess.call([
  sys.executable,
  '../vsr-duf/test.py',
  '--in_dir', vsr_duf_inputs_dir,
  '--out_dir', vsr_duf_results_dir,
  '--network_dir', os.path.abspath('../vsr-duf'),
])

print(">>>> merge framez")
subprocess.call([
  sys.executable,
  '../vsr-duf/merge_frames.py',
  '--step', step,
  '--in_dir', vsr_duf_results_dir,
  '--render_dir', out_dir,
  '--out_fn', opt.out_fn,
  '--mov',
])

if opt.folder_id > 0:
  mp4 = os.path.join(out_dir, opt.out_fn + '.mp4')
  print("upload {}".format(mp4))
  endpoint = os.getenv('API_REMOTE') + '/api/folder/{}/upload/'.format(opt.folder_id)
  subprocess.call([
    "curl",
    "-X", "POST",
    "-F", "module=pix2pixhd",
    "-F", "activity=uprez",
    "-F", "folder_id={}".format(opt.folder_id),
    "-F", "generated=1",
    "-F", "dataset={}".format(dataset),
    "-F", "datatype=video",
    "-F", "file=\@{}".format(mp4),
    endpoint
  ])

finish_time = time.time()

print("took %02f s.".format(finish_time - start_time))