# ======================================================================================================================
##
# @file         remotemachine/_aws.py
#
# @author       Patrik Kupčo   \n
#               Faculty of Information Technology \n
#               Brno University of Technology \n
#               kupco.patrik16@gmail.com
#
# @brief        AWS HPC Cluster.
#
# @version      1.3
#
# @date         30 July    2023 
#
# @copyright    Copyright (C) 2023 SC\@FIT Research Group, Brno University of Technology, CZ.
#
# ======================================================================================================================

# ----------------------------------------------------------------------------------------------------------------------
# =========================================== AWS HPC Cluster Cluster Class ============================================
# ----------------------------------------------------------------------------------------------------------------------

from remotemachine import RemoteMachine
from common import JobStatus, TPStatus, HpcStates
import exceptions

import logging
import fabric.contrib.project as fcp
import fabric.api
import fabric.exceptions
import os
from   string import Template
from   pathlib import Path

## AWS class.
#  This is a subclass of the RemoteMachine class.
#  @param [in]  RemoteMachine     Parent class.
class AWS(RemoteMachine):

    ## @var logger
    #  Class attribute.
    logger   = logging.getLogger(__name__)

# ----------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------- Getters and Setters --------------------------------------------------
# ----------------------------------------------------------------------------------------------------------------------

    ## Getter of the local home path.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  Local home path.
    @property
    def localHomePath(self):
        return (Template(self._homePathTemplate).substitute(adminName = self.adminName))

    ## Getter of the local work path.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  Local work path.
    @property
    def localWorkPath(self):
        return (Template(self._workPathTemplate).substitute(adminName = self.adminName))

    ## Getter of the remote user's work path.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  Remote work path.
    @property
    def remoteWorkPath(self):
        return (Template(self._remoteUserPathTemplate).substitute(adminName = self.adminName,
                                                                  userName  = self.userName))

    ## Getter of the remote scratch path.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  Remote scratch path.
    @property
    def remoteScratchPath(self):
        return (Template(self._remoteScratchPathTemplate).substitute(userName = self.userName))

    ## Getter of the remote path where the results are located.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  Path to the results stored remotely.
    @property
    def resultsPath(self):
        return (os.path.join((Template(self._remoteScratchPathTemplate).substitute(userName = self.userName)),
                             "results"))

    ## Getter of the remote path where the submission job scripts are located.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  Path to the results stored remotely.
    @property
    def jobScriptsPath(self):
        return (os.path.join((Template(self._remoteUserPathTemplate).substitute(adminName = self.adminName,
                                                                                userName  = self.userName)),
                             "job_scripts"))

    ## Getter of the remote path where program binaries are located.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  Binaries locations.
    @property
    def programsPath(self):
        return (Template(self._programsPath).substitute(adminName = self.adminName))

    ## Getter of the local SSH key file path.
    #  @param   self The object pointer.
    #  @rtype   string
    #  @return  SSH key file path.
    @property
    def sshKeyFilePath(self):
        return (Template(self._sshKeyFilePath).substitute(adminName = self.adminName))

    ## The initialization method of the class.
    #  @param           self                    The object pointer.
    #  @param   [in]    treatmentPlanRecord     (*PeeWee object*) - Treatment plan database object.
    #  @param   [in]    userName                (*string*) - User name (login).
    def __init__(self, treatmentPlanRecord, userName):
# ----------------------------------------------------------------------------------------------------------------------
        self._tpRecord                  = treatmentPlanRecord
        self._userName                  = userName
        self._maxSubmittedJobs          = 100 # Maximum number of jobs in queues per user
        # Set class variables
        self._name                      = treatmentPlanRecord.hpc.name
        self._adminName                 = ''
        self._remote                    = treatmentPlanRecord.hpc.hostname
        # Set specific remote paths
        self._remoteUserPathTemplate    = '/sharedStorageEfs/dispatcher/$adminName/$userName/tp_%d'
        # Fill in the project name and tp id
        self._remoteScratchPathTemplate = '/sharedStorageEfs/dispatcher/$userName/%s/tp_%d'
        self._programsPath              = '/sharedStorage/software/MATLAB/MatlabRuntime'
        self._sshKeyFilePath            = '/home/xkupco00/.ssh/WSL-KEY2.pem'
        self._jobSubmit                 = 'sbatch'
        self._jobDependency             = '-d afterok'
        self._jobAlter                  = 'scontrol'
        self._jobDeletion               = 'scancel'
        self._jobStat                   = 'squeue'
        self._modulesLoadCmd            = 'ml'
        self._jobIdSuffix               = ''
        self._jobScriptHeader           = '''#!/bin/bash
#SBATCH -p {{ queue }}
#SBATCH -A {{ project }}
#SBATCH -N {{ nodes }}
#SBATCH --cpus-per-task={{ cpus }}
#SBATCH -t {{ walltime }}
#SBATCH -J {{ job_name }}
#SBATCH -e {{ results_path }}/{{ job_name }}.err
#SBATCH -o {{ results_path }}/{{ job_name }}.log'''

        # Set Fabric's environment
        fabric.api.env['host_string']    = '%s@%s' % (self._adminName, self._remote)
        fabric.api.env['username']       = self._userName
        fabric.api.env['user']           = self._userName
        fabric.api.env['key_filename']   = Template(self._sshKeyFilePath).substitute(adminName=self._adminName)
        fabric.api.env['warn_only']      = True
        fabric.api.env['use_ssh_config'] = True

        # Check the existence of SSH keys
        if not Path(Template(self._sshKeyFilePath).substitute(adminName=self._adminName)).is_file():
            self._tpRecord.setStatus(TPStatus.NEW)
            self._tpRecord.errorNote = "Missing keys for selected HPC."
            self._tpRecord.incRestartsNumber()
            self._tpRecord.hpc.setStatus(HpcStates.DOWN)
            self.logger.warning("HPC[%d]: Missing SSH keys. HPC going to DOWN." % self._tpRecord.hpc)
            raise RemoteMachine.RemoteMachineException(RemoteMachine.Errors.SSHKEY_ERR.value)
        # Try to connect to the cluster
        self.tryConnection()

    ## Requests a remote job's status and parses the output.
    #
    #  Taken from the section 11.9, "Job Exit Codes", on page 984 of the PBS Professional Administrator’s Guide.
    #  @param   self    The object pointer.
    #  @param   [in]    remoteId    (*int*) - Id of the job on the remote machine.
    #  @rtype   tuple
    #  @return  Tuple of *JobStatus enumeration class* item, *float* (consumed core hours), *float* walltime,
    #           and *string* (comment).
    def getJobStatus(self, remoteId):
# ----------------------------------------------------------------------------------------------------------------------
        self.logger.debug("Get status for remote id %s" % remoteId)
        jobHistory = {}  # convert fabric string into dictionary
        with fabric.api.hide('running', 'stdout', 'stderr'):
            # each line in fabric string is ended by \r\n
            try:
                string = fabric.api.run('%s -j %s -X --format state,cputime,ncpus,comment' % ("sacct", remoteId)).rsplit("\r\n")
                headers = string[0].split()
                data = string[2].split()
                for header, value in zip(headers, data):
                    jobHistory[header.lower()] = value
            except Exception as e: # Job not yet accounted for, should be by the next run, set as QUEUED
                print(e)
                pass 
        status = jobHistory.get('state')

        if status in ['PENDING', 'REQUEUED'] or status == None:
            status = JobStatus.QUEUED
        elif status == 'RUNNING':
            status = JobStatus.RUNNING
        elif status == 'COMPLETED':
            status = JobStatus.COMPLETE
        elif status == 'SUSPENDED':
            status = JobStatus.HOLD
        elif status.startswith('CANCELLED'): # Is followed with the id of user, who cancelled the job
            status = JobStatus.DELETED
        elif status in ['FAILED', 'RESIZING', 'REVOKED', 'BOOT_FAIL', 'DEADLINE', 'NODE_FAIL', 'OUT_OF_MEMORY', 'PREEMPTED', 'TIMEOUT']:
            status = JobStatus.ERROR

        # consumed core-hours
        # cpu time
        consumedCoreHours = None
        walltime = None
        if jobHistory.get('cputime') is not None:
            # walltime
            wh, wm, ws = (jobHistory.get('cputime')).split(':')
            wh = float(wh)
            wm = float(wm)
            ws = float(ws)
            walltime = wh + wm / 60 + ws / 3600

            consumedCoreHours = walltime * self._tpRecord.hpc.normalizationFactor

        comment = None
        comment = jobHistory.get('comment')
        return (status, consumedCoreHours, walltime, comment)

    ## Uploads a file to a remote machine.
    #  @param   self    The object pointer.
    #  @param   [in]    localFilePath   (*string*) - Local path to the file (including) on k-Dispatch.
    #  @param   [in]    remoteFilePath  (*string*) - Remote path to the file (including).
    def uploadFile(self, localFilePath, remoteFilePath):
# ----------------------------------------------------------------------------------------------------------------------
        try:
            with fabric.api.hide('running', 'stdout', 'stderr'):
                result = fabric.api.run('mkdir -p %s' % remoteFilePath)
                if result.failed:
                    raise fabric.exceptions.NetworkError("[TP ID = %d]: Folder could not be created on the remote site."
                                                         % self._tpRecord.getId)

                try:
                    fcp.rsync_project(remote_dir   = remoteFilePath,
                                      local_dir    = localFilePath,
                                      default_opts = '-azvP',
                                      ssh_opts     = '-o StrictHostKeyChecking=no')
                except (fabric.exceptions.NetworkError, fabric.exceptions.CommandTimeout):
                    raise fabric.exceptions.NetworkError("[TP ID = %d]: Planning file failed to be transferred to "
                                                         "remote." % self._tpRecord.getId)
                else:
                    # Delete the file from the temporary folder on the local disk
                    result = fabric.api.local('rm %s' % localFilePath)
                    #if result.failed:
                    #    raise fabric.exceptions.NetworkError("[TP ID = %d]: Planning file was not deleted from the "
                    #                                         "local disk." % self._tpRecord.getId)
        except Exception as e:
            raise

    ## Downloads the result file from a remote machine.
    #  Target result folder is zipped and then downloaded.
    #  @param   self    The object pointer.
    def downloadFile(self):
# ----------------------------------------------------------------------------------------------------------------------
        try:
            self.logger.info("Downloading the result file (TP ID: %d) ..." % self._tpRecord.getId)
            if self._tpRecord.resultFileName is None:
                self.logger.error("[TP ID = %d]: Result file is not specified. Could not be downloaded."
                                  % self._tpRecord.getId)
            else:
                with fabric.api.hide('running', 'stdout', 'stderr'):
                    # Add TP ID to paths
                    resultFolderPath  = self.resultsPath       % (self._tpRecord.getAllocation.name.lower(),
                                                                  self._tpRecord.getId)
                    remoteScratchPath = self.remoteScratchPath % (self._tpRecord.getAllocation.name.lower(),
                                                                  self._tpRecord.getId)
                    # Zip the results folder and store the file into the remote path
                    result = fabric.api.run('zip -rj %s %s' % (os.path.join(remoteScratchPath,
                                                                            self._tpRecord.resultFileName),
                                                                resultFolderPath))
                    if result.failed:
                        self.logger.error('[TP ID = %d]: Result file could not be zipped (%d).'
                                          % (self._tpRecord.getId, result.return_code))
                        raise Exception("[TP ID = %d]: Result file could not be zipped." % self._tpRecord.getId)
                    else:
                        self.logger.info('[TP ID = %d]: Zip and download: %s' % (self._tpRecord.getId, result.succeeded))
                    # Download the zipped file
                    fabric.api.local('mkdir -p %s' % self._tpRecord.baseFolderPath)
                try:
                    result = fcp.rsync_project(remote_dir   = os.path.join(remoteScratchPath,
                                                                           self._tpRecord.resultFileName),
                                               local_dir    = os.path.join(self._tpRecord.baseFolderPath,
                                                                           self._tpRecord.resultFileName),
                                               default_opts = '-azvP', # Override default rsync parameters
                                               ssh_opts     = '-o StrictHostKeyChecking=no',
                                               upload       = False,   # Downloading
                                               capture      = True)    # Otherwise, the local method throws an error
                    if result.failed:
                        self.logger.error('[TP ID = %d]: Result file download crashed (%d).'
                                          % (self._tpRecord.getId, result.return_code))
                        raise Exception("[TP ID = %d]: Result file download crashed." % self._tpRecord.getId)
                except:
                    raise
                else:
                    self.logger.debug('[TP ID = %d]: Result file downloaded.' % self._tpRecord.getId)
                    print("Result file downloaded", self._tpRecord.getId)
        except (fabric.exceptions.NetworkError, Exception) as e:
            raise


    ## Submits a remote job to a computational queue.
    #  @param           self    The object pointer.
    #  @param   [in]    scriptName     (*string*) - Name of the file (excluding its path).
    #  @param   [in]    predecessors   (*list*)   - List of strings, i.e., preceding job IDs.
    #  @rtype   string
    #  @return  Remote job identifier.
    #  @throw   Exception. When job submission fails.
    def submitJob(self, scriptName, predecessors):
# ----------------------------------------------------------------------------------------------------------------------
        jobId = None
        if predecessors == None:
            jobId = fabric.api.run('%s %s/%s' % (self.jobSubmitCmd, self.jobScriptsPath % self._tpRecord.getId,
                                                 scriptName))
        elif type(predecessors) is list:
            dependString = ""
            for p in predecessors:
                dependString += ":%s" % p
            jobId = fabric.api.run('%s %s%s %s' % (self.jobSubmitCmd,
                                                   self.jobDependencyCmd,
                                                   dependString,
                                                   os.path.join(self.jobScriptsPath % self._tpRecord.getId,
                                                                scriptName)))
        else: # just a single predecessor
            jobId = fabric.api.run('%s %s%s %s' % (self.jobSubmitCmd,
                                                   self.jobDependencyCmd,
                                                   ":%s" % predecessors,
                                                   os.path.join(self.jobScriptsPath % self._tpRecord.getId,
                                                                scriptName)))
        # Parse jobId
        try:
            jobId = int(jobId[len('Submitted batch job '):])
            """
            number, suffix = jobId.split('.')
            if not (int(number) and self._jobIdSuffix in suffix):
                raise
            """
        except Exception:
            self._tpRecord.setStatus(HpcStates.ERROR)
            self._tpRecord.errorNote = jobId
            raise Exception("sbatch returned: %s" % jobId)
        return(jobId)

    ## Alters (restarts) the submitted job and returns its new remote ID.
    #  qalter does not return anything. PBS ID remains.
    #  @param           self    The object pointer.
    #  @param   [in]    remoteId       (*string*) - Id of the job on the remote machine.
    #  @param   [in]    predecessors   (*string* | *list*) - List of strings, i.e., preceding job IDs.
    #                                   May be None.
    def restartJob(self, remoteId, predecessors):
# ----------------------------------------------------------------------------------------------------------------------
        if predecessors is None:
            fabric.api.run('%s %s %s' % (self.jobAlterCmd,
                                         self.jobDependencyCmd,
                                         remoteId))
        elif type(predecessors) is list:
            dependString = ""
            for p in predecessors:
                dependString += ":%s" % p
            fabric.api.run('%s %s%s %s' % (self.jobAlterCmd,
                                           self.jobDependencyCmd,
                                           dependString,
                                           remoteId))
        else: # just a single predecessor
            fabric.api.run('%s %s%s %s' % (self.jobAlterCmd,
                                           self.jobDependencyCmd,
                                           ":%s" % predecessors,
                                           remoteId))



    ## Cancels the remote job or jobs.
    #  @param           self            The object pointer.
    #  @param   [in]    remoteIds       (*string* | *list*) - A single remote job ID (string) or a list of strings as a
    #  list of job IDs to be cancelled.
    def cancelJob(self, remoteIds):
# ----------------------------------------------------------------------------------------------------------------------
        try:  # try to cancel them
            jobsToCancelIdString = ''
            if type(remoteIds) is list:
                for id in remoteIds:
                    jobsToCancelIdString = '%s %s ' % (jobsToCancelIdString, id)
            else:
                jobsToCancelIdString = remoteIds

            with fabric.api.hide('running', 'stdout', 'stderr'):
                fabric.api.run('%s %s' % (self.jobDeleteCmd, jobsToCancelIdString))
        except:  # Meantime, job could have finished or they have not been submitted.
            pass  # We do not wanna receive errors like this.

    ## Cleans remote disk space.
    #  @param   self    The object pointer.
    def cleanRemoteSpace(self):
# ----------------------------------------------------------------------------------------------------------------------
        remoteHomePath    = self.remoteWorkPath    % self._tpRecord.getId
        remoteScratchPath = self.remoteScratchPath % (self._tpRecord.getAllocation.name.lower(),
                                                      self._tpRecord.getId)

        with fabric.api.hide('running', 'stdout', 'stderr', 'warnings'), \
             fabric.api.settings(warn_only=True): # avoid abortion
            resultHome    = fabric.api.run('rm -r %s' % remoteHomePath)
            resultScratch = fabric.api.run('rm -r %s' % remoteScratchPath)

        if resultHome.failed and resultScratch.failed and not self._tpRecord.remoteCleaned:
            self.logger.debug("[TP ID = %d]: Remote home space has been cleaned." % self._tpRecord.getId)
            self._tpRecord.remoteCleaned = True
        elif resultHome.failed or resultScratch.failed:
            self.logger.info("[TP ID = %d]: Remote space could not been cleaned again. "
                             "Home already cleaned: %s. Scratch already cleaned: %s."
                             % (self._tpRecord.getId, resultHome.failed, resultScratch.failed))

    ## Creates folders for results, planning file and submission scripts on the remote machine.
    #  These folders are located under the dispatcher/\<user\>/tp_\<tpID\>/ folder.
    #  @param self The object pointer.
    #  @throw fabric.exceptions.NetworkError Network error (SSH).
    def createFoldersOnRemoteMachine(self):
# ----------------------------------------------------------------------------------------------------------------------
        try:
            if fabric.api.env.machine_name == 'mock':
                raise exceptions.MockException
            with fabric.api.hide('running', 'stdout', 'stderr'):
                resultPath = self.resultsPath       % (self._tpRecord.getAllocation.name.lower(), self._tpRecord.getId)
                planPath   = self.remoteScratchPath % (self._tpRecord.getAllocation.name.lower(), self._tpRecord.getId)
                scriptPath = self.jobScriptsPath    % self._tpRecord.getId

                res = fabric.api.run('mkdir -p %s' % resultPath)
                if res.failed:
                    raise fabric.exceptions.NetworkError
                res = fabric.api.run('mkdir -p %s' % planPath)
                if res.failed:
                    raise fabric.exceptions.NetworkError
                res = fabric.api.run('mkdir -p %s' % scriptPath)
                if res.failed:
                    raise fabric.exceptions.NetworkError
        except exceptions.MockException:
            pass
        except Exception as e:
            raise

    ## Function returns information of how many jobs are currently in the queues and what is the maximum number of
    #  submitted jobs.
    #  @param self The object pointer.
    #  @rtype tuple
    #  @return A tuple of a number of currently submitted jobs, and the maximum number of jobs that can be submitted.
    #  @throw fabric.exceptions.NetworkError Network error (SSH).
    def getSubmittedJobsLimits(self):
        try:
            with fabric.api.hide('running', 'stdout', 'stderr'):
                submittedJobs = fabric.api.run("%s -u $USER | wc -l" % (self.jobStatCmd))
            return(int(submittedJobs) - 1, self._maxSubmittedJobs)
        except:
            # In case of error number of submitted jobs are not known
            # This prevents submitting new jobs to cluster - no space
            return(self._maxSubmittedJobs, self._maxSubmittedJobs)
            raise
