from collections import defaultdict
from typing import Dict, List

from ray.autoscaler._private.prom_metrics import AutoscalerPrometheusMetrics
from ray.autoscaler.v2.instance_manager.common import InstanceUtil
from ray.autoscaler.v2.instance_manager.config import NodeTypeConfig
from ray.autoscaler.v2.schema import NodeType
from ray.core.generated.instance_manager_pb2 import Instance as IMInstance


class AutoscalerMetricsReporter:
    def __init__(self, prom_metrics: AutoscalerPrometheusMetrics) -> None:
        self._prom_metrics = prom_metrics

    def inc_stopped_nodes(self, count: int) -> None:
        self._prom_metrics.stopped_nodes.inc(count)

    def report_instances(
        self,
        instances: List[IMInstance],
        node_type_configs: Dict[NodeType, NodeTypeConfig],
    ):
        """
        Record autoscaler metrics for:
            - pending_nodes: Nodes that are launching/pending ray start
            - active_nodes: Active nodes (nodes running ray)
            - recently_failed_nodes: Nodes that are being terminated.
        """
        # map of instance type to a dict of status to count.
        status_count_by_type: Dict[NodeType, Dict[str, int]] = {}

        def _new_status_count() -> Dict[str, int]:
            return {
                "pending": 0,
                "running": 0,
                "terminating": 0,
                "terminated": 0,
            }

        # initialize the status count by type.
        for instance_type in node_type_configs.keys():
            status_count_by_type[instance_type] = _new_status_count()

        for instance in instances:
            status_count = status_count_by_type.get(instance.instance_type)
            if status_count is None:
                status_count = _new_status_count()
                status_count_by_type[instance.instance_type] = status_count
            if InstanceUtil.is_ray_pending(instance.status):
                status_count["pending"] += 1
            elif InstanceUtil.is_ray_running(instance.status):
                status_count["running"] += 1
            elif instance.status == IMInstance.TERMINATING:
                status_count["terminating"] += 1
            elif instance.status == IMInstance.TERMINATED:
                status_count["terminated"] += 1

        for instance_type, status_count in status_count_by_type.items():
            self._prom_metrics.pending_nodes.labels(
                SessionName=self._prom_metrics.session_name, NodeType=instance_type
            ).set(status_count["pending"])

            self._prom_metrics.active_nodes.labels(
                SessionName=self._prom_metrics.session_name, NodeType=instance_type
            ).set(status_count["running"])

            self._prom_metrics.recently_failed_nodes.labels(
                SessionName=self._prom_metrics.session_name, NodeType=instance_type
            ).set(status_count["terminating"])

    def report_resources(
        self,
        instances: List[IMInstance],
        node_type_configs: Dict[NodeType, NodeTypeConfig],
    ):
        """
        Record autoscaler metrics for:
            - pending_resources: Pending resources
            - cluster_resources: Cluster resources (resources running on the cluster)
        """
        # pending resources.
        pending_resources = defaultdict(float)
        cluster_resources = defaultdict(float)

        def _add_resources(resource_map, node_type_configs, node_type, count):
            node_resources = node_type_configs[node_type].resources
            for resource_name, resource_value in node_resources.items():
                resource_map[resource_name] += resource_value * count

        for instance in instances:
            if instance.instance_type not in node_type_configs:
                continue
            if InstanceUtil.is_ray_pending(instance.status):
                _add_resources(
                    pending_resources, node_type_configs, instance.instance_type, 1
                )
            elif InstanceUtil.is_ray_running(instance.status):
                _add_resources(
                    cluster_resources, node_type_configs, instance.instance_type, 1
                )

        for resource_name, resource_value in pending_resources.items():
            self._prom_metrics.pending_resources.labels(
                SessionName=self._prom_metrics.session_name, resource=resource_name
            ).set(resource_value)

        for resource_name, resource_value in cluster_resources.items():
            self._prom_metrics.cluster_resources.labels(
                SessionName=self._prom_metrics.session_name, resource=resource_name
            ).set(resource_value)