import hashlib
import sys
from getpass import getuser

import pytest

from ray.autoscaler._private.command_runner import (
    DockerCommandRunner,
    SSHCommandRunner,
    _with_environment_variables,
)
from ray.autoscaler.command_runner import CommandRunnerInterface
from ray.autoscaler.sdk import get_docker_host_mount_location
from ray.tests.test_autoscaler import MockProcessRunner, MockProvider

auth_config = {
    "ssh_user": "ray",
    "ssh_private_key": "8265.pem",
}


def test_environment_variable_encoder_strings():
    env_vars = {"var1": 'quote between this " and this', "var2": "123"}
    res = _with_environment_variables("echo hello", env_vars)
    expected = """export var1='"quote between this \\" and this"';export var2='"123"';echo hello"""  # noqa: E501
    assert res == expected


def test_environment_variable_encoder_dict():
    env_vars = {"value1": "string1", "value2": {"a": "b", "c": 2}}
    res = _with_environment_variables("echo hello", env_vars)

    expected = """export value1='"string1"';export value2='{"a":"b","c":2}';echo hello"""  # noqa: E501
    assert res == expected


def test_command_runner_interface_abstraction_violation():
    """Enforces the CommandRunnerInterface functions on the subclasses.

    This is important to make sure the subclasses do not violate the
    function abstractions. If you need to add a new function to one of
    the CommandRunnerInterface subclasses, you have to add it to
    CommandRunnerInterface and all of its subclasses.
    """

    cmd_runner_interface_public_functions = dir(CommandRunnerInterface)
    allowed_public_interface_functions = {
        func
        for func in cmd_runner_interface_public_functions
        if not func.startswith("_")
    }
    for subcls in [SSHCommandRunner, DockerCommandRunner]:
        subclass_available_functions = dir(subcls)
        subclass_public_functions = {
            func for func in subclass_available_functions if not func.startswith("_")
        }
        assert allowed_public_interface_functions == subclass_public_functions


def test_ssh_command_runner():
    process_runner = MockProcessRunner()
    provider = MockProvider()
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
    ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
    ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
        ssh_user_hash[:10], ssh_control_hash[:10]
    )
    args = {
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
    }
    cmd_runner = SSHCommandRunner(**args)

    env_vars = {"var1": 'quote between this " and this', "var2": "123"}
    cmd_runner.run(
        "echo helloo", port_forward=[(8265, 8265)], environment_variables=env_vars
    )

    expected = [
        "ssh",
        "-tt",
        "-L",
        "8265:localhost:8265",
        "-i",
        "8265.pem",
        "-o",
        "StrictHostKeyChecking=no",
        "-o",
        "UserKnownHostsFile=/dev/null",
        "-o",
        "IdentitiesOnly=yes",
        "-o",
        "ExitOnForwardFailure=yes",
        "-o",
        "ServerAliveInterval=5",
        "-o",
        "ServerAliveCountMax=3",
        "-o",
        "ControlMaster=auto",
        "-o",
        "ControlPath={}/%C".format(ssh_control_path),
        "-o",
        "ControlPersist=10s",
        "-o",
        "ConnectTimeout=120s",
        "ray@1.2.3.4",
        "bash",
        "--login",
        "-c",
        "-i",
        """'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"quote between this \\" and this"'"'"';export var2='"'"'"123"'"'"';echo helloo)'""",  # noqa: E501
    ]

    # Much easier to debug this loop than the function call.
    for x, y in zip(process_runner.calls[0], expected):
        assert x == y
    process_runner.assert_has_call("1.2.3.4", exact=expected)


def test_docker_command_runner():
    process_runner = MockProcessRunner()
    provider = MockProvider()
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    ssh_control_hash = hashlib.sha256(cluster_name.encode()).hexdigest()
    ssh_user_hash = hashlib.sha256(getuser().encode()).hexdigest()
    ssh_control_path = "/tmp/ray_ssh_{}/{}".format(
        ssh_user_hash[:10], ssh_control_hash[:10]
    )
    docker_config = {"container_name": "container"}
    args = {
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
        "docker_config": docker_config,
    }
    cmd_runner = DockerCommandRunner(**args)
    assert len(process_runner.calls) == 0, "No calls should be made in ctor"

    env_vars = {"var1": 'quote between this " and this', "var2": "123"}
    cmd_runner.run("echo hello", environment_variables=env_vars)

    # This string is insane because there are an absurd number of embedded
    # quotes. While this is a ridiculous string, the escape behavior is
    # important and somewhat difficult to get right for environment variables.
    cmd = """'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (docker exec -it  container /bin/bash -c '"'"'bash --login -c -i '"'"'"'"'"'"'"'"'source ~/.bashrc; export OMP_NUM_THREADS=1 PYTHONWARNINGS=ignore && (export var1='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"quote between this \\" and this"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';export var2='"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"123"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"'"';echo hello)'"'"'"'"'"'"'"'"''"'"' )'"""  # noqa: E501

    expected = [
        "ssh",
        "-tt",
        "-i",
        "8265.pem",
        "-o",
        "StrictHostKeyChecking=no",
        "-o",
        "UserKnownHostsFile=/dev/null",
        "-o",
        "IdentitiesOnly=yes",
        "-o",
        "ExitOnForwardFailure=yes",
        "-o",
        "ServerAliveInterval=5",
        "-o",
        "ServerAliveCountMax=3",
        "-o",
        "ControlMaster=auto",
        "-o",
        "ControlPath={}/%C".format(ssh_control_path),
        "-o",
        "ControlPersist=10s",
        "-o",
        "ConnectTimeout=120s",
        "ray@1.2.3.4",
        "bash",
        "--login",
        "-c",
        "-i",
        cmd,
    ]
    # Much easier to debug this loop than the function call.
    for x, y in zip(process_runner.calls[0], expected):
        print(f"expeted:\t{y}")
        print(f"actual: \t{x}")
        assert x == y
    process_runner.assert_has_call("1.2.3.4", exact=expected)


def test_docker_rsync():
    process_runner = MockProcessRunner()
    provider = MockProvider()
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    docker_config = {"container_name": "container"}
    args = {
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
        "docker_config": docker_config,
    }
    cmd_runner = DockerCommandRunner(**args)

    local_mount = "/home/ubuntu/base/mount/"
    remote_mount = "/root/protected_mount/"
    docker_mount_prefix = get_docker_host_mount_location(cluster_name)
    remote_host_mount = f"{docker_mount_prefix}{remote_mount}"

    local_file = "/home/ubuntu/base-file"
    remote_file = "/root/protected-file"
    remote_host_file = f"{docker_mount_prefix}{remote_file}"

    process_runner.respond_to_call("docker inspect -f", ["true"])
    cmd_runner.run_rsync_up(
        local_mount, remote_mount, options={"docker_mount_if_possible": True}
    )

    # Make sure we do not copy directly to raw destination
    process_runner.assert_not_has_call(
        "1.2.3.4", pattern=f"-avz {local_mount} ray@1.2.3.4:{remote_mount}"
    )
    process_runner.assert_not_has_call("1.2.3.4", pattern=f"mkdir -p {remote_mount}")
    # No rsync -e.*docker exec -i for file_mounts
    process_runner.assert_not_has_call("1.2.3.4", pattern="rsync -e.*docker exec -i")
    process_runner.assert_has_call(
        "1.2.3.4", pattern=f"-avz {local_mount} ray@1.2.3.4:{remote_host_mount}"
    )
    process_runner.clear_history()
    ##############################

    process_runner.respond_to_call("docker inspect -f", ["true"])
    cmd_runner.run_rsync_up(
        local_file, remote_file, options={"docker_mount_if_possible": False}
    )

    # Make sure we do not copy directly to raw destination
    process_runner.assert_not_has_call(
        "1.2.3.4", pattern=f"-avz {local_file} ray@1.2.3.4:{remote_file}"
    )
    process_runner.assert_not_has_call("1.2.3.4", pattern=f"mkdir -p {remote_file}")

    process_runner.assert_has_call("1.2.3.4", pattern="rsync -e.*docker exec -i")
    process_runner.assert_has_call(
        "1.2.3.4", pattern=f"-avz {local_file} ray@1.2.3.4:{remote_host_file}"
    )
    process_runner.clear_history()
    ##############################

    cmd_runner.run_rsync_down(
        remote_mount, local_mount, options={"docker_mount_if_possible": True}
    )

    process_runner.assert_not_has_call("1.2.3.4", pattern="rsync -e.*docker exec -i")
    process_runner.assert_not_has_call(
        "1.2.3.4", pattern=f"-avz ray@1.2.3.4:{remote_mount} {local_mount}"
    )
    process_runner.assert_has_call(
        "1.2.3.4", pattern=f"-avz ray@1.2.3.4:{remote_host_mount} {local_mount}"
    )

    process_runner.clear_history()
    ##############################

    cmd_runner.run_rsync_down(
        remote_file, local_file, options={"docker_mount_if_possible": False}
    )

    process_runner.assert_has_call("1.2.3.4", pattern="rsync -e.*docker exec -i")
    process_runner.assert_not_has_call(
        "1.2.3.4", pattern=f"-avz ray@1.2.3.4:{remote_file} {local_file}"
    )
    process_runner.assert_has_call(
        "1.2.3.4", pattern=f"-avz ray@1.2.3.4:{remote_host_file} {local_file}"
    )


def test_rsync_exclude_and_filter():
    process_runner = MockProcessRunner()
    provider = MockProvider()
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    args = {
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
    }
    cmd_runner = SSHCommandRunner(**args)

    local_mount = "/home/ubuntu/base/mount/"
    remote_mount = "/root/protected_mount/"

    process_runner.respond_to_call("docker inspect -f", ["true"])
    cmd_runner.run_rsync_up(
        local_mount,
        remote_mount,
        options={
            "docker_mount_if_possible": True,
            "rsync_exclude": ["test"],
            "rsync_filter": [".ignore"],
        },
    )

    process_runner.assert_has_call(
        "1.2.3.4", pattern="--exclude test --filter dir-merge,- .ignore"
    )


def test_rsync_without_exclude_and_filter():
    process_runner = MockProcessRunner()
    provider = MockProvider()
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"
    args = {
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
    }
    cmd_runner = SSHCommandRunner(**args)

    local_mount = "/home/ubuntu/base/mount/"
    remote_mount = "/root/protected_mount/"

    process_runner.respond_to_call("docker inspect -f", ["true"])
    cmd_runner.run_rsync_up(
        local_mount,
        remote_mount,
        options={
            "docker_mount_if_possible": True,
        },
    )

    process_runner.assert_not_has_call("1.2.3.4", pattern="--exclude test")
    process_runner.assert_not_has_call(
        "1.2.3.4", pattern="--filter dir-merge,- .ignore"
    )


@pytest.mark.parametrize("run_option_type", ["run_options", "head_run_options"])
def test_docker_shm_override(run_option_type):
    process_runner = MockProcessRunner()
    provider = MockProvider()
    provider.create_node({}, {}, 1)
    cluster_name = "cluster"

    docker_config = {
        "container_name": "container",
        "image": "rayproject/ray:latest",
        run_option_type: ["--shm-size=80g"],
    }
    args = {
        "log_prefix": "prefix",
        "node_id": "0",
        "provider": provider,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": False,
        "docker_config": docker_config,
    }
    cmd_runner = DockerCommandRunner(**args)

    process_runner.respond_to_call("json .Config.Env", 2 * ["[]"])
    cmd_runner.run_init(as_head=True, file_mounts={}, sync_run_yet=True)

    # Ensure the user-provided SHM size is used.
    process_runner.assert_has_call("1.2.3.4", pattern="--shm-size=80g")

    # Ensure that SHM auto detection is bypassed
    process_runner.assert_not_has_call("1.2.3.4", pattern="/proc/meminfo")


if __name__ == "__main__":
    sys.exit(pytest.main(["-sv", __file__]))