#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
#
"""
Scheduler-side Environment to run scripts locally
and upload/download data to the shared storage.
"""

import logging

from datetime import datetime
from string import Template
from typing import Any, Dict, List, Generator, Iterable, Mapping, Optional, Tuple

from mlos_bench.services.base_service import Service
from mlos_bench.services.types.local_exec_type import SupportsLocalExec
from mlos_bench.services.types.fileshare_type import SupportsFileShareOps
from mlos_bench.environments.status import Status
from mlos_bench.environments.local.local_env import LocalEnv
from mlos_bench.tunables.tunable import TunableValue
from mlos_bench.tunables.tunable_groups import TunableGroups

_LOG = logging.getLogger(__name__)


class LocalFileShareEnv(LocalEnv):
    """
    Scheduler-side Environment that runs scripts locally
    and uploads/downloads data to the shared file storage.
    """

    def __init__(self,
                 *,
                 name: str,
                 config: dict,
                 global_config: Optional[dict] = None,
                 tunables: Optional[TunableGroups] = None,
                 service: Optional[Service] = None):
        """
        Create a new application environment with a given config.

        Parameters
        ----------
        name: str
            Human-readable name of the environment.
        config : dict
            Free-format dictionary that contains the benchmark environment
            configuration. Each config must have at least the "tunable_params"
            and the "const_args" sections.
            `LocalFileShareEnv` must also have at least some of the following
            parameters: {setup, upload, run, download, teardown,
            dump_params_file, read_results_file}
        global_config : dict
            Free-format dictionary of global parameters (e.g., security credentials)
            to be mixed in into the "const_args" section of the local config.
        tunables : TunableGroups
            A collection of tunable parameters for *all* environments.
        service: Service
            An optional service object (e.g., providing methods to
            deploy or reboot a VM, etc.).
        """
        super().__init__(name=name, config=config, global_config=global_config, tunables=tunables, service=service)

        assert self._service is not None and isinstance(self._service, SupportsLocalExec), \
            "LocalEnv requires a service that supports local execution"
        self._local_exec_service: SupportsLocalExec = self._service

        assert self._service is not None and isinstance(self._service, SupportsFileShareOps), \
            "LocalEnv requires a service that supports file upload/download operations"
        self._file_share_service: SupportsFileShareOps = self._service

        self._upload = self._template_from_to("upload")
        self._download = self._template_from_to("download")

    def _template_from_to(self, config_key: str) -> List[Tuple[Template, Template]]:
        """
        Convert a list of {"from": "...", "to": "..."} to a list of pairs
        of string.Template objects so that we can plug in self._params into it later.
        """
        return [
            (Template(d['from']), Template(d['to']))
            for d in self.config.get(config_key, [])
        ]

    @staticmethod
    def _expand(from_to: Iterable[Tuple[Template, Template]],
                params: Mapping[str, TunableValue]) -> Generator[Tuple[str, str], None, None]:
        """
        Substitute $var parameters in from/to path templates.
        Return a generator of (str, str) pairs of paths.
        """
        return (
            (path_from.safe_substitute(params), path_to.safe_substitute(params))
            for (path_from, path_to) in from_to
        )

    def setup(self, tunables: TunableGroups, global_config: Optional[dict] = None) -> bool:
        """
        Run setup scripts locally and upload the scripts and data to the shared storage.

        Parameters
        ----------
        tunables : TunableGroups
            A collection of tunable OS and application parameters along with their
            values. In a local environment these could be used to prepare a config
            file on the scheduler prior to transferring it to the remote environment,
            for instance.
        global_config : dict
            Free-format dictionary of global parameters of the environment
            that are not used in the optimization process.

        Returns
        -------
        is_success : bool
            True if operation is successful, false otherwise.
        """
        self._is_ready = super().setup(tunables, global_config)
        if self._is_ready:
            assert self._temp_dir is not None
            params = self._get_env_params()
            params["PWD"] = self._temp_dir
            for (path_from, path_to) in self._expand(self._upload, params):
                self._file_share_service.upload(self._config_loader_service.resolve_path(
                    path_from, extra_paths=[self._temp_dir]), path_to)
        return self._is_ready

    def _download_files(self, ignore_missing: bool = False) -> None:
        """
        Download files from the shared storage.

        Parameters
        ----------
        ignore_missing : bool
            If True, raise an exception when some file cannot be downloaded.
            If False, proceed with downloading other files and log a warning.
        """
        assert self._temp_dir is not None
        params = self._get_env_params()
        params["PWD"] = self._temp_dir
        for (path_from, path_to) in self._expand(self._download, params):
            try:
                self._file_share_service.download(
                    path_from, self._config_loader_service.resolve_path(
                        path_to, extra_paths=[self._temp_dir]))
            except FileNotFoundError as ex:
                _LOG.warning("Cannot download: %s", path_from)
                if not ignore_missing:
                    raise ex

    def run(self) -> Tuple[Status, Optional[Dict[str, float]]]:
        """
        Download benchmark results from the shared storage
        and run post-processing scripts locally.

        Returns
        -------
        (status, output) : (Status, dict)
            A pair of (Status, output) values, where `output` is a dict
            with the results or None if the status is not COMPLETED.
            If run script is a benchmark, then the score is usually expected to
            be in the `score` field.
        """
        self._download_files()
        return super().run()

    def status(self) -> Tuple[Status, List[Tuple[datetime, str, Any]]]:
        self._download_files(ignore_missing=True)
        return super().status()
