#!/usr/bin/env python3

import argparse
import os
import platform
import socket
import subprocess
import sys
import tempfile
from multiprocessing import Process
from pathlib import Path
from threading import Thread

parser = argparse.ArgumentParser()
parser.add_argument('run_dirs', nargs='*')
args = parser.parse_args()


def get_open_port():
    s = socket.socket()
    s.bind(('', 0))
    port = s.getsockname()[1]
    s.close()
    return port


def communicate_process(p):
    while True:
        p.communicate()


class NoUniquePartsOfPath(Exception):
    pass


def get_unique_part_of_path(paths, path):
    n = -1
    while True:
        try:
            paths_field = set(Path(p).parts[n] for p in paths)
        except IndexError:
            raise NoUniquePartsOfPath
        if len(paths_field) == len(paths):
            break
        n -= 1
    return Path(path).parts[n]

port = get_open_port()
if len(args.run_dirs) == 1:
    if not os.path.exists(args.run_dirs[0]):
        print(f"Error: path '{args.run_dirs[0]}' does not exist",
              file=sys.stderr)
        exit(1)
    tensorboard_dir = args.run_dirs[0]
else:
    tensorboard_dir = tempfile.mkdtemp()
    for path in args.run_dirs:
        if not os.path.exists(path):
            print(f"Error: path '{path}' does not exist", file=sys.stderr)
            exit(1)
        if not os.path.isdir(path):
            print(f"Error: path '{path}' is not a directory", file=sys.stderr)
            exit(1)
        path = os.path.normpath(path)  # Strip trailing slashes
        try:
            name = get_unique_part_of_path(args.run_dirs, path)
        except NoUniquePartsOfPath:
            name = path
        link = os.path.join(tensorboard_dir, name)
        os.makedirs(os.path.dirname(link), exist_ok=True)
        os.symlink(os.path.abspath(path), link)

subprocess.run(['pkill', '-f', 'tensorboard --port {}'.format(port)])
tmp_dir = '/tmp/matthew'
os.makedirs(tmp_dir, exist_ok=True)
env = dict(os.environ)
env['TMPDIR'] = tmp_dir
proc1 = subprocess.Popen(['tensorboard',
                          '--samples_per_plugin', 'scalars=0',
                          '--port', str(port), '--logdir', tensorboard_dir],
                         stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                         env=env)
print("Waiting for TensorBoard to start...")
first_line = proc1.stderr.readline().decode()
if 'Traceback' in first_line:
    print(first_line + proc1.stderr.read().decode(), end='')
    exit(1)
proc2 = Process(target=communicate_process, args=(proc1,))
proc2.start()

if platform.system() == 'Darwin':
    subprocess.run(['open', '-na', 'Google Chrome', '--args', '--new-window', f'http://localhost:{port}'])
else:
    proc3 = subprocess.Popen(['h', str(port)])

try:
    input()
except KeyboardInterrupt:
    pass
proc1.terminate()
proc2.terminate()
