class Tool:
    def __init__(
        self,
        name: str,
        depends: Union[List[str], str] = [],
        load: Union[List[str], str] = "",
        unload: Union[List[str], str] = "",
        setup: Union[List[str], str] = "",
        options: Dict[str, any] = {},
    ):
        """
        A tool is a software that can be loaded/unloaded in the task script.
        The tool can be a module, a package, a virtual environment, etc.
        A tool provides three scripts:
        - load: script to load the tool
        - unload: script to unload the tool
        - setup: script to install the tool (in the deploy family)

        Parameters
        ----------
        name: str
            name of the tool
        depends: list of str
            list of tools that must be loaded before this tool
        load: str or list of str
            bash script to load the tool
        unload: str or list of str
            bash script to unload the tool
        setup: str or list of str
            bash script to setup the tool
        options: dict
            dictionary of options for the tool
        """
        self.name = name
        self.options = options
        self.depends = depends if isinstance(depends, list) else [depends]
        self.scripts = {}
        self.scripts["load"] = load
        self.scripts["unload"] = unload
        self.scripts["setup"] = setup

class ModuleTool(Tool):
    def __init__(
        self,
        name: str,
        module_name: str,
        version: str,
        depends: List[str] = [],
        options: Dict[str, any] = {},
    ):
        """
        A module tool is a tool that can be loaded/unloaded using the
        module system.

        Parameters
        ----------
        name: str
            name of the tool
        module_name: str
            name of the module
        version: str
            version of the module
        depends: list of str
            list of tools that must be loaded before this tool
        options: dict
            dictionary of options for the tool
        """
        load = pf.TemplateScript(
            module_load, NAME=module_name, VERSION=version
        )
        unload = pf.TemplateScript(module_unload, NAME=module_name)
        super().__init__(name, depends, load, unload, options=options)

class PrivateModuleTool(Tool):
    def __init__(
        self,
        name: str,
        module_name: str,
        version: str,
        modulefiles: str,
        depends: List[str] = [],
        options: Dict[str, any] = {},
    ):
        """
        Extension of the ModuleTool class to load a module from a custom
        modulefile path.

        Parameters
        ----------
        name: str
            name of the tool
        module_name: str
            name of the module
        version: str
            version of the module
        modulefiles: str
            path to the custom modulefile
        depends: list of str
            list of tools that must be loaded before this tool
        options: dict
            dictionary of options for the tool
        """
        load = pf.TemplateScript(
            module_use + module_load,
            NAME=module_name,
            VERSION=version,
            MODULEFILES=modulefiles,
        )
        unload = pf.TemplateScript(module_unload, NAME=module_name)
        super().__init__(name, depends, load, unload, options=options)

class EnvVarTool(Tool):
    def __init__(
        self,
        name: str,
        var: str,
        value: str,
        setup: Union[List[str], str] = None,
        depends: List[str] = [],
        options: Dict[str, any] = {},
    ):
        """
        Environment variable tool.
        When loading, it set the environment variable to the given value.
        If the value contains the string PATH (for instance PATH or
        PYTHONPATH), then the value is added to the environment variable.
        For instance:
        export PATH=/this/his/the/path:$PATH

        Parameters
        ----------
        name: str
            name of the tool
        var: str
            name of the envirpnment variable
        value: str
            path to set to the environment variable
        setup: str or list of str
            setup script used to deploy the tool
        depends: list of str
            list of tools that must be loaded before this tool
        options: dict
            dictionary of options for the tool
        """
        load_tplt = env_var_load
        unload_tplt = env_var_unload
        if "PATH" in var.upper():
            load_tplt = env_path_load
            unload_tplt = env_path_unload

        load = pf.TemplateScript(load_tplt, NAME=var, VALUE=value)
        unload = pf.TemplateScript(unload_tplt, NAME=var, VALUE=value)
        super().__init__(
            name, depends, load, unload, setup=setup, options=options
        )

class PackageTool(Tool):
    def __init__(self, name: str, lib_dir: str, options: Dict[str, any]):
        """
        Package tool that installs a user package in an environment.
        The package is downloaded in the lib_dir/build directory and the user
        must provide a script to install the package in the environment.

        Parameters
        ----------
        name: str
            name of the tool.
        lib_dir: str
            path to the lib directory.
        options: dict
            dictionary of options for the tool.
        """
        depends = options.get("depends", [])
        build_dir = options.get("build_dir")
        setup = deploy_package_script(name, options, lib_dir, build_dir)
        super().__init__(name, depends, setup=setup, options=options)

class FolderTool(EnvVarTool):
    def __init__(
        self,
        name: str,
        lib_dir: str,
        depends: List[str] = [],
        options: Dict[str, any] = {},
    ):
        """
        Folder tool extending the environment variable tool.
        It updates the PATH environment variable to include the folder located
        in lib_dir/name.

        Parameters
        ----------
        name: str
            name of the tool.
        lib_dir: str
            path to the lib directory of the suite.
        depends: list of str
            list of tools that must be loaded before this tool.
        options: dict
            dictionary of options for the tool.
        """
        folder_path = path.join(lib_dir, name)
        setup = [
            f"rm -rf {folder_path}",
            f"mkdir -p {folder_path}",
        ]
        super().__init__(
            name, "PATH", folder_path, setup, depends, options=options
        )

class VirtualEnvTool(Tool):
    def __init__(
        self,
        name: str,
        lib_dir: str,
        venv_options: Union[str, List[str]] = "",
        extra_packages: Union[str, List[str]] = [],
        depends: List[str] = [],
        options: Dict[str, any] = {},
    ):
        """
        A tool that creates a python virtual environment and sets the right
        scripts to load it, unload it and install it.

        Parameters
        ----------
        name : str
            The name of the tool.
        lib_dir : str
            The path to the lib directory.
        venv_options : Union[str, List[str]], optional
            Additional options to pass to the `python3 -m venv` command, by
            default "".
        depends : List[str], optional
            A list of dependencies for the tool, by default [].
        options : Dict[str], optional
            A dictionary of options for the tool, by default {}.
        """
        env_root = path.join(lib_dir, name)
        load = [
            "\nsource {}".format(path.join(env_root, "bin", "activate")),
            "export LD_LIBRARY_PATH={}:${{LD_LIBRARY_PATH:=}}".format(
                path.join(env_root, "lib")
            ),
        ]
        unload = "deactivate"
        if not isinstance(venv_options, list):
            venv_options = [venv_options]
        opts = " ".join(venv_options)
        setup = [
            f"rm -rf {env_root}",
            f"python3 -m venv {env_root} {opts}",
        ]
        if extra_packages:
            setup.extend(load)
            pkgs = " ".join([shlex.quote(pkg) for pkg in extra_packages])
            setup.append(f"pip install {pkgs}")
        super().__init__(name, depends, load, unload, setup, options=options)

class SystemEnvTool(VirtualEnvTool):
    def __init__(
        self,
        name: str,
        lib_dir: str,
        venv_options: Union[str, List[str]] = "",
        extra_packages: Union[str, List[str]] = [],
        depends: List[str] = [],
        options: Dict[str, any] = {},
    ):
        """
        An extension of the VirtualEnvTool class that creates a python virtual
        environment that extends the current environment, i.e. with the
        --system-site-packages option.

        Parameters
        ----------
        name : str
            The name of the tool.
        lib_dir : str
            The path to the lib directory.
        venv_options : Union[str, List[str]], optional
            Additional options to pass to the `python3 -m venv` command, by
            default "".
        depends : List[str], optional
            A list of dependencies for the tool, by default [].
        options : Dict[str], optional
            A dictionary of options for the tool, by default {}.
        """
        to_add = [" --system-site-packages"]
        if venv_options:
            if isinstance(venv_options, str):
                venv_options = [venv_options]
            to_add.extend(venv_options)

        super().__init__(
            name,
            lib_dir,
            venv_options=to_add,
            extra_packages=extra_packages,
            depends=depends,
            options=options,
        )

class CondaEnvTool(Tool):
    def __init__(
        self,
        name: str,
        environment: str,
        setup: str = None,
        depends: List[str] = [],
        conda_activate_cmd: str = "conda",
        options: Dict[str, str] = {},
    ):
        """
        An environment tool that can load and unload a conda environment.

        Parameters
        ----------
        name : str
            The name of the tool.
        environment : str
            The name or path of the conda environment.
        setup : str, optional
            The setup script, by default None.
        depends : List[str], optional
            A list of dependencies for the tool, by default [].
        conda_activate_cmd : str, optional
            The conda command to use to load the environment,
            some login-nodes require "source" instead of "conda"
            by default "conda".
        options : Dict[str, str], optional
            A dictionary of options for the tool, by default {}.
        """
        load = pf.TemplateScript(
            conda_load,
            TARGET=environment,
            CONDA_ACTIVATE_CMD=conda_activate_cmd,
        )
        unload = f"{conda_activate_cmd} deactivate"
        super().__init__(name, depends, load, unload, setup, options=options)

class FileCondaEnvTool(CondaEnvTool):
    def __init__(
        self,
        name: str,
        lib_dir: str,
        env_file: str,
        depends: List[str] = [],
        conda_cmd: str = "conda",
        conda_activate_cmd: str = "conda",
        options: Dict[str, any] = {},
    ):
        """
        An extension of the CondaEnvTool class that also creates a conda
        environment in the lib directory.
        The conda environment is created from a conda environment file.
        The path to the environment file is specified as a dictionary following
        the StaticData options format, for instance:
        ```yaml
        my_env:
            env_file:
                type: copy
                source: /path/to/environment.yml
        my_env2:
            env_file:
                type: git
                source: git@github.com:myrepo.git
                branch: master
                files: myenvironment.yml
        ```

        Parameters
        ----------
        name : str
            The name of the tool.
        lib_dir : str
            The path to the lib directory of the suite.
        env_file : str
            The options to retrieve the environment file.
        depends : List[str], optional
            A list of dependencies for the tool, by default [].
        conda_cmd : str, optional
            The conda command to use, by default "conda".
        options : Dict[str, str], optional
            A dictionary of options for the tool, by default {}.
        """
        env_root = path.join(lib_dir, name)
        build_dir = options.get("build_dir", path.join(lib_dir, "build"))
        fname = env_file.get("files", os.path.basename(env_file["source"]))
        file_setup = data.parse_data_item(build_dir, name, env_file)

        env_file_path = path.join(build_dir, name, fname)

        setup = pf.TemplateScript(
            conda_setup_yaml,
            ENV_FILE=env_file_path,
            ENV_DIR=env_root,
            CONDA_CMD=conda_cmd,
        )
        setup_script = [file_setup.script, setup]
        super().__init__(
            name,
            env_root,
            setup_script,
            depends,
            conda_activate_cmd,
            options=options,
        )

class SimpleCondaEnvTool(CondaEnvTool):
    def __init__(
        self,
        name: str,
        lib_dir: str,
        packages: Union[str, List[str]],
        depends: List[str] = [],
        conda_cmd: str = "conda",
        conda_activate_cmd: str = "conda",
        options: Dict[str, any] = {},
    ):
        """
        An extension of the CondaEnvTool class that also creates a conda
        environment in the lib directory.
        The conda environment is created from a list of conda packages.

        Parameters
        ----------
        name : str
            The name of the tool.
        lib_dir : str
            The path to the lib directory of the suite.
        packages : str or list of str
            The options to retrieve the environment file.
        depends : List[str], optional
            A list of dependencies for the tool, by default [].
        conda_cmd : str, optional
            The conda command to use, by default "conda".
        options : Dict[str, str], optional
            A dictionary of options for the tool, by default {}.
        """
        env_root = path.join(lib_dir, name)
        packages_str = " ".join([shlex.quote(pkg) for pkg in packages])
        setup = pf.TemplateScript(
            conda_create,
            ENV_DIR=env_root,
            PACKAGES=packages_str,
            CONDA_CMD=conda_cmd,
        )
        super().__init__(
            name, env_root, setup, depends, conda_activate_cmd, options=options
        )

class ToolStore:
    def __init__(self, lib_dir: str, options: Dict[str, any]):
        """
        The ToolStore class is a container for all the tools of a suite.
        The constructor parses the options and creates the tools.

        Parameters
        ----------
        lib_dir : str
            The path to the lib directory of the suite.
        options : Dict[str, str], optional
            A dictionary of options containing all the tools and their options.
        """
        if options is None:
            options = {}
        store_options = options.pop("options", {})
        self.modules = options.get("modules", {})
        self.packages = options.get("packages", {})
        self.environments = options.get("environments", {})
        self.env_vars = options.get("env_variables", {})
        self.dir = lib_dir
        self.list_modules = store_options.get("list_modules", True)

        # Build tools
        self.tools = {}

        # modules:
        for name, options in self.modules.items():
            module = parse_module(name, options)
            self.add_tool(name, module)

        # packages
        for name, options in self.packages.items():
            package = parse_package(lib_dir, name, options)
            self.add_tool(name, package)

        # environments
        for name, options in self.environments.items():
            env = parse_environment(lib_dir, name, options)
            self.add_tool(name, env)

        # environment variables
        for name, options in self.env_vars.items():
            var = parse_env_var(name, options)
            self.add_tool(name, var)

    def add_tool(self, toolname: str, tool: Tool):
        """
        Adds a tool to the tool store.

        Parameters
        ----------
        toolname : str
            The name of the tool.
        tool : Tool
            The tool object to add to the dictionary of tools.
        """
        self.tools[toolname] = tool

    def tool_script(self, script_type: str, toolname: str) -> dict:
        """
        Returns the script of a tool and its dependencies.

        Parameters
        ----------
        script_type : str
            The type of script (load, unload or setup).
        toolname : str
            The name of the tool.

        Returns
        -------
        dict: a dictionary of tools -> scripts
        """
        tool = self.tools[toolname]
        script = {}
        for item in tool.depends:
            script.update(self.tool_script(script_type, item))
        if tool.scripts[script_type]:
            script[toolname] = tool.scripts[script_type]
        return script

    def load(self, tools: Union[List[str], str]) -> list:
        """
        Returns the load script of one or multiple tools.

        Parameters
        ----------
        tools : str or list of str
            The tools to load.

        Returns
        -------
        list : a list of scripts to load the tools
        """
        if not isinstance(tools, list):
            tools = [tools]
        script = {}
        script["head"] = "# load tools and activate environment"
        for item in tools:
            script.update(self.tool_script("load", item))
        load_scripts = list(script.values())
        if self.list_modules:
            load_scripts.append("module list")
        return load_scripts

    def unload(self, tools: Union[List[str], str]):
        """
        Returns the unload script of one or multiple tools.
        In this case we reverse the dependency order to unload the tools.

        Parameters
        ----------
        tools : str or list of str
            The tools to unload.

        Returns
        -------
        list : a list of scripts to unload the tools
        """
        if not isinstance(tools, list):
            tools = [tools]
        script = {}
        script["head"] = "# unload tools and deactivate environment"
        for item in tools:
            # reverse dictionary to unload
            unload_script = self.tool_script("unload", item)
            unload_script = unload_script.fromkeys(
                reversed(list(unload_script.keys()))
            )
            script.update(unload_script)
        return list(script.values())

    def setup(self, toolname: str) -> list:
        """
        Returns the setup script of the tool.

        Parameters
        ----------
        toolname : str
            The tool to setup.

        Returns
        -------
        list : a list of scripts to setup the tool
        """
        script = {}
        script["head"] = "# load tools and activate environment"
        tool = self.tools[toolname]
        for item in tool.depends:
            script.update(self.tool_script("load", item))
        script[toolname] = tool.scripts["setup"]
        return list(script.values())

    def depends(self, toolname: str) -> list:
        """
        Returns the dependencies of a tool.

        Parameters
        ----------
        toolname : str
            The name of the tool.

        Returns
        -------
        list : the list of dependcies of the tool
        """
        tool = self.tools[toolname]
        depends = []
        for item in tool.depends:
            depends += [item]
            depends += self.depends(item)
        return depends

    def __getitem__(self, item):
        return self.tools[item]

    def items(self):
        return self.tools.items()