#!/usr/bin/env python3

import itertools
import pathlib
import subprocess

import click
import yaml


class Variable:
    def __init__(self, name, values):
        self._name = name
        self._values = values

    @property
    def name(self):
        return self._name

    @property
    def values(self):
        return self._values

    @classmethod
    def new_scalar(cls, name, value):
        return cls(name, [value])

    @classmethod
    def new_multi_valued(cls, name, values):
        return cls(name, values)

    @classmethod
    def new(cls, name, value_config):
        if isinstance(value_config, list):
            return cls.new_multi_valued(name, value_config)
        elif not isinstance(value_config, dict):  # scalar value
            return cls.new_scalar(name, value_config)

        var_type = value_config["type"]
        assert var_type in ["file", "range"], "Unsupported variable type."

        if var_type == "file":
            directory = pathlib.Path(value_config["directory"]).expanduser()
            assert directory.exists() and directory.is_dir()

            values = sorted(list(directory.iterdir()))

            if value_config.get("basename", False):
                values = [v.name for v in values]

        if var_type == "range":
            assert "max" in value_config, f"Missing \"max\" value for range variable \"{name}\""
            upper = value_config["max"]
            lower = value_config.get("min", 0)
            step = value_config.get("step", 1)

            values = list(range(lower, upper + 1, step))

        return cls.new_multi_valued(name, values)


class Task:
    def __init__(self, name, work_dir, base_command, parameters):
        self._name = name
        self._work_dir = work_dir
        self._base_command = base_command
        self._parameters = parameters

    @property
    def name(self):
        return self._name

    @property
    def work_dir(self):
        return self._work_dir

    def run(self, variables):
        args = [self._base_command]

        for p in self._parameters:
            if not isinstance(p, dict):
                args.append(p.format(**variables))
                continue

            assert "value" in p, "An extended paramter has to contain a value field"

            if "if-variable" in p:
                v = p["if-variable"]

                if v not in variables.keys() or (isinstance(variables[v], bool) and not variables[v]):
                    continue  # skip

            if "if-not-variable" in p:
                v = p["if-not-variable"]

                if v in variables.keys() and (not isinstance(variables[v], bool) or variables[v]):
                    continue  # skip

            args.append(p["value"].format(**variables))

        # print()
        # print(f"pushd {self._work_dir}")
        # print(" ".join(args))
        # print("popd")

        try:
            subprocess.run(
                " ".join(args),
                shell = True,
                cwd = self._work_dir,
                timeout = 3600,  # FIXME
            )

            return False
        except subprocess.TimeoutExpired:
            print()
            print("Task timed out!")
            print()

            return True

    @classmethod
    def from_dict(cls, data):
        name = data.get("name", "Unnamed Task")

        assert "base-command" in data, f"No base command specified for task \"{name}\""
        base_command = data["base-command"]

        parameters = data.get("parameters", [])

        if "work-dir" in data:
            work_dir = pathlib.Path(data["work-dir"]).expanduser()
        else:
            work_dir = pathlib.Path.cwd()

        assert work_dir.exists(), f"Working directory \"{work_dir}\" for task \"{name}\" does not exist!"

        return cls(name, work_dir, base_command, parameters)


class Config:
    def __init__(self, name, repetitions, variables, tasks):
        self._name = name
        self._variables = variables
        self._repetitions = repetitions
        self._tasks = tasks

    @property
    def name(self):
        return self._name

    @property
    def repetitions(self):
        return self._repetitions

    @property
    def variables(self):
        return self._variables

    @property
    def tasks(self):
        return self._tasks

    @classmethod
    def from_file(cls, filepath: pathlib.Path):
        with filepath.open('r') as file:
            raw_config = yaml.safe_load(file)

        name = raw_config.get("name", str(filepath))
        repetitions = raw_config.get("repetitions", 1)
        variables = [Variable.new(k, v) for k, v in raw_config.get("variables", {}).items()]
        tasks = [Task.from_dict(t) for t in raw_config.get("tasks", [])]

        return cls(name, repetitions, variables, tasks)


@click.command()
@click.argument("config_file", type = click.Path(
    exists = True, file_okay = True, dir_okay = False, readable = True, path_type = pathlib.Path))
def main(config_file):
    c = Config.from_file(config_file)

    print(f"> Running experiment \"{c.name}\"")
    print(f">    with {len(c.variables)} variables")
    print(f">     and {len(c.tasks)} tasks")
    print(f">     and {c.repetitions} repetitions")
    print()

    variable_map = {v.name: v.values for v in c.variables}
    value_combinations = itertools.product(*variable_map.values())
    run_values = [dict(zip(variable_map.keys(), vals)) for vals in value_combinations]

    for variables in run_values:
        print()
        print("Running tasks with variable map: ")
        print(variables)
        print()

        for rep in range(c.repetitions):
            print(f"Repetition {rep}")
            timed_out = False

            for task in c.tasks:
                print(f"Running task {task.name}")
                timed_out = task.run(variables)
            print()

            if timed_out:  # repetitions are useless if tasks time out
                break


if __name__ == "__main__":
    main()