import json
import logging
import sys
from threading import RLock
from typing import Any, Dict, Optional

import requests

from ray._common.network_utils import build_address
from ray.autoscaler.node_launch_exception import NodeLaunchException
from ray.autoscaler.node_provider import NodeProvider
from ray.autoscaler.tags import (
    NODE_KIND_HEAD,
    NODE_KIND_WORKER,
    STATUS_SETTING_UP,
    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__)

HEAD_NODE_ID = 0
HEAD_NODE_TYPE = "ray.head.default"


class SparkNodeProvider(NodeProvider):
    """A node provider that implements provider for nodes of Ray on spark."""

    def __init__(self, provider_config, cluster_name):
        NodeProvider.__init__(self, provider_config, cluster_name)
        self.lock = RLock()

        self._nodes = {
            str(HEAD_NODE_ID): {
                "tags": {
                    TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
                    TAG_RAY_USER_NODE_TYPE: HEAD_NODE_TYPE,
                    TAG_RAY_NODE_NAME: HEAD_NODE_ID,
                    TAG_RAY_NODE_STATUS: STATUS_UP_TO_DATE,
                }
            },
        }
        self._next_node_id = 0

        self.ray_head_ip = self.provider_config["ray_head_ip"]
        # The port of spark job server. We send http request to spark job server
        # to launch spark jobs, ray worker nodes are launched by spark task in
        # spark jobs.
        spark_job_server_port = self.provider_config["spark_job_server_port"]
        self.spark_job_server_url = (
            f"http://{build_address(self.ray_head_ip, spark_job_server_port)}"
        )
        self.ray_head_port = self.provider_config["ray_head_port"]
        # The unique id for the Ray on spark cluster.
        self.cluster_id = self.provider_config["cluster_unique_id"]

    def get_next_node_id(self):
        with self.lock:
            self._next_node_id += 1
            return self._next_node_id

    def non_terminated_nodes(self, tag_filters):
        with self.lock:
            nodes = []

            died_nodes = []
            for node_id in self._nodes:
                if node_id == str(HEAD_NODE_ID):
                    status = "running"
                else:
                    status = self._query_node_status(node_id)

                if status == "running":
                    if (
                        self._nodes[node_id]["tags"][TAG_RAY_NODE_STATUS]
                        == STATUS_SETTING_UP
                    ):
                        self._nodes[node_id]["tags"][
                            TAG_RAY_NODE_STATUS
                        ] = STATUS_UP_TO_DATE
                        logger.info(
                            f"Spark node provider node {node_id} starts running."
                        )

                if status == "terminated":
                    died_nodes.append(node_id)
                else:
                    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)

            for died_node_id in died_nodes:
                self._nodes.pop(died_node_id)

            return nodes

    def _query_node_status(self, node_id):
        spark_job_group_id = self._gen_spark_job_group_id(node_id)

        response = requests.post(
            url=self.spark_job_server_url + "/query_task_status",
            json={"spark_job_group_id": spark_job_group_id},
        )
        response.raise_for_status()

        decoded_resp = response.content.decode("utf-8")
        json_res = json.loads(decoded_resp)
        return json_res["status"]

    def is_running(self, node_id):
        with self.lock:
            return (
                node_id in self._nodes
                and self._nodes[node_id]["tags"][TAG_RAY_NODE_STATUS]
                == STATUS_UP_TO_DATE
            )

    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):
        assert node_id in self._nodes
        self._nodes[node_id]["tags"].update(tags)

    def create_node(
        self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
    ) -> Optional[Dict[str, Any]]:
        raise AssertionError("This method should not be called.")

    def _gen_spark_job_group_id(self, node_id):
        return (
            f"ray-cluster-{self.ray_head_port}-{self.cluster_id}"
            f"-worker-node-{node_id}"
        )

    def create_node_with_resources_and_labels(
        self, node_config, tags, count, resources, labels
    ):
        for _ in range(count):
            self._create_node_with_resources_and_labels(
                node_config, tags, resources, labels
            )

    def _create_node_with_resources_and_labels(
        self, node_config, tags, resources, labels
    ):
        from ray.util.spark.cluster_init import _append_resources_config

        with self.lock:
            resources = resources.copy()
            node_type = tags[TAG_RAY_USER_NODE_TYPE]
            # NOTE:
            #  "NODE_ID_AS_RESOURCE" value must be an integer,
            #  but `node_id` used by autoscaler must be a string.
            node_id = str(self.get_next_node_id())
            resources["NODE_ID_AS_RESOURCE"] = int(node_id)

            conf = self.provider_config.copy()

            num_cpus_per_node = resources.pop("CPU")
            num_gpus_per_node = resources.pop("GPU")
            heap_memory_per_node = resources.pop("memory")
            object_store_memory_per_node = resources.pop("object_store_memory")

            conf["worker_node_options"] = _append_resources_config(
                conf["worker_node_options"], resources
            )
            response = requests.post(
                url=self.spark_job_server_url + "/create_node",
                json={
                    "spark_job_group_id": self._gen_spark_job_group_id(node_id),
                    "spark_job_group_desc": (
                        "This job group is for spark job which runs the Ray "
                        f"cluster worker node {node_id} connecting to ray "
                        f"head node {build_address(self.ray_head_ip, self.ray_head_port)}"
                    ),
                    "using_stage_scheduling": conf["using_stage_scheduling"],
                    "ray_head_ip": self.ray_head_ip,
                    "ray_head_port": self.ray_head_port,
                    "ray_temp_dir": conf["ray_temp_dir"],
                    "num_cpus_per_node": num_cpus_per_node,
                    "num_gpus_per_node": num_gpus_per_node,
                    "heap_memory_per_node": heap_memory_per_node,
                    "object_store_memory_per_node": object_store_memory_per_node,
                    "worker_node_options": conf["worker_node_options"],
                    "collect_log_to_path": conf["collect_log_to_path"],
                    "node_id": resources["NODE_ID_AS_RESOURCE"],
                },
            )

            try:
                # Spark job server is locally launched, if spark job server request
                # failed, it is unlikely network error but probably unrecoverable
                # error, so we make it fast-fail.
                response.raise_for_status()
            except Exception:
                raise NodeLaunchException(
                    "Node creation failure",
                    f"Starting ray worker node {node_id} failed",
                    sys.exc_info(),
                )

            self._nodes[node_id] = {
                "tags": {
                    TAG_RAY_NODE_KIND: NODE_KIND_WORKER,
                    TAG_RAY_USER_NODE_TYPE: node_type,
                    TAG_RAY_NODE_NAME: node_id,
                    TAG_RAY_NODE_STATUS: STATUS_SETTING_UP,
                },
            }
            logger.info(f"Spark node provider creates node {node_id}.")

    def terminate_node(self, node_id):
        if node_id in self._nodes:
            response = requests.post(
                url=self.spark_job_server_url + "/terminate_node",
                json={"spark_job_group_id": self._gen_spark_job_group_id(node_id)},
            )
            response.raise_for_status()

        with self.lock:
            if node_id in self._nodes:
                self._nodes.pop(node_id)

        logger.info(f"Spark node provider terminates node {node_id}")

    @staticmethod
    def bootstrap_config(cluster_config):
        return cluster_config