"""
Created on 21 March 2019

@author: galdam
"""

import eHive

import random
import string
import os
import subprocess
import shutil
import json
from typing import Dict, List, Tuple


class Singularity(eHive.BaseRunnable):
    """
    The singularity runnable is intended as a base class for running workflow tasks inside singularity images.
    In most cases, it's sufficient to override the following parameters:
        - PIPELINE: The name of the pipeline module. Used in logging and as part of the default prefix
        - CMD: a new-style python format string. This will be the command that should be run in the singularity image.
            Reserved substitutions are:
                - WORKING_DIR: the temporary working directory that has been opened for the workflow
                - PREFIX: generated by the get_prefix function
                - ARGS: the unpacked CMD_ARGS
        - CMD_ARGS: required arguments. These will be searched for in the parameters and substituted in the CMD.
        - CMD_KWARGS: a list of parameter keys that will be unpacked, if they exist in the available params
        - FILES:
    The following parameters must be provided:
        - singularity_executable: the executable for singularity (eg '/bin/singularity')
        - singularity_cache: A directory the contains singularity images
        - singularity_image: Name of the image to use
        - root_output_dir: a root location to write to
    """
    PIPELINE = None
    CMD = None
    CMD_ARGS = []
    CMD_KWARGS = []
    FILES = dict()

    __session_suffix = None

    def run(self):
        """
        The main method of the runnable which defines the steps that will be taken.
        """
        self.setup()

        # Build the commands
        s_cmd = self.make_singularity_command()
        t_cmd = self.make_task_command()
        cmd = f"{s_cmd} {t_cmd}"

        # Execute the command
        self.execute_program(cmd)

        # Store the outputs and delete the temporary directory
        self.close_working_dir()

    # ---------------------------------------------------------------------------------------------------
    #
    def setup(self):
        """
        Setup the general parameters: define the output directory, prefix and verify the file list.
        """
        # Verify that the runnable is properly configured
        assert self.PIPELINE is not None, "Pipeline name has not been defined."

        # Open all directories
        # Creating a parameter here to assist with logging
        self.__directories = {
            'output_directory': self.output_directory,
            'working_directory': self.working_directory,
            'log_directory': self.log_directory
        }
        for d in self.__directories.values():
            os.makedirs(d, exist_ok=True)

        # Log the state of the input parameters
        with open(f"{self.log_basename}.ehive.log", 'w') as log:
            log.write(json.dumps(self, default=lambda o: getattr(o, '__dict__', str(o))))

        # Verify that the file list is valid (fail fast)
        self.get_output_file_list()

    # ================================================================ #
    # Naming Properties
    @property
    def session_suffix(self) -> str:
        """
        A randomised code to differentiate repeats of the runnable.
        :return: str, a randomised set of characters
        """
        if self.__session_suffix is None:
            self.__session_suffix = random_suffix()
        return self.__session_suffix

    @property
    def base_prefix(self) -> str:
        """
        The base name for the samples.
        :return: str, the basename for the sample
        """
        return self.param_required('basename').replace(':', '.')

    @property
    def prefix(self) -> str:
        """
        A value to be prefixed to the output files
        :return: str, the prefix to use for files
        """
        return f"{self.base_prefix}.{self.PIPELINE}"

    # ================================================================ #
    # Directory Properties
    @property
    def output_directory(self) -> str:
        """
        Gives a location to write files to. If the dir_label_params are defined,
        the values will be used to create a directory structure.
        :return: str, the path to the output directory
        """
        root_output = self.param_required('root_output_dir')
        if self.param_is_defined('dir_label_params'):
            dir_label_params = self.param('dir_label_params')
            output_dir = os.path.join(root_output, *[self.param(d).replace(':', '.') for d in dir_label_params])
        else:
            output_dir = root_output
        return output_dir

    @property
    def working_directory(self) -> str:
        """
        A temporary directory where the work is conducted
        :return: str, the working directory
        """
        return os.path.join(self.output_directory, f"{self.prefix}_{self.session_suffix}")

    @property
    def log_directory(self) -> str:
        """
        The location where log files are stored.
        :return: str, the root logging directory
        """
        return os.path.join(self.output_directory, 'ehive_log')

    @property
    def log_basename(self) -> str:
        """
        A path and file name for logging to. Requires a log type at the end.
        :return: str, the basic logging path and log file name
        """
        return os.path.join(self.log_directory, f"{self.base_prefix}.{self.PIPELINE}.{self.session_suffix}")

    def get_output_file_list(self) -> Dict[str, str]:
        """
        Turns the FILES dictionary into the final list. Each file is appended with the output directory and
        formatted with the PREFIX (assuming that prefix is included as a new style python format parameter)
        This function can be overridden to include files that need to be treated differently
        :return: {file label: file path, ...}
        """
        return {k: os.path.join(self.output_directory, v.format(PREFIX=self.prefix)) for k, v in self.FILES.items()}

    def make_singularity_command(self) -> str:
        """
        Sets up the singularity command. Includes selection of the executable, image and setting the working directory.
        :return: str, command for singularity
        """
        command = [
            self.param_required('singularity_executable'),
            "exec",
            f"--pwd {self.working_directory}",
            f"{self.param_required('singularity_cache')}/{self.param_required('singularity_image')}",
        ]
        command = ' '.join(command)
        return command

    def make_task_command(self) -> str:
        """
        Builds the command that will be run within the singularity image.
        This is a combination of the CMD parameter, the WORKING_DIR (internal), PREFIX (internal),
        ARGS (unpacked internally) and the CMD_ARGS
        :return: str, command within singularity
        """
        cmd = self.CMD.format(WORKING_DIR=self.working_directory, PREFIX=self.prefix,
                              ARGS=self.unpack_cmd_kwargs(), **self.get_cmd_args())
        return cmd

    def get_cmd_args(self) -> Dict[str, str]:
        """
        Gathers the arguments from the CMD_ARGS.
        Values that are not defined in both the CMD_ARGS and parameters are skipped.
        :return: dict{str: str, ...}, a dictionary of command keys and values
        """
        return {k: self.param(k) for k in self.CMD_ARGS}

    def unpack_cmd_kwargs(self) -> str:
        """
        Utility for unpacking key-value parameters in the task command.
        :return: str, parameters
        """
        return ' '.join([f"--{k} {self.param(k)}" for k in self.CMD_KWARGS if self.param_is_defined(k)])

    def execute_program(self, cmd: str):
        """
        Execute the provided command line and check the output.
        :param cmd: str, the command to run
        """
        with open(f'{self.log_basename}.stderr.txt', 'wb') as stderr_fh, \
                open(f'{self.log_basename}.stdout.txt', 'wb') as stdout_fh:
            try:
                p = subprocess.run(cmd, shell=True, check=True,
                                   stdout=stdout_fh, stderr=stderr_fh, universal_newlines=True)
            except subprocess.CalledProcessError as err:
                stderr_fh.flush()
                stdout_fh.flush()
                msg = f"\nERROR: {str(err)}\n".encode('utf-8')
                stderr_fh.write(msg)
                stdout_fh.write(msg)
                raise err

    def close_working_dir(self) -> List[Tuple[str, str]]:
        """
        Move the files from the temporary directory to the output directory.
        This step attempts to be atomic by checking before moving files, but
        :return: list[tuple(str, str), ..], moved files
        """
        file_conflict_msg = "Cannot close working directory, files already exist: {}"
        files = os.listdir(self.working_directory)
        transactions = [(os.path.join(self.working_directory, file_name),
                         os.path.join(self.output_directory, file_name))
                        for file_name in files]
        # Check first (it violates the pythonic principles of EAFP, but is more atomic than hitting an issue half way)
        conflicts = [dst for src, dst in transactions if os.path.exists(dst)]
        if conflicts:
            raise shutil.Error(file_conflict_msg.format(' ; '.join(conflicts)))
        for src, dst in transactions:
            try:
                shutil.move(src, dst)
            except shutil.Error as err:
                raise shutil.Error(file_conflict_msg.format(dst)) from err
        shutil.rmtree(self.working_directory)
        return transactions

    def write_output(self):
        """
        Define the output to pass back to eHive.
        The list of files are returned by get_output_file_list and passed to stream #1
        """
        self.warning('Work is done!')
        output_files = self.get_output_file_list()
        for f in output_files.values():
            assert os.path.exists(f), f"Missing file: {f}"
        self.dataflow(output_files, 1)


def random_suffix(suffix_length=5) -> str:
    """
    Generates a randomised string of numbers and lowercase characters.
    This can be used to differentiate temporary directories
    :param suffix_length: int, the number of characters to generate, default: 5
    :return: string, a random string
    """
    return ''.join(random.choices(string.ascii_lowercase + string.digits, k=suffix_length))

