#!/usr/bin/env python3

import argparse
import datetime
import fnmatch
import glob
import io
import os
import platform
import random
import shutil
import subprocess
import sys
import tempfile
import time
import traceback
import uuid
from threading import Thread

import boto3
import daemon

session = boto3.Session(profile_name='runs_s3')
s3 = session.client('s3')
s3.list_objects(Bucket='mrahtz-runs-2')


def find_files_matching_pattern(pattern, path):
    result = []
    for root, dirs, files in os.walk(path, followlinks=True):
        for name in files:
            if fnmatch.fnmatch(name, pattern):
                result.append(os.path.join(root, name))
    return result


def split_preserving_seps(s, seps=('_', '-')):
    """
    'foo-bar_1' => ['foo', '-', 'bar', '_', '1']
    """
    arr = []
    cur = ""
    for c in s:
        if c in seps:
            if cur:
                arr.append(cur)
            arr.append(c)
            cur = ""
        else:
            cur += c
    arr.append(cur)
    return arr


def common_name(names):
    """
    ['foo-1-bar', 'foo-2-bar'] => 'foo-*-bar'
    """
    fields = [split_preserving_seps(name) for name in names]
    if len({len(field) for field in fields}) == 1:
        n_fields = len(fields[0])
        common_name = ""
        for n in range(n_fields):
            if len({field[n] for field in fields}) == 1:
                # This field the same in all names
                common_name += fields[0][n]
            else:
                common_name += '*'
    else:
        longest_common_prefix = os.path.commonprefix(names)
        longest_common_suffix = os.path.commonprefix([d[::-1] for d in names])[::-1]
        common_name = longest_common_prefix + '*' + longest_common_suffix
    return common_name


def monitor(paths):
    data_name = common_name(paths)
    while True:
        if any([not os.path.exists(p) for p in paths]):
            break

        try:
            make_plots(paths)
            data_files = find_videos(paths) + find_files(paths, 'events*') + find_files(paths, '*log')
            data_files = [d for d in data_files if not d.endswith('.tmp')]
            with tempfile.TemporaryDirectory() as d:
                d_subdir = os.path.join(d, data_name)
                os.makedirs(d_subdir)
                for f in data_files:
                    target_dir = os.path.join(d_subdir, os.path.dirname(f))
                    os.makedirs(target_dir, exist_ok=True)
                    shutil.copy2(f, target_dir)
                shutil.make_archive(data_name, 'xztar', d)

            print("Uploading...")
            for f in [data_name + '.png', data_name + '.tar.xz']:
                with open(f, 'rb') as fd:
                    b = fd.read()
                s3.upload_fileobj(io.BytesIO(b), 'mrahtz-runs-2', os.path.basename(f), ExtraArgs={'ACL': 'public-read'})

            print("Sleeping...")
            stagger_minutes = random.randint(1, 5)
            time.sleep((5 + stagger_minutes) * 60)
        except Exception:
            traceback.print_exc()
            time.sleep(60)

def find_videos(paths):
    videos_all_paths = []
    for path in paths:
        path = os.path.normpath(path)  # Strip trailing slash
        video_files = find_files_matching_pattern('*mp4', os.path.join(path, 'test_env'))
        video_files += find_files_matching_pattern('*mp4', os.path.join(path, 'train_env'))
        video_files = sorted(video_files, key=lambda f: os.path.getmtime(f))
        video_files = video_files[-20:-1]
        videos_all_paths.extend(video_files)
    return videos_all_paths


def find_files(paths, pattern):
    files_all_paths = []
    for path in paths:
        path = os.path.normpath(path)  # Strip trailing slash
        paths = find_files_matching_pattern(pattern, path)
        files_all_paths.extend(paths)
    return files_all_paths


def make_plots(paths):
    print("Plotting...")
    try:
        subprocess.run(['tbplot'] + paths, check=True)
    except subprocess.CalledProcessError as e:
        print(e)


def download_file(s3, filename, last_modified, target_dir):
    target_file = os.path.join(target_dir, filename)
    print(f"Downloading {filename}...")
    s3.download_file('mrahtz-runs-2', filename, target_file)
    os.utime(target_file, (last_modified, last_modified))


def download_files_from_bucket(get_data, time_limit_hours):
    tmp_dir = tempfile.mkdtemp()
    print("Listing bucket...")
    files = s3.list_objects(Bucket='mrahtz-runs-2')
    threads = []
    print("Downloading files...")
    for content in files['Contents']:
        filename = content['Key']
        if (not get_data) and '.tar.xz' in filename:
            continue
        last_modified = content['LastModified']
        last_modified = last_modified.replace(tzinfo=datetime.timezone.utc).timestamp()
        if time.time() - last_modified < (time_limit_hours * 60 * 60):
            t = Thread(target=download_file, args=(s3, filename, last_modified, tmp_dir))
            threads.append(t)
            t.start()
    for t in threads:
        t.join()
    return tmp_dir


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--foreground', action='store_true')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('--get', action='store_true')
    group.add_argument('--paths', nargs='+')
    parser.add_argument('--get_data', action='store_true')
    parser.add_argument('--time_limit_hours', type=float, default=2.)
    args = parser.parse_args()

    if args.get:
        tmp_dir = download_files_from_bucket(args.get_data, args.time_limit_hours)
        if args.get_data:
            for p in glob.glob(os.path.join(tmp_dir, '*.tar.xz')):
                print(f"Unpacking {p}...")
                shutil.unpack_archive(p, tmp_dir)
                os.remove(p)
        if platform.system() == 'Darwin':
            subprocess.run(['open', tmp_dir])
        else:
            print(f"Files downloaded to {tmp_dir}")
    elif args.paths:
        for p in args.paths:
            if not os.path.exists(p):
                print(f"Error: path '{p}' does not exist", file=sys.stderr)
                exit(1)
        dirs = [p for p in args.paths if os.path.isdir(p)]
        dirs = [os.path.normpath(d) for d in dirs]  # Strip trailing slashes
        if args.foreground:
            monitor(dirs)
        else:
            name = str(uuid.uuid4())
            log_file = open('/tmp/runmon-' + name + '.log', 'w')
            with daemon.DaemonContext(stdout=log_file, stderr=log_file, working_directory=os.getcwd()):
                monitor(dirs)
    else:
        parser.print_help()


if __name__ == '__main__':
    main()
