import logging
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set

from ray.autoscaler._private.constants import (
    DISABLE_LAUNCH_CONFIG_CHECK_KEY,
    DISABLE_NODE_UPDATERS_KEY,
    FOREGROUND_NODE_LAUNCH_KEY,
)
from ray.autoscaler._private.util import NodeID, NodeIP, NodeKind, NodeStatus, NodeType
from ray.autoscaler.node_provider import NodeProvider
from ray.autoscaler.tags import (
    NODE_KIND_HEAD,
    TAG_RAY_NODE_KIND,
    TAG_RAY_NODE_STATUS,
    TAG_RAY_REPLICA_INDEX,
    TAG_RAY_USER_NODE_TYPE,
)

logger = logging.getLogger(__name__)


@dataclass
class ScaleRequest:
    """Stores desired scale computed by the autoscaler.

    Attributes:
        desired_num_workers: Map of worker NodeType to desired number of workers of
            that type.
        workers_to_delete: List of ids of nodes that should be removed.
    """

    desired_num_workers: Dict[NodeType, int] = field(default_factory=dict)
    workers_to_delete: Set[NodeID] = field(default_factory=set)


@dataclass
class NodeData:
    """Stores all data about a Ray node needed by the autoscaler.

    Attributes:
        kind: Whether the node is the head or a worker.
        type: The user-defined type of the node.
        replica_index: An identifier for nodes in a replica of a TPU worker group.
            This value is set as a Pod label by a GKE webhook when TPUs are requested
        ip: Cluster-internal ip of the node. ip can be None if the ip
            has not yet been assigned.
        status: The status of the node. You must adhere to the following semantics
            for status:
            * The status must be "up-to-date" if and only if the node is running.
            * The status must be "update-failed" if and only if the node is in an
                unknown or failed state.
            * If the node is in a pending (starting-up) state, the status should be
                a brief user-facing description of why the node is pending.
    """

    kind: NodeKind
    type: NodeType
    ip: Optional[NodeIP]
    status: NodeStatus
    replica_index: Optional[str] = None


class BatchingNodeProvider(NodeProvider):
    """Abstract subclass of NodeProvider meant for use with external cluster managers.

    Batches reads of cluster state into a single method, get_node_data, called at the
    start of an autoscaling update.

    Batches modifications to cluster state into a single method, submit_scale_request,
    called at the end of an autoscaling update.

    Implementing a concrete subclass of BatchingNodeProvider only requires overriding
    get_node_data() and submit_scale_request().

    See the method docstrings for more information.

    Note that an autoscaling update may be conditionally
    cancelled using the optional method safe_to_scale()
    of the root NodeProvider.
    """

    def __init__(
        self,
        provider_config: Dict[str, Any],
        cluster_name: str,
    ) -> None:
        NodeProvider.__init__(self, provider_config, cluster_name)
        self.node_data_dict: Dict[NodeID, NodeData] = {}

        # These flags enforce correct behavior for single-threaded node providers
        # which interact with external cluster managers:
        assert (
            provider_config.get(DISABLE_NODE_UPDATERS_KEY, False) is True
        ), f"To use BatchingNodeProvider, must set `{DISABLE_NODE_UPDATERS_KEY}:True`."
        assert provider_config.get(DISABLE_LAUNCH_CONFIG_CHECK_KEY, False) is True, (
            "To use BatchingNodeProvider, must set "
            f"`{DISABLE_LAUNCH_CONFIG_CHECK_KEY}:True`."
        )
        assert (
            provider_config.get(FOREGROUND_NODE_LAUNCH_KEY, False) is True
        ), f"To use BatchingNodeProvider, must set `{FOREGROUND_NODE_LAUNCH_KEY}:True`."

        # self.scale_change_needed tracks whether we need to update scale.
        # set to True in create_node and terminate_nodes calls
        # reset to False in non_terminated_nodes, which occurs at the start of the
        #   autoscaling update. For good measure, also set to false in post_process.
        self.scale_change_needed = False

        self.scale_request = ScaleRequest()

        # Initialize map of replica indices to nodes in that replica
        self.replica_index_to_nodes = defaultdict(list[str])

    def get_node_data(self) -> Dict[NodeID, NodeData]:
        """Queries cluster manager for node info. Returns a mapping from node id to
        NodeData.

        Each NodeData value must adhere to the semantics of the NodeData docstring.
        (Note in particular the requirements for NodeData.status.)

        Consistency requirement:
        If a node id was present in ScaleRequest.workers_to_delete of a previously
        submitted scale request, it should no longer be present as a key in
        get_node_data.
        (Node termination must be registered immediately when submit_scale_request
        returns.)
        """
        raise NotImplementedError

    def submit_scale_request(self, scale_request: ScaleRequest) -> None:
        """Tells the cluster manager which nodes to delete and how many nodes of
        each node type to maintain.

        Consistency requirement:
        If a node id was present in ScaleRequest.workers_to_delete of a previously
        submitted scale request, it should no longer be present as key in get_node_data.
        (Node termination must be registered immediately when submit_scale_request
        returns.)
        """
        raise NotImplementedError

    def post_process(self) -> None:
        """Submit a scale request if it is necessary to do so."""
        if self.scale_change_needed:
            self.submit_scale_request(self.scale_request)
        self.scale_change_needed = False

    def non_terminated_nodes(self, tag_filters: Dict[str, str]) -> List[str]:
        self.scale_change_needed = False
        self.node_data_dict = self.get_node_data()

        # Initialize ScaleRequest
        self.scale_request = ScaleRequest(
            desired_num_workers=self.cur_num_workers(),  # Current scale
            workers_to_delete=set(),  # No workers to delete yet
        )
        all_nodes = list(self.node_data_dict.keys())
        self.replica_index_to_nodes.clear()
        for node_id in all_nodes:
            replica_index = self.node_data_dict[node_id].replica_index
            # Only add node to map if it belongs to a multi-host podslice
            if replica_index is not None:
                self.replica_index_to_nodes[replica_index].append(node_id)
        # Support filtering by TAG_RAY_NODE_KIND, TAG_RAY_NODE_STATUS, and
        # TAG_RAY_USER_NODE_TYPE.
        # The autoscaler only uses tag_filters={},
        # but filtering by the these keys is useful for testing.
        filtered_nodes = [
            node
            for node in all_nodes
            if tag_filters.items() <= self.node_tags(node).items()
        ]
        return filtered_nodes

    def cur_num_workers(self):
        """Returns dict mapping node type to the number of nodes of that type."""
        # Factor like this for convenient re-use.
        return self._cur_num_workers(self.node_data_dict)

    def _cur_num_workers(self, node_data_dict: Dict[str, Any]):
        num_workers_dict = defaultdict(int)
        for node_data in node_data_dict.values():
            if node_data.kind == NODE_KIND_HEAD:
                # Only track workers.
                continue
            num_workers_dict[node_data.type] += 1
        return num_workers_dict

    def node_tags(self, node_id: str) -> Dict[str, str]:
        node_data = self.node_data_dict[node_id]
        tags = {
            TAG_RAY_NODE_KIND: node_data.kind,
            TAG_RAY_NODE_STATUS: node_data.status,
            TAG_RAY_USER_NODE_TYPE: node_data.type,
        }
        if node_data.replica_index is not None:
            tags[TAG_RAY_REPLICA_INDEX] = node_data.replica_index
        return tags

    def internal_ip(self, node_id: str) -> str:
        return self.node_data_dict[node_id].ip

    def create_node(
        self, node_config: Dict[str, Any], tags: Dict[str, str], count: int
    ) -> Optional[Dict[str, Any]]:
        node_type = tags[TAG_RAY_USER_NODE_TYPE]
        self.scale_request.desired_num_workers[node_type] += count
        self.scale_change_needed = True

    def terminate_node(self, node_id: str) -> Optional[Dict[str, Any]]:
        # Sanity check: We should never try to delete the same node twice.
        if node_id in self.scale_request.workers_to_delete:
            logger.warning(
                f"Autoscaler tried to terminate node {node_id} twice in the same update"
                ". Skipping termination request."
            )
            return

        # Sanity check: We should never try to delete a node we haven't seen.
        if node_id not in self.node_data_dict:
            logger.warning(
                f"Autoscaler tried to terminate unkown node {node_id}"
                ". Skipping termination request."
            )
            return

        node_type = self.node_data_dict[node_id].type

        # Sanity check: Don't request less than 0 nodes.
        if self.scale_request.desired_num_workers[node_type] <= 0:
            # This is logically impossible.
            raise AssertionError(
                "NodeProvider attempted to request less than 0 workers of type "
                f"{node_type}. Skipping termination request."
            )

        # Terminate node
        self.scale_request.desired_num_workers[node_type] -= 1
        self.scale_request.workers_to_delete.add(node_id)

        # Scale down all nodes in replica if node_id is part of a multi-host podslice
        tags = self.node_tags(node_id)
        if TAG_RAY_REPLICA_INDEX in tags:
            node_replica_index = tags[TAG_RAY_REPLICA_INDEX]
            for worker_id in self.replica_index_to_nodes[node_replica_index]:
                # Check if worker has already been scheduled to delete
                if worker_id not in self.scale_request.workers_to_delete:
                    self.scale_request.workers_to_delete.add(worker_id)
                    logger.info(
                        f"Autoscaler terminating node {worker_id} "
                        f"in multi-host replica {node_replica_index}."
                    )
        self.scale_change_needed = True