import copy
import json
import logging
import os
import subprocess
import sys
import time
from threading import RLock
from types import ModuleType
from typing import Any, Dict, Optional

import yaml

import ray
import ray._private.ray_constants as ray_constants
from ray.autoscaler._private.fake_multi_node.command_runner import (
    FakeDockerCommandRunner,
)
from ray.autoscaler.command_runner import CommandRunnerInterface
from ray.autoscaler.node_provider import NodeProvider
from ray.autoscaler.tags import (
    NODE_KIND_HEAD,
    NODE_KIND_WORKER,
    STATUS_UP_TO_DATE,
    TAG_RAY_NODE_KIND,
    TAG_RAY_NODE_NAME,
    TAG_RAY_NODE_STATUS,
    TAG_RAY_USER_NODE_TYPE,
)

logger = logging.getLogger(__name__)

# We generate the node ids deterministically in the fake node provider, so that
# we can associate launched nodes with their resource reports. IDs increment
# starting with fffff*00000 for the head node, fffff*00001, etc. for workers.
FAKE_HEAD_NODE_ID = "fffffffffffffffffffffffffffffffffffffffffffffffffff00000"
FAKE_HEAD_NODE_TYPE = "ray.head.default"

FAKE_DOCKER_DEFAULT_GCS_PORT = 16379
FAKE_DOCKER_DEFAULT_OBJECT_MANAGER_PORT = 18076
FAKE_DOCKER_DEFAULT_CLIENT_PORT = 10002

DOCKER_COMPOSE_SKELETON = {
    "version": "3.9",
    "services": {},
    "networks": {"ray_local": {}},
}

DOCKER_NODE_SKELETON = {
    "networks": ["ray_local"],
    "mem_limit": "3000m",
    "mem_reservation": "3000m",
    "shm_size": "1200m",
    "volumes": [],
}

DOCKER_HEAD_CMD = (
    'bash -c "'
    "sudo mkdir -p {volume_dir} && "
    "sudo chmod 777 {volume_dir} && "
    "touch {volume_dir}/.in_docker && "
    "sudo chown -R ray:users /cluster/node && "
    "sudo chmod -R 777 /cluster/node && "
    "sudo chown -R ray:users /cluster/shared && "
    "sudo chmod -R 777 /cluster/shared && "
    "sudo chmod 700 ~/.ssh && "
    "sudo chmod 600 ~/.ssh/authorized_keys && "
    "sudo chmod 600 ~/ray_bootstrap_key.pem && "
    "sudo chown ray:users "
    "~/.ssh ~/.ssh/authorized_keys ~/ray_bootstrap_key.pem && "
    "{ensure_ssh} && "
    "sleep 1 && "
    "RAY_FAKE_CLUSTER=1 ray start --head "
    "--autoscaling-config=~/ray_bootstrap_config.yaml "
    "--object-manager-port=8076 "
    "--num-cpus {num_cpus} "
    "--num-gpus {num_gpus} "
    # "--resources='{resources}' "
    '--block"'
)

DOCKER_WORKER_CMD = (
    'bash -c "'
    "sudo mkdir -p {volume_dir} && "
    "sudo chmod 777 {volume_dir} && "
    "touch {volume_dir}/.in_docker && "
    "sudo chown -R ray:users /cluster/node && "
    "sudo chmod -R 777 /cluster/node && "
    "sudo chmod 700 ~/.ssh && "
    "sudo chmod 600 ~/.ssh/authorized_keys && "
    "sudo chown ray:users ~/.ssh ~/.ssh/authorized_keys && "
    "{ensure_ssh} && "
    "sleep 1 && "
    f"ray start --address={FAKE_HEAD_NODE_ID}:6379 "
    "--object-manager-port=8076 "
    "--num-cpus {num_cpus} "
    "--num-gpus {num_gpus} "
    # "--resources='{resources}' "
    '--block"'
)


def host_dir(container_dir: str):
    """Replace local dir with potentially different host dir.

    E.g. in docker-in-docker environments, the host dir might be
    different to the mounted directory in the container.

    This method will do a simple global replace to adjust the paths.
    """
    ray_tempdir = os.environ.get("RAY_TEMPDIR", None)
    ray_hostdir = os.environ.get("RAY_HOSTDIR", None)

    if not ray_tempdir or not ray_hostdir:
        return container_dir

    return container_dir.replace(ray_tempdir, ray_hostdir)


def create_node_spec(
    head: bool,
    docker_image: str,
    mounted_cluster_dir: str,
    mounted_node_dir: str,
    num_cpus: int = 2,
    num_gpus: int = 0,
    resources: Optional[Dict] = None,
    env_vars: Optional[Dict] = None,
    host_gcs_port: int = 16379,
    host_object_manager_port: int = 18076,
    host_client_port: int = 10002,
    volume_dir: Optional[str] = None,
    node_state_path: Optional[str] = None,
    docker_status_path: Optional[str] = None,
    docker_compose_path: Optional[str] = None,
    bootstrap_config_path: Optional[str] = None,
    private_key_path: Optional[str] = None,
    public_key_path: Optional[str] = None,
):
    node_spec = copy.deepcopy(DOCKER_NODE_SKELETON)
    node_spec["image"] = docker_image

    bootstrap_cfg_path_on_container = "/home/ray/ray_bootstrap_config.yaml"
    bootstrap_key_path_on_container = "/home/ray/ray_bootstrap_key.pem"

    resources = resources or {}

    ensure_ssh = (
        (
            "((sudo apt update && sudo apt install -y openssh-server && "
            "sudo service ssh start) || true)"
        )
        if not bool(int(os.environ.get("RAY_HAS_SSH", "0") or "0"))
        else "sudo service ssh start"
    )

    cmd_kwargs = dict(
        ensure_ssh=ensure_ssh,
        num_cpus=num_cpus,
        num_gpus=num_gpus,
        resources=json.dumps(resources, indent=None),
        volume_dir=volume_dir,
        autoscaling_config=bootstrap_cfg_path_on_container,
    )

    env_vars = env_vars or {}

    # Set to "auto" to mount current autoscaler directory to nodes for dev
    fake_cluster_dev_dir = os.environ.get("FAKE_CLUSTER_DEV", "")
    if fake_cluster_dev_dir:
        if fake_cluster_dev_dir == "auto":
            local_ray_dir = os.path.dirname(ray.__file__)
        else:
            local_ray_dir = fake_cluster_dev_dir
        os.environ["FAKE_CLUSTER_DEV"] = local_ray_dir

        mj = sys.version_info.major
        mi = sys.version_info.minor

        fake_modules_str = os.environ.get("FAKE_CLUSTER_DEV_MODULES", "autoscaler")
        fake_modules = fake_modules_str.split(",")

        docker_ray_dir = f"/home/ray/anaconda3/lib/python{mj}.{mi}/site-packages/ray"

        node_spec["volumes"] += [
            f"{local_ray_dir}/{module}:{docker_ray_dir}/{module}:ro"
            for module in fake_modules
        ]
        env_vars["FAKE_CLUSTER_DEV"] = local_ray_dir
        env_vars["FAKE_CLUSTER_DEV_MODULES"] = fake_modules_str
        os.environ["FAKE_CLUSTER_DEV_MODULES"] = fake_modules_str

    if head:
        node_spec["command"] = DOCKER_HEAD_CMD.format(**cmd_kwargs)
        # Expose ports so we can connect to the cluster from outside
        node_spec["ports"] = [
            f"{host_gcs_port}:{ray_constants.DEFAULT_PORT}",
            f"{host_object_manager_port}:8076",
            f"{host_client_port}:10001",
        ]
        # Mount status and config files for the head node
        node_spec["volumes"] += [
            f"{host_dir(node_state_path)}:{node_state_path}",
            f"{host_dir(docker_status_path)}:{docker_status_path}",
            f"{host_dir(docker_compose_path)}:{docker_compose_path}",
            f"{host_dir(bootstrap_config_path)}:" f"{bootstrap_cfg_path_on_container}",
            f"{host_dir(private_key_path)}:{bootstrap_key_path_on_container}",
        ]

        # Create file if it does not exist on local filesystem
        for filename in [node_state_path, docker_status_path, bootstrap_config_path]:
            if not os.path.exists(filename):
                with open(filename, "wt") as f:
                    f.write("{}")
    else:
        node_spec["command"] = DOCKER_WORKER_CMD.format(**cmd_kwargs)
        node_spec["depends_on"] = [FAKE_HEAD_NODE_ID]

    # Mount shared directories and ssh access keys
    node_spec["volumes"] += [
        f"{host_dir(mounted_cluster_dir)}:/cluster/shared",
        f"{host_dir(mounted_node_dir)}:/cluster/node",
        f"{host_dir(public_key_path)}:/home/ray/.ssh/authorized_keys",
    ]

    # Pass these environment variables (to the head node)
    # These variables are propagated by the `docker compose` command.
    env_vars.setdefault("RAY_HAS_SSH", os.environ.get("RAY_HAS_SSH", ""))
    env_vars.setdefault("RAY_TEMPDIR", os.environ.get("RAY_TEMPDIR", ""))
    env_vars.setdefault("RAY_HOSTDIR", os.environ.get("RAY_HOSTDIR", ""))

    node_spec["environment"] = [f"{k}={v}" for k, v in env_vars.items()]

    return node_spec


class FakeMultiNodeProvider(NodeProvider):
    """A node provider that implements multi-node on a single machine.

    This is used for laptop mode testing of autoscaling functionality."""

    def __init__(self, provider_config, cluster_name):
        NodeProvider.__init__(self, provider_config, cluster_name)
        self.lock = RLock()
        if "RAY_FAKE_CLUSTER" not in os.environ:
            raise RuntimeError(
                "FakeMultiNodeProvider requires ray to be started with "
                "RAY_FAKE_CLUSTER=1 ray start ..."
            )

        self._nodes = {
            FAKE_HEAD_NODE_ID: {
                "tags": {
                    TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
                    TAG_RAY_USER_NODE_TYPE: FAKE_HEAD_NODE_TYPE,
                    TAG_RAY_NODE_NAME: FAKE_HEAD_NODE_ID,
                    TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
                }
            },
        }
        self._next_node_id = 0

    def _next_hex_node_id(self):
        self._next_node_id += 1
        base = "fffffffffffffffffffffffffffffffffffffffffffffffffff"
        return base + str(self._next_node_id).zfill(5)

    def non_terminated_nodes(self, tag_filters):
        with self.lock:
            nodes = []
            for node_id in self._nodes:
                tags = self.node_tags(node_id)
                ok = True
                for k, v in tag_filters.items():
                    if tags.get(k) != v:
                        ok = False
                if ok:
                    nodes.append(node_id)

            return nodes

    def is_running(self, node_id):
        with self.lock:
            return node_id in self._nodes

    def is_terminated(self, node_id):
        with self.lock:
            return node_id not in self._nodes

    def node_tags(self, node_id):
        with self.lock:
            return self._nodes[node_id]["tags"]

    def _get_ip(self, node_id: str) -> Optional[str]:
        return node_id

    def external_ip(self, node_id):
        return self._get_ip(node_id)

    def internal_ip(self, node_id):
        return self._get_ip(node_id)

    def set_node_tags(self, node_id, tags):
        raise AssertionError("Readonly node provider cannot be updated")

    def create_node_with_resources_and_labels(
        self, node_config, tags, count, resources, labels
    ):
        with self.lock:
            node_type = tags[TAG_RAY_USER_NODE_TYPE]
            next_id = self._next_hex_node_id()
            ray_params = ray._private.parameter.RayParams(
                min_worker_port=0,
                max_worker_port=0,
                dashboard_port=None,
                num_cpus=resources.pop("CPU", 0),
                num_gpus=resources.pop("GPU", 0),
                object_store_memory=resources.pop("object_store_memory", None),
                resources=resources,
                labels=labels,
                redis_address="{}:6379".format(
                    ray._private.services.get_node_ip_address()
                ),
                gcs_address="{}:6379".format(
                    ray._private.services.get_node_ip_address()
                ),
                env_vars={
                    "RAY_OVERRIDE_NODE_ID_FOR_TESTING": next_id,
                    ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: json.dumps(resources),
                    ray_constants.LABELS_ENVIRONMENT_VARIABLE: json.dumps(labels),
                },
            )
            node = ray._private.node.Node(
                ray_params, head=False, shutdown_at_exit=False, spawn_reaper=False
            )
            self._nodes[next_id] = {
                "tags": {
                    TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
                    TAG_RAY_USER_NODE_TYPE: node_type,
                    TAG_RAY_NODE_NAME: next_id,
                    TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
                },
                "node": node,
            }

    def terminate_node(self, node_id):
        with self.lock:
            try:
                node = self._nodes.pop(node_id)
            except Exception as e:
                raise e

            self._terminate_node(node)

    def _terminate_node(self, node):
        node["node"].kill_all_processes(check_alive=False, allow_graceful=True)

    @staticmethod
    def bootstrap_config(cluster_config):
        return cluster_config


class FakeMultiNodeDockerProvider(FakeMultiNodeProvider):
    """A node provider that implements multi-node on a single machine.

    This is used for laptop mode testing of multi node functionality
    where each node has their own FS and IP."""

    def __init__(self, provider_config, cluster_name):
        super(FakeMultiNodeDockerProvider, self).__init__(provider_config, cluster_name)

        fake_head = copy.deepcopy(self._nodes)

        self._project_name = self.provider_config["project_name"]
        self._docker_image = self.provider_config["image"]

        self._host_gcs_port = self.provider_config.get(
            "host_gcs_port", FAKE_DOCKER_DEFAULT_GCS_PORT
        )
        self._host_object_manager_port = self.provider_config.get(
            "host_object_manager_port", FAKE_DOCKER_DEFAULT_OBJECT_MANAGER_PORT
        )
        self._host_client_port = self.provider_config.get(
            "host_client_port", FAKE_DOCKER_DEFAULT_CLIENT_PORT
        )

        self._head_resources = self.provider_config["head_resources"]

        # subdirs:
        #  - ./shared (shared filesystem)
        #  - ./nodes/<node_id> (node-specific mounted filesystem)
        self._volume_dir = self.provider_config["shared_volume_dir"]
        self._mounted_cluster_dir = os.path.join(self._volume_dir, "shared")

        if not self.in_docker_container:
            # Only needed on host
            os.makedirs(self._mounted_cluster_dir, mode=0o755, exist_ok=True)

        self._boostrap_config_path = os.path.join(
            self._volume_dir, "bootstrap_config.yaml"
        )

        self._private_key_path = os.path.join(self._volume_dir, "bootstrap_key.pem")
        self._public_key_path = os.path.join(self._volume_dir, "bootstrap_key.pem.pub")

        if not self.in_docker_container:
            # Create private key
            if not os.path.exists(self._private_key_path):
                subprocess.check_output(
                    f'ssh-keygen -b 2048 -t rsa -q -N "" '
                    f"-f {self._private_key_path}",
                    shell=True,
                )

            # Create public key
            if not os.path.exists(self._public_key_path):
                subprocess.check_output(
                    f"ssh-keygen -y "
                    f"-f {self._private_key_path} "
                    f"> {self._public_key_path}",
                    shell=True,
                )

        self._docker_compose_config_path = os.path.join(
            self._volume_dir, "docker-compose.yaml"
        )
        self._docker_compose_config = None

        self._node_state_path = os.path.join(self._volume_dir, "nodes.json")
        self._docker_status_path = os.path.join(self._volume_dir, "status.json")

        self._load_node_state()
        if FAKE_HEAD_NODE_ID not in self._nodes:
            # Reset
            self._nodes = copy.deepcopy(fake_head)

        self._nodes[FAKE_HEAD_NODE_ID][
            "node_spec"
        ] = self._create_node_spec_with_resources(
            head=True, node_id=FAKE_HEAD_NODE_ID, resources=self._head_resources
        )
        self._possibly_terminated_nodes = dict()

        self._cleanup_interval = provider_config.get("cleanup_interval", 9.5)

        self._docker_status = {}

        self._update_docker_compose_config()
        self._update_docker_status()
        self._save_node_state()

    @property
    def in_docker_container(self):
        return os.path.exists(os.path.join(self._volume_dir, ".in_docker"))

    def _create_node_spec_with_resources(
        self, head: bool, node_id: str, resources: Dict[str, Any]
    ):
        resources = resources.copy()

        # Create shared directory
        node_dir = os.path.join(self._volume_dir, "nodes", node_id)
        os.makedirs(node_dir, mode=0o777, exist_ok=True)

        resource_str = json.dumps(resources, indent=None)

        return create_node_spec(
            head=head,
            docker_image=self._docker_image,
            mounted_cluster_dir=self._mounted_cluster_dir,
            mounted_node_dir=node_dir,
            num_cpus=resources.pop("CPU", 0),
            num_gpus=resources.pop("GPU", 0),
            host_gcs_port=self._host_gcs_port,
            host_object_manager_port=self._host_object_manager_port,
            host_client_port=self._host_client_port,
            resources=resources,
            env_vars={
                "RAY_OVERRIDE_NODE_ID_FOR_TESTING": node_id,
                ray_constants.RESOURCES_ENVIRONMENT_VARIABLE: resource_str,
                **self.provider_config.get("env_vars", {}),
            },
            volume_dir=self._volume_dir,
            node_state_path=self._node_state_path,
            docker_status_path=self._docker_status_path,
            docker_compose_path=self._docker_compose_config_path,
            bootstrap_config_path=self._boostrap_config_path,
            public_key_path=self._public_key_path,
            private_key_path=self._private_key_path,
        )

    def _load_node_state(self) -> bool:
        if not os.path.exists(self._node_state_path):
            return False
        try:
            with open(self._node_state_path, "rt") as f:
                nodes = json.load(f)
        except Exception:
            return False
        if not nodes:
            return False
        self._nodes = nodes
        return True

    def _save_node_state(self):
        with open(self._node_state_path, "wt") as f:
            json.dump(self._nodes, f)

        # Make sure this is always writeable from inside the containers
        if not self.in_docker_container:
            # Only chmod from the outer container
            os.chmod(self._node_state_path, 0o777)

    def _update_docker_compose_config(self):
        config = copy.deepcopy(DOCKER_COMPOSE_SKELETON)
        config["services"] = {}
        for node_id, node in self._nodes.items():
            config["services"][node_id] = node["node_spec"]

        with open(self._docker_compose_config_path, "wt") as f:
            yaml.safe_dump(config, f)

    def _update_docker_status(self):
        if not os.path.exists(self._docker_status_path):
            return
        with open(self._docker_status_path, "rt") as f:
            self._docker_status = json.load(f)

    def _update_nodes(self):
        for node_id in list(self._nodes):
            if not self._is_docker_running(node_id):
                self._possibly_terminated_nodes.setdefault(node_id, time.monotonic())
            else:
                self._possibly_terminated_nodes.pop(node_id, None)
        self._cleanup_nodes()

    def _cleanup_nodes(self):
        for node_id, timestamp in list(self._possibly_terminated_nodes.items()):
            if time.monotonic() > timestamp + self._cleanup_interval:
                if not self._is_docker_running(node_id):
                    self._nodes.pop(node_id, None)
                self._possibly_terminated_nodes.pop(node_id, None)
        self._save_node_state()

    def _container_name(self, node_id):
        node_status = self._docker_status.get(node_id, {})
        timeout = time.monotonic() + 60
        while not node_status:
            if time.monotonic() > timeout:
                raise RuntimeError(f"Container for {node_id} never became available.")
            time.sleep(1)
            self._update_docker_status()
            node_status = self._docker_status.get(node_id, {})

        return node_status["Name"]

    def _is_docker_running(self, node_id):
        self._update_docker_status()

        return self._docker_status.get(node_id, {}).get("State", None) == "running"

    def non_terminated_nodes(self, tag_filters):
        self._update_nodes()
        return super(FakeMultiNodeDockerProvider, self).non_terminated_nodes(
            tag_filters
        )

    def is_running(self, node_id):
        with self.lock:
            self._update_nodes()

            return node_id in self._nodes and self._is_docker_running(node_id)

    def is_terminated(self, node_id):
        with self.lock:
            self._update_nodes()

            return node_id not in self._nodes and not self._is_docker_running(node_id)

    def get_command_runner(
        self,
        log_prefix: str,
        node_id: str,
        auth_config: Dict[str, Any],
        cluster_name: str,
        process_runner: ModuleType,
        use_internal_ip: bool,
        docker_config: Optional[Dict[str, Any]] = None,
    ) -> CommandRunnerInterface:
        if self.in_docker_container:
            return super(FakeMultiNodeProvider, self).get_command_runner(
                log_prefix,
                node_id,
                auth_config,
                cluster_name,
                process_runner,
                use_internal_ip,
            )

        # Else, host command runner:
        common_args = {
            "log_prefix": log_prefix,
            "node_id": node_id,
            "provider": self,
            "auth_config": auth_config,
            "cluster_name": cluster_name,
            "process_runner": process_runner,
            "use_internal_ip": use_internal_ip,
        }

        docker_config["container_name"] = self._container_name(node_id)
        docker_config["image"] = self._docker_image

        return FakeDockerCommandRunner(docker_config, **common_args)

    def _get_ip(self, node_id: str) -> Optional[str]:
        for i in range(3):
            self._update_docker_status()
            ip = self._docker_status.get(node_id, {}).get("IP", None)
            if ip:
                return ip
            time.sleep(3)
        return None

    def set_node_tags(self, node_id, tags):
        assert node_id in self._nodes
        self._nodes[node_id]["tags"].update(tags)

    def create_node_with_resources_and_labels(
        self, node_config, tags, count, resources, labels
    ):
        with self.lock:
            is_head = tags[TAG_RAY_NODE_KIND] == NODE_KIND_HEAD

            if is_head:
                next_id = FAKE_HEAD_NODE_ID
            else:
                next_id = self._next_hex_node_id()

            self._nodes[next_id] = {
                "tags": tags,
                "node_spec": self._create_node_spec_with_resources(
                    head=is_head, node_id=next_id, resources=resources
                ),
            }
            self._update_docker_compose_config()
            self._save_node_state()

    def create_node(
        self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
    ) -> Optional[Dict[str, Any]]:
        resources = self._head_resources
        return self.create_node_with_resources_and_labels(
            node_config, tags, count, resources, {}
        )

    def _terminate_node(self, node):
        self._update_docker_compose_config()
        self._save_node_state()

    @staticmethod
    def bootstrap_config(cluster_config):
        return cluster_config
