import os
from copy import deepcopy
from typing import List, Optional, Union

from .common import S3Artifact
from .config import config
from .io import InputArtifact
from .op_template import OPTemplate
from .utils import randstr, upload_s3

try:
    from argo.workflows.client import (V1HostPathVolumeSource, V1Volume,
                                       V1VolumeMount)
except Exception:
    pass


class Executor(object):
    """
    Executor
    """

    def render(
            self,
            template: OPTemplate,
    ) -> OPTemplate:
        """
        render original template and return a new template, do not modify
        self in this method to make the executor reusable
        """
        raise NotImplementedError()


def run_script(image, cmd, docker=None, singularity=None, podman=None):
    if docker is not None:
        return "%s pull %s && %s run -v$(pwd)/tmp:/tmp "\
            "-v$(pwd)/script:/script %s %s /script" % (
                docker, image, docker, image, " ".join(cmd))
    elif singularity is not None:
        return "%s pull image.sif %s && %s run -B$(pwd)/tmp:/tmp "\
            "-B$(pwd)/script:/script image.sif %s /script && rm image.sif" % (
                singularity, image, singularity, " ".join(cmd))
    elif podman is not None:
        return "%s pull %s && %s run -v$(pwd)/tmp:/tmp "\
            "-v$(pwd)/script:/script %s %s /script" % (
                podman, image, podman, image, " ".join(cmd))
    else:
        return "%s script" % " ".join(cmd)


def render_script_with_tmp_root(template, tmp_root):
    if hasattr(template, "render_script"):
        tmp_template = deepcopy(template)
        tmp_template.tmp_root = tmp_root
        tmp_template.render_script()
        return tmp_template.script
    else:
        return template.script.replace("/tmp", tmp_root)


class RemoteExecutor(Executor):
    def __init__(
            self,
            host: str,
            port: int = 22,
            username: str = "root",
            password: Optional[str] = None,
            private_key_file: Optional[os.PathLike] = None,
            workdir: str = "~/dflow/workflows/{{workflow.name}}/{{pod.name}}",
            command: Union[str, List[str]] = None,
            remote_command: Union[str, List[str]] = None,
            image: Optional[str] = None,
            image_pull_policy: Optional[str] = None,
            map_tmp_dir: bool = True,
            docker_executable: Optional[str] = None,
            singularity_executable: Optional[str] = None,
            podman_executable: Optional[str] = None,
            action_retries: int = -1,
    ) -> None:
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.private_key_file = private_key_file
        self.workdir = workdir
        if isinstance(command, str):
            command = [command]
        if command is None:
            command = ["sh"]
        self.command = command
        if isinstance(remote_command, str):
            remote_command = [remote_command]
        self.remote_command = remote_command
        if image is None:
            image = config["extender_image"]
        self.image = image
        if image_pull_policy is None:
            image_pull_policy = config["extender_image_pull_policy"]
        self.image_pull_policy = image_pull_policy
        self.map_tmp_dir = map_tmp_dir
        self.docker_executable = docker_executable
        self.singularity_executable = singularity_executable
        self.podman_executable = podman_executable
        if self.docker_executable is not None or \
                self.singularity_executable is not None or \
                self.podman_executable is not None:
            self.map_tmp_dir = False
        self.action_retries = action_retries

    def execute(self, cmd):
        # add '' in case shell will expand ~
        return "execute %s '%s'" % (self.action_retries, cmd)

    def upload(self, src, dst):
        # add '' in case shell will expand ~
        return "upload %s '%s' '%s'" % (self.action_retries, src, dst)

    def download(self, src, dst):
        # add '' in case shell will expand ~
        return "download %s '%s' '%s'" % (self.action_retries, src, dst)

    def run(self, image, remote_command):
        script = "cd %s && " % self.workdir
        return self.execute(script + run_script(
            image, remote_command, self.docker_executable,
            self.singularity_executable, self.podman_executable)) + \
            " || exit 1\n"

    def upload_if_exists(self, path):
        return "if [ -e %s ]; then " % path \
            + self.upload(path, "%s/%s" % (
                self.workdir, os.path.dirname(path))) + "; fi\n"

    def mkdir_and_download(self, path):
        return "mkdir -p %s" % os.path.dirname(path) + "\n" + \
            self.download("%s/%s" % (self.workdir, path),
                          os.path.dirname(path)) + "\n"

    def get_script(self, template):
        if self.map_tmp_dir:
            remote_script = render_script_with_tmp_root(
                template, "%s/tmp" % self.workdir)
        else:
            remote_script = template.script
        remote_command = template.command if self.remote_command is None \
            else self.remote_command
        ssh_pass = "sshpass -p %s " % self.password if self.password is not \
            None else ""
        identity = "-i /root/.ssh/%s " % os.path.basename(
            self.private_key_file) if self.private_key_file is not None else ""
        script = """
execute() {
    if [ $1 != 0 ]; then
        %sssh %s-C -o StrictHostKeyChecking=no -p %s %s@%s -- $2
        if [ $? != 0 ]; then
            echo retry: $1
            sleep 1
            execute $(($1-1)) '$2'
        fi
    fi
}
""" % (ssh_pass, identity, self.port, self.username, self.host)
        script += """
upload() {
    if [ $1 != 0 ]; then
        %sssh %s-C -o StrictHostKeyChecking=no -p %s %s@%s -- mkdir -p $3 && \
        %sscp %s-C -o StrictHostKeyChecking=no -P %s -r $2 %s@%s:$3
        if [ $? != 0 ]; then
            echo retry: $1
            sleep 1
            upload $(($1-1)) $2 $3
        fi
    fi
}
""" % (ssh_pass, identity, self.port, self.username, self.host,
            ssh_pass, identity, self.port, self.username, self.host)
        script += """
download() {
    if [ $1 != 0 ]; then
        %sscp %s-C -o StrictHostKeyChecking=no -P %s -r %s@%s:$2 $3
        if [ $? != 0 ]; then
            echo retry: $1
            sleep 1
            download $(($1-1)) $2 $3
        fi
    fi
}
""" % (ssh_pass, identity, self.port, self.username, self.host)
        script += "cat <<'EOF'> script\n" + remote_script + "\nEOF\n"
        for art in template.inputs.artifacts.values():
            script += self.upload_if_exists(art.path)
        for par in template.inputs.parameters.values():
            if par.save_as_artifact:
                script += self.upload_if_exists(par.path)
        script += self.upload_if_exists("script")
        script += self.run(template.image, remote_command)
        for art in template.outputs.artifacts.values():
            script += self.mkdir_and_download(art.path)
        for par in template.outputs.parameters.values():
            if par.save_as_artifact:
                script += self.mkdir_and_download(par.path)
            elif par.value_from_path is not None:
                script += self.mkdir_and_download(par.value_from_path)
        return script

    def render(self, template):
        new_template = deepcopy(template)
        new_template.name += "-" + randstr()
        new_template.image = self.image
        new_template.image_pull_policy = self.image_pull_policy
        new_template.command = self.command
        new_template.script = self.get_script(template)
        if self.password is not None:
            pass
        elif self.private_key_file is not None:
            key = upload_s3(self.private_key_file)
            private_key_artifact = S3Artifact(key=key)
            new_template.inputs.artifacts["dflow_private_key"] = InputArtifact(
                path="/root/.ssh/" + os.path.basename(self.private_key_file),
                source=private_key_artifact, mode=0o600)
        else:
            new_template.volumes.append(V1Volume(
                name="dflow-private-key", host_path=V1HostPathVolumeSource(
                    path=config["private_key_host_path"])))
            new_template.mounts.append(V1VolumeMount(
                name="dflow-private-key", mount_path="/root/.ssh"))
        return new_template
