# SPDX-FileCopyrightText: 2022 German Aerospace Center (DLR)
#
# SPDX-License-Identifier: MIT

"""
Created on 25.01.2017

@author: freu_se
"""
import getpass
import os
import re
import shutil
import subprocess
import tarfile
import zlib
from datetime import datetime

from patme.service.duration import Duration
from patme.service.logger import log
from patme.service.systemutils import dos2unix, tarGzDirectory
from patme.sshtools.cara import copyClusterFilesSCP, get_default_slurm_args, sshCluster, sshClusterJob
from patme.sshtools.sshcall import callSSH

duration = Duration()


class ClusterCaller:
    """This class represents a general interface to call several fe and other functions on the cluster.

    The remote jobs are performed by copying the files of the local input directory to the
    cluster on "\\\\cluster.fa.bs.dlr.de\\<username>\\job\\<runDirName>".
    Then the program creates an ssh connection to the cluster. More information about
    the ssh connection can be found in service.utilities.callSSH.
    After the completion of the job the result is copied back to the local runDir.

    :param runDir: absolute or relative path to the folder where the fe run should be executed
    """

    REMOTE_SCRIPT = "run_sbatch.sh"

    clusterFuncDict = {
        "cara": {
            "copyClusterFilesSCP": copyClusterFilesSCP,
            "sshCluster": sshCluster,
            "sshClusterJob": sshClusterJob,
        },
    }

    def __init__(self, runDir, **kwargs):
        self.runDir = runDir
        self.remoteBaseDir = None
        """Directory where files will be copied remotely. Defaults to solverName in runRemote
        remoteRunDir = ~/remoteBaseDir/self.runDir[-1]"""
        self.remoteCmds = None
        self.solverName = None
        self.preCommands = []
        """List of commands that are sent to the cluster prior to the remote command.
        It is intended to setup the environment for the call that will be done
        """
        self.ignorePatternsHostToLocal = {".svn"}
        """Files and folders that shall not copied from the cluster to the local machine on remote runs."""

        self.reuseResults = False
        """If True, the resulting tar.gz file is also extracted on the remote machine.
        This way, the result folder can be reused on a next remote call without copying inputs."""

        self._clusterName = kwargs.pop("clusterName", "cara")
        self.__checkClusterName(self._clusterName)

        if "wsl" in self.clusterName:
            self.remoteCmds = [
                "bash",
                "-i",
                self.REMOTE_SCRIPT,
            ]
        elif "cara" in self.clusterName:
            self.remoteCmds = [
                f"chmod a+x {self.REMOTE_SCRIPT} ; ",
                "sbatch",
                self.REMOTE_SCRIPT,
            ]

    def _getClusterFunc(self, funcName):
        """doc"""
        if self.clusterName not in self.clusterFuncDict:
            raise Exception(f"Unknown cluster {self.clusterName}")

        if funcName not in self.clusterFuncDict[self.clusterName]:
            raise Exception(f"Unknown cluster func {funcName}")

        return self.clusterFuncDict[self.clusterName][funcName]

    def runRemote(self, copyFilesLocalToHost=True, jobName="job", **kwargs):
        """doc"""
        log.info(f"call {self.solverName} remotely on '{self.clusterName}'")

        self.addClusterSpecificBatchCmds()

        remoteRunDir = kwargs.get("remoteRunDir", None)
        if not remoteRunDir:
            ws_dir_name, remoteRunDir = self.getRemoteRunDir(**kwargs)

        # copy rundir to network folder
        if copyFilesLocalToHost:
            log.info(f"Copy files to remote machine '{self.clusterName}'")
            with log.switchLevelTemp(log.WARN):
                self._transferInputToCluster(remoteRunDir, **kwargs)

        # Replace <jobName> two times since it used two times in the command list
        remoteCmds = [re.sub("<jobName>", jobName, elem) for elem in self.remoteCmds]

        # change to run dir of fe call
        remoteCmdsJoin = " ".join(remoteCmds)
        if "wsl" not in self.clusterName:
            command = f"cd {remoteRunDir};{remoteCmdsJoin}"
        else:
            command = remoteCmdsJoin

        if self.preCommands:
            precmd = ";".join(self.preCommands)
            command = f"{precmd};{command}"

        log.debug(f'Call "{self.solverName}" with the following command:')
        log.debug(f"{command}", longMesageDelim=";")

        if "cara" in self.clusterName:
            # =======================================================================
            # run remote
            # =======================================================================
            if self.clusterName not in self.clusterFuncDict:
                raise Exception(f"Unknown cluster {self.clusterName}")

            sshClusterFunc = self._getClusterFunc("sshClusterJob")

            jobId = sshClusterFunc(command, printOutput=False, **kwargs)
        else:

            jobId = "WSL_NOJOBID"
            command += " > cluster.r00.log"
            remoteRunDirAsWinPath = self.__getWSLRemoteDirAsWinPath(remoteRunDir)
            wsl_cmd = f"wsl --cd {remoteRunDirAsWinPath} {command}"
            p = subprocess.run(wsl_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        log.debug(f"Transfer files from remote to local working dir")
        with log.switchLevelTemp(log.WARN):
            self._transferFromCluster(remoteRunDir, jobId, **kwargs)

        self.cleanRemoteDir(ws_dir_name, remoteRunDir, **kwargs)

        log.info(f"Remote call {self.solverName} with jobId {str(jobId)} done")

    def cleanRemoteDir(self, ws_name, remoteRunDir, **kwargs):
        """Clean scratch directory and release scratch"""
        if self.clusterName == "cara":

            from patme.sshtools.cara import _getClusterAuthentication

            hostname, hostKeyString, privateKey = _getClusterAuthentication()
            username = kwargs.get("username", getpass.getuser())
            cmd = f"rm -rf {remoteRunDir}/*;ws_release {ws_name}"
            _ = callSSH(hostname, cmd, privateKey, hostKeyString=hostKeyString, username=username, printOutput=False)

        elif "wsl" in self.clusterName:

            shutil.rmtree(self.__getWSLRemoteDirAsWinPath(remoteRunDir))

    def getRemoteRunDir(self, **kwargs):
        """doc"""
        if self.clusterName == "cara":

            from patme.sshtools.cara import _getClusterAuthentication

            hostname, hostKeyString, privateKey = _getClusterAuthentication()
            username = kwargs.get("username", getpass.getuser())
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
            ws_dir_name = f"run_{timestamp}"
            cmd = f"ws_allocate {ws_dir_name}"
            log.info("Allocate workspace: " + ws_dir_name)
            with log.switchLevelTemp(log.WARN):
                scratch_dir = callSSH(
                    hostname, cmd, privateKey, hostKeyString=hostKeyString, username=username, printOutput=False
                )

            remoteRunDir = scratch_dir.strip()

        elif "wsl" in self.clusterName:

            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
            ws_dir_name = f"run_{timestamp}"
            remoteRunDir = f"~/tmp/{ws_dir_name}"
            os.makedirs(self.__getWSLRemoteDirAsWinPath(remoteRunDir))

        else:
            raise Exception(f"Cluster with name {self.clusterName} not supported!")

        return ws_dir_name, remoteRunDir

    def _transferInputToCluster(self, remoteRunDir, **kwargs):
        """Transfer the input to the cluster by performing 4 steps.

        1. tar.gz the whole input directory
        2. transfer this tarfile to the cluster via sftp
        3. extract the file on the cluster
        4. remove the tarfiles on both ends
        """
        log.info(f"Copy files to remote machine: {remoteRunDir}")
        tarFileName = "clusterInput.tar.gz"
        tarFileNameFull = tarGzDirectory(self.runDir, tarFileName)

        if "cara" in self.clusterName:
            username = kwargs.get("username", getpass.getuser())
            if self.clusterName not in self.clusterFuncDict:
                raise Exception(f"Unknown cluster {self.clusterName}")

            clusterFuncCopy = self._getClusterFunc("copyClusterFilesSCP")
            sshClusterFunc = self._getClusterFunc("sshCluster")

            clusterFuncCopy(
                [tarFileName], srcBaseDir=self.runDir, destBaseDir=remoteRunDir, mode="put", username=username
            )

            # unzip files at cluster and remove zip files on both sides
            extract = [f"cd {remoteRunDir}", f"tar -xf {tarFileName}", f"rm {tarFileName}"]

            sshClusterFunc(";".join(extract), printOutput=False, username=username)

        elif "wsl" in self.clusterName:

            remoteRunDirAsWinPath = self.__getWSLRemoteDirAsWinPath(remoteRunDir)
            shutil.copy(tarFileNameFull, remoteRunDirAsWinPath)

            cmd = f"wsl --cd {remoteRunDir} tar -xf {tarFileName} -C .;rm {tarFileName}"
            subprocess.run(cmd, shell=True)

        else:
            raise NotImplementedError(f"Unknown cluster name or wsl: {self.clusterName}")

        os.remove(os.path.join(self.runDir, tarFileName))

        log.debug(f"Copy files to remote machine '{self.clusterName}' done")

    def __getWSLRemoteDirAsWinPath(self, remoteDir):
        """doc"""
        p = subprocess.run(["wsl", "wslpath", "-w", remoteDir], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        return p.stdout.decode().strip()

    def _transferFromCluster(self, remoteRunDir, jobId, **kwargs):
        """Transfer the calculation result from the cluster by performing these steps.

        1. copy the result tar.gz file to self.runDir
        2. extract the file locally
        3. extract tar file remotely if the remote folder should be reused
        4. remove the tarfile locally
        4.1 optionally extract tarfile remotely to reuse the calculation results
        5. remove the tarfile remote"""

        def getFileAndExtract(tarFileName, tarFileNamePath, clusterLog, remoteRunDir, copyFunc, username=None):
            """doc"""
            for ffile in [tarFileName, clusterLog]:
                copyFunc([ffile], remoteRunDir, self.runDir, "get", username=username)

            with tarfile.open(tarFileNamePath) as f:
                f.extractall(self.runDir, filter="data")

        if "wsl" in self.clusterName:

            remoteRunDirAsWinPath = self.__getWSLRemoteDirAsWinPath(remoteRunDir)
            dir_creation_time = os.path.getctime(remoteRunDirAsWinPath)
            for dirpath, _, filenames in os.walk(remoteRunDirAsWinPath):
                for filename in filenames:

                    filepath = os.path.join(dirpath, filename)
                    try:
                        file_mtime = os.path.getmtime(filepath)
                        # Compare the file's modification time with the directory's creation time
                        if file_mtime > dir_creation_time:
                            rel_dir = os.path.relpath(dirpath, remoteRunDirAsWinPath)
                            dirInRunDir = os.path.join(self.runDir, rel_dir)
                            os.makedirs(dirInRunDir, exist_ok=True)
                            shutil.copy(filepath, dirInRunDir)

                    except FileNotFoundError:
                        # Skip files that might disappear during walk
                        continue

        else:
            username = kwargs.get("username", getpass.getuser())
            copyFunc = self._getClusterFunc("copyClusterFilesSCP")
            sshClusterFunc = self._getClusterFunc("sshCluster")

            try:
                log.info("copy result from cluster")
                prefix = f"cluster.r{jobId}"
                tarFileName = f"{prefix}.tar.gz"
                clusterLog = f"{prefix}.log"
                tarFileNamePath = os.path.join(self.runDir, tarFileName)

                try:
                    getFileAndExtract(
                        tarFileName, tarFileNamePath, clusterLog, remoteRunDir, copyFunc, username=username
                    )

                except tarfile.ReadError:
                    msg = "Could not open the result tar file. "
                    msg += "Probably the calculation was cancelled."
                    log.error(msg)
                    return  # retruning to make a local calculation possible

                except zlib.error:
                    msg = "Got error while extracting result file. "
                    msg += "Try again to copy and extract"
                    log.error(msg)
                    try:
                        getFileAndExtract(tarFileName, tarFileNamePath, clusterLog, remoteRunDir)
                    except tarfile.ReadError:
                        msg = "Could not open the result tar file. "
                        msg += "Probably the calculation was cancelled."
                        log.error(msg)
                        return  # retruning to make a local calculation possible

                os.remove(tarFileNamePath)

                cleanupCommand = f"cd {remoteRunDir}"
                if self.reuseResults:
                    cleanupCommand += f";tar -xf {tarFileName}"

                cleanupCommand += f";rm {tarFileName}"

                sshClusterFunc(cleanupCommand, printOutput=False, username=username)
                log.debug("copy result from cluster done")

            except OSError as exception:

                msg = "Retrieved IOError when copying results from remote computer. "
                msg += "Maybe some items are not copied properly. "
                msg += "See debug.log for more details."

                log.warning(msg)
                log.debug("error message to the above warning: " + str(exception))

    def addClusterSpecificBatchCmds(self, *args, **kwargs):
        """Create bash cmds which for remote sbatch call"""
        runScriptCmds = ["#!/bin/bash"]
        if self.clusterName == "cara":

            slurm_args = get_default_slurm_args()
            cara_partion = os.environ.get("CARA_PARTITION", slurm_args["partition"])
            slurm_args["partition"] = cara_partion
            if cara_partion == "ppp":
                slurm_args["ntasks"] = "1"

            for sarg, sval in slurm_args.items():
                cmd = f"#SBATCH --{sarg}"
                cmd += "" if not sval else f"={sval}"
                runScriptCmds.append(cmd)

            runScriptCmds += [
                ". /usr/share/lmod/lmod/init/sh",
                ". /etc/profile.d/10-dlr-modulepath.sh",
                "module --force purge",
            ]

        runScriptCmds += [
            'curtime=$(date +"%Y-%m-%d %H:%M:%S")',
            # find can only distinguish between timestamps which have seconds as highest precision
            "sleep 1s",
            'tarFile="cluster.r$SLURM_JOB_ID.tar.gz"',
        ]
        runScriptCmds += self.generateSolverSpecificRemoteCmds()
        runScriptCmds += [
            'tar -czf $tarFile -C . $(find . -newermt "$curtime" -type f)',
        ]

        runScriptFile = os.path.join(self.runDir, self.REMOTE_SCRIPT)
        with open(runScriptFile, "w") as f:
            f.write("\n".join(runScriptCmds))

        dos2unix(runScriptFile)

    def __checkClusterName(self, cname):
        if not re.search("cara|wsl", cname):
            raise NotImplementedError(f"No remote caller implemented for cluster '{cname}'")

    def _getClusterName(self):
        return self._clusterName

    def _setClusterName(self, cname):
        self.__checkClusterName(cname)
        self._clusterName = cname

    clusterName = property(fget=_getClusterName, fset=_setClusterName)


class PythonModuleCaller(ClusterCaller):
    """This class handles calling target functions like those in buckling.bucklingtargetfunction on the cluster"""

    def __init__(self, runDir, jobCommands):
        """
        :param jobCommands: list with strings defining the command to call a python script excluding the leading python executable
        """
        ClusterCaller.__init__(self, runDir)
        self.solverName = "python35"
        clusterCmds = ["py35", "-d", "-O", "tgz", "-J", "<jobName>"]
        self.remoteCmds = clusterCmds + jobCommands
        self.preCommands = [". ${HOME}/.profile;echo PYTHONPATH:${PYTHONPATH}"]


class MatlabCaller(ClusterCaller):
    """This class handles calling target functions like those in buckling.bucklingtargetfunction on the cluster"""

    solverName = "matlab"

    def __init__(self, runDir, jobCommands):
        """
        :param jobCommands: list with strings defining the command to call a matlab script excluding the leading matlabe executable
        """
        ClusterCaller.__init__(self, runDir)
        clusterCmds = ["mat2013a", "-d", "-O", "tgz", "-J", "<jobName>", "--", "-r"]
        self.remoteCmds = clusterCmds + jobCommands
