from enum import Enum
from typing import Dict, Optional

from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag


class ServeUsageTag(Enum):
    API_VERSION = TagKey.SERVE_API_VERSION
    NUM_DEPLOYMENTS = TagKey.SERVE_NUM_DEPLOYMENTS
    GCS_STORAGE = TagKey.GCS_STORAGE
    NUM_GPU_DEPLOYMENTS = TagKey.SERVE_NUM_GPU_DEPLOYMENTS
    FASTAPI_USED = TagKey.SERVE_FASTAPI_USED
    DAG_DRIVER_USED = TagKey.SERVE_DAG_DRIVER_USED
    HTTP_ADAPTER_USED = TagKey.SERVE_HTTP_ADAPTER_USED
    GRPC_INGRESS_USED = TagKey.SERVE_GRPC_INGRESS_USED
    REST_API_VERSION = TagKey.SERVE_REST_API_VERSION
    NUM_APPS = TagKey.SERVE_NUM_APPS
    NUM_REPLICAS_LIGHTWEIGHT_UPDATED = TagKey.SERVE_NUM_REPLICAS_LIGHTWEIGHT_UPDATED
    USER_CONFIG_LIGHTWEIGHT_UPDATED = TagKey.SERVE_USER_CONFIG_LIGHTWEIGHT_UPDATED
    AUTOSCALING_CONFIG_LIGHTWEIGHT_UPDATED = (
        TagKey.SERVE_AUTOSCALING_CONFIG_LIGHTWEIGHT_UPDATED
    )
    DEPLOYMENT_HANDLE_API_USED = TagKey.SERVE_DEPLOYMENT_HANDLE_API_USED
    DEPLOYMENT_HANDLE_TO_OBJECT_REF_API_USED = (
        TagKey.SERVE_DEPLOYMENT_HANDLE_TO_OBJECT_REF_API_USED
    )
    MULTIPLEXED_API_USED = TagKey.SERVE_MULTIPLEXED_API_USED
    HTTP_PROXY_USED = TagKey.SERVE_HTTP_PROXY_USED
    GRPC_PROXY_USED = TagKey.SERVE_GRPC_PROXY_USED
    SERVE_STATUS_API_USED = TagKey.SERVE_STATUS_API_USED
    SERVE_GET_APP_HANDLE_API_USED = TagKey.SERVE_GET_APP_HANDLE_API_USED
    SERVE_GET_DEPLOYMENT_HANDLE_API_USED = TagKey.SERVE_GET_DEPLOYMENT_HANDLE_API_USED
    APP_CONTAINER_RUNTIME_ENV_USED = TagKey.SERVE_APP_CONTAINER_RUNTIME_ENV_USED
    DEPLOYMENT_CONTAINER_RUNTIME_ENV_USED = (
        TagKey.SERVE_DEPLOYMENT_CONTAINER_RUNTIME_ENV_USED
    )
    NUM_NODE_COMPACTIONS = TagKey.SERVE_NUM_NODE_COMPACTIONS
    AUTO_NUM_REPLICAS_USED = TagKey.SERVE_AUTO_NUM_REPLICAS_USED
    CUSTOM_REQUEST_ROUTER_USED = TagKey.SERVE_CUSTOM_REQUEST_ROUTER_USED
    NUM_REPLICAS_VIA_API_CALL_UPDATED = TagKey.SERVE_NUM_REPLICAS_VIA_API_CALL_UPDATED
    NUM_REPLICAS_USING_ASYNCHRONOUS_INFERENCE = (
        TagKey.SERVE_NUM_REPLICAS_USING_ASYNCHRONOUS_INFERENCE
    )
    CUSTOM_AUTOSCALING_POLICY_USED = TagKey.SERVE_CUSTOM_AUTOSCALING_POLICY_USED

    def record(self, value: str):
        """Record telemetry value."""
        record_extra_usage_tag(self.value, value)

    def get_value_from_report(self, report: Dict) -> Optional[str]:
        """Returns `None` if the tag isn't in the report."""
        if "extra_usage_tags" not in report:
            return None

        return report["extra_usage_tags"].get(TagKey.Name(self.value).lower(), None)