#!/usr/bin/env python3
import pathlib
import argparse
import sys
import shutil
import exman
import os

parser = argparse.ArgumentParser()
commands = parser.add_subparsers(title="commands", help="Experiment manager CLI")

mark = commands.add_parser("mark")


def key_validator(key):
    try:
        path = pathlib.Path(key)
    except (TypeError, NotImplementedError):
        raise argparse.ArgumentError(
            'Mark "{}" should be a valid relative path'.format(key)
        )
    else:
        if path.is_absolute():
            raise argparse.ArgumentError(
                'Mark "{}" should be a valid relative path'.format(key)
            )
    return path


mark.add_argument("key", type=key_validator, help="key for mark")


class Mark(argparse.Action):
    def __init__(self, option_strings, dest=argparse.SUPPRESS, nargs="+", **kwargs):
        super().__init__(option_strings, dest=dest, nargs=nargs, type=int, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        selected = set(values)
        dest = pathlib.Path("marked") / namespace.key
        for run in pathlib.Path("runs").iterdir():
            ind = int(run.name.split("-", 1)[0])
            if ind in selected:
                os.makedirs(dest, exist_ok=True)
                rel_param_symlink = (
                    pathlib.Path("..", *([".."] * len(namespace.key.parts)))
                    / "runs"
                    / run.name
                )
                (dest / run.name).symlink_to(
                    rel_param_symlink, target_is_directory=True
                )
                selected.remove(ind)
                print("Created symlink from", dest / run.name, "->", rel_param_symlink)
        if selected:
            sys.stderr.write("warning: runs {} were not found\n".format(selected))
        parser.exit(0)


mark.add_argument("runs", help="runs to mark", action=Mark)

delete = commands.add_parser("delete")
delete.add_argument(
    "--all", action="store_true", help="Delete all associated files too"
)


class Delete(argparse.Action):
    def __init__(self, option_strings, dest=argparse.SUPPRESS, nargs="+", **kwargs):
        super().__init__(option_strings, dest=dest, nargs=nargs, type=int, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        selected = set(values)
        for index in pathlib.Path("index").iterdir():
            ind = int(index.name.split("-", 1)[0])
            if ind in selected:
                index.unlink()
                if not namespace.all:
                    selected.remove(ind)
        if namespace.all:
            for run in pathlib.Path("runs").iterdir():
                ind = int(run.name.split("-", 1)[0])
                if ind in selected:
                    shutil.rmtree(run, ignore_errors=True)
                    selected.remove(ind)
        if selected:
            sys.stderr.write("warning: runs {} were not found\n".format(selected))
        parser.exit(0)


delete.add_argument("runs", action=Delete, help="runs to delete")
if __name__ == "__main__":
    exman.parser.ExmanDirectory(".", mode="validate")
    parser.parse_args()
