#!/usr/bin/env python3

import argparse
import subprocess
import sys
import time

PROJECT = "waystotrain"
IMAGE = PROJECT


def sh(cmd):
    if subprocess.call(cmd, shell=True):
        sys.stderr.write(f"[[./run]] Command failed: {cmd}\n")
        sys.exit(1)


def docker(cmd, *args, detach=False):
    mode = "-d" if detach else "--rm -it"
    docker_args = " ".join(args)
    sh(
        f"docker run {mode} --gpus all"
        f" -v `pwd`:/app -w /app -e PYTHONPATH=/app"
        f" {docker_args} {IMAGE} {cmd}"
    )


def build():
    sh(f"docker build --rm -t {IMAGE} .")


def run(args):
    docker(" ".join(args))


def test(filter):
    cov_or_filter = (
        "-s -k " + filter
        if filter
        else f"--cov={PROJECT} --no-cov-on-fail --cov-report term-missing"
    )
    docker(f"pytest {PROJECT} {cov_or_filter}")


def check():
    test(filter=None)
    docker(f"flake8 {PROJECT} run")
    docker(f"pylint --rcfile=setup.cfg --errors-only {PROJECT} run")


def notebook(port):
    name = f"$USER-{PROJECT}-notebook"
    sh(f"docker rm -f {name} || true")
    docker(
        f"jupyter notebook --allow-root --port {port} --ip 0.0.0.0",
        f"--name {name} -p {port}:{port} -e JUPYTER_ENABLE_LAB=true",
        detach=True,
    )

    print("Waiting for server to start, so we can tell you your token...")
    time.sleep(3)
    token = (
        subprocess.check_output(
            f"docker logs {name} 2>&1 | grep -oE '\\?token=.+$' | head -1", shell=True
        )
        .decode("utf8")
        .strip()
    )
    print(f"\nSee: http://localhost:{port}/lab{token}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=f"Development runner for {PROJECT}")
    parser.set_defaults(action=parser.print_help)
    subs = parser.add_subparsers()

    subs.add_parser("build", help="build the Docker image").set_defaults(action=build)

    p = subs.add_parser("run", help="run an arbitrary command in the image")
    p.add_argument("args", nargs="*", help="command/arguments to pass to the image")
    p.set_defaults(action=run)

    p = subs.add_parser("test", help="run unit tests")
    p.add_argument("-k", "--filter", help="filter tests matching this pattern")
    p.set_defaults(action=test)

    subs.add_parser("check", help="run all checks").set_defaults(action=check)

    p = subs.add_parser("notebook", help="run dev notebook")
    p.add_argument("--port", default="8888", help="port to serve on")
    p.set_defaults(action=notebook)

    args = vars(parser.parse_args())
    args.pop("action")(**args)
