#! /usr/bin/env python

import argparse
from functools import partial
import json
from pathlib import Path
import subprocess
from typing import Optional


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("host", type=str)
    parser.add_argument("-v", "--verbose", action="store_true")
    parser.add_argument("--user", type=Optional[str])
    parser.add_argument("--rootdir", type=Optional[str])
    args = parser.parse_args()

    run = partial(subprocess.run, shell=True, check=True, text=True)
    host = args.host if args.user is None else f"{args.user}@{args.host}"

    def run_remote_cmd(cmd):
        return subprocess.check_output(f"ssh {host} {cmd}", shell=True, text=True)

    def ls(p):
        out = run_remote_cmd(f"ls {p}")
        return out.strip().split("\n")[::-1]

    def ask(l, info=None):
        print(
            "\n".join(
                [
                    f"{i:{len(str(len(l)))}d}: {d}"
                    + (f" ({info[d]})" if info is not None else "")
                    for i, d in enumerate(l, 1)
                ]
            )
        )
        while True:
            i = input("\nEnter a number: ")
            if i.isdigit() and 1 <= int(i) <= len(l):
                break
            print("\n/!\\ Invalid choice\n")
        return l[int(i) - 1]

    def ask_if_verbose(question, default):
        if not args.verbose:
            return default
        suffix = "[Y|n]" if default else "[y|N]"
        answer = input(f"{question} {suffix} ").lower()

        return (answer != "n") if default else (answer == "y")

    def get_info(rundir):
        return json.loads(
            run_remote_cmd(f"cat {rundir}/checkpoints/info_for_import_script.json")
        )

    if args.rootdir is None:
        for p in Path(__file__).resolve().parents:
            if (p / ".git").is_dir():
                break
        else:
            raise RuntimeError("This file is not in a git repository")
        out = run_remote_cmd(f"find -type d -name {p.name}").strip().split("\n")
        assert len(out) == 1
        rootdir = out[0]
    else:
        rootdir = f'{args.rootdir.strip().strip("/")}'

    dates = ls(f"{rootdir}/outputs")
    date = ask(dates)
    times = ls(f"{rootdir}/outputs/{date}")

    infos = {
        time: get_info(rundir=f"{rootdir}/outputs/{date}/{time}") for time in times
    }
    time = ask(times, infos)

    src = f"{rootdir}/outputs/{date}/{time}"

    dst = Path(args.host) / date
    dst.mkdir(exist_ok=True, parents=True)

    exclude = [
        "*.log",
        "checkpoints/*",
        "checkpoints_tmp",
        ".hydra",
        "media",
        "__pycache__",
        "wandb",
    ]

    include = ["checkpoints/agent_versions"]

    if ask_if_verbose("Download only last checkpoint?", default=True):
        last_ckpt = ls(f"{src}/checkpoints/agent_versions")[0]
        exclude.append("checkpoints/agent_versions/*")
        include.append(f"checkpoints/agent_versions/{last_ckpt}")

    if not ask_if_verbose("Download train dataset?", default=False):
        exclude.append("dataset/train")

    if not ask_if_verbose("Download test dataset?", default=False):
        exclude.append("dataset/test")

    cmd = "rsync -av"
    for i in include:
        cmd += f' --include="{i}"'
    for e in exclude:
        cmd += f' --exclude="{e}"'

    cmd += f" {host}:{src} {str(dst)}"
    run(cmd)

    path = (dst / time).absolute()
    print(f"\n--> Run imported in:\n{path}")
    run(f"echo {path} | xclip")


if __name__ == "__main__":
    main()