import copy
import os
from typing import Any, Dict

import yaml

import ray
from ray.autoscaler._private.aws.cloudwatch.cloudwatch_helper import CloudwatchHelper
from ray.autoscaler._private.aws.node_provider import AWSNodeProvider
from ray.autoscaler._private.commands import prepare_config, validate_config
from ray.autoscaler.tags import (
    NODE_KIND_HEAD,
    NODE_KIND_WORKER,
    TAG_RAY_CLUSTER_NAME,
    TAG_RAY_NODE_KIND,
    TAG_RAY_USER_NODE_TYPE,
)
from ray.tests.aws.utils.constants import (
    DEFAULT_CLUSTER_NAME,
    DEFAULT_NODE_PROVIDER_INSTANCE_TAGS,
)


def get_aws_example_config_file_path(file_name):
    import ray.autoscaler.aws

    return os.path.join(os.path.dirname(ray.autoscaler.aws.__file__), file_name)


def load_aws_example_config_file(file_name):
    config_file_path = get_aws_example_config_file_path(file_name)
    return yaml.safe_load(open(config_file_path).read())


def fake_fillout_available_node_types_resources(config: Dict[str, Any]) -> None:
    """A cheap way to fill out the resources field (the same way a node
    provider would autodetect them) as far as schema validation is concerned."""
    available_node_types = config.get("available_node_types", {})
    for label, value in available_node_types.items():
        value["resources"] = value.get("resources", {"filler": 1})


def bootstrap_aws_config(config):
    config = prepare_config(config)
    fake_fillout_available_node_types_resources(config)
    validate_config(config)
    config["cluster_name"] = DEFAULT_CLUSTER_NAME
    return ray.autoscaler._private.aws.config.bootstrap_aws(config)


def bootstrap_aws_example_config_file(file_name):
    config = load_aws_example_config_file(file_name)
    return bootstrap_aws_config(config)


def node_provider_tags(config: dict, type_name: str) -> dict:
    """
    Returns a copy of DEFAULT_NODE_PROVIDER_INSTANCE_TAGS with the Ray node
    kind and Ray user node type filled in from the input config and node type
    name.

    Args:
        config: autoscaler config
        type_name: node type name
    Returns:
        tags: node provider tags
    """
    tags = copy.copy(DEFAULT_NODE_PROVIDER_INSTANCE_TAGS)
    head_name = config["head_node_type"]
    node_kind = NODE_KIND_HEAD if type_name is head_name else NODE_KIND_WORKER
    tags[TAG_RAY_NODE_KIND] = node_kind
    tags[TAG_RAY_USER_NODE_TYPE] = type_name
    return tags


def apply_node_provider_config_updates(
    config: dict, node_cfg: dict, node_type_name: str, max_count: int
) -> None:
    """
    Applies default updates made by AWSNodeProvider to node_cfg during node
    creation. This should only be used for testing purposes.

    Args:
        config: autoscaler config
        node_cfg: node config
        node_type_name: node type name
        max_count: max nodes of the given type to launch
    """
    tags = node_provider_tags(config, node_type_name)
    tags[TAG_RAY_CLUSTER_NAME] = DEFAULT_CLUSTER_NAME
    user_tag_specs = node_cfg.get("TagSpecifications", [])
    tag_specs = [
        {
            "ResourceType": "instance",
            "Tags": [{"Key": k, "Value": v} for k, v in sorted(tags.items())],
        }
    ]
    node_provider_cfg_updates = {
        "MinCount": 1,
        "MaxCount": max_count,
        "TagSpecifications": tag_specs,
    }
    tags.pop(TAG_RAY_CLUSTER_NAME)
    node_cfg.update(node_provider_cfg_updates)
    # merge node provider tag specs with user overrides
    AWSNodeProvider._merge_tag_specs(tag_specs, user_tag_specs)


def get_cloudwatch_agent_config_file_path():
    return get_aws_example_config_file_path(
        "cloudwatch/example-cloudwatch-agent-config.json"
    )


def get_cloudwatch_dashboard_config_file_path():
    return get_aws_example_config_file_path(
        "cloudwatch/example-cloudwatch-dashboard-config.json"
    )


def get_cloudwatch_alarm_config_file_path():
    return get_aws_example_config_file_path(
        "cloudwatch/example-cloudwatch-alarm-config.json"
    )


def load_cloudwatch_example_config_file():
    config = load_aws_example_config_file("example-cloudwatch.yaml")
    cw_cfg = config["provider"]["cloudwatch"]
    cw_cfg["agent"]["config"] = get_cloudwatch_agent_config_file_path()
    cw_cfg["dashboard"]["config"] = get_cloudwatch_dashboard_config_file_path()
    cw_cfg["alarm"]["config"] = get_cloudwatch_alarm_config_file_path()
    return config


def get_cloudwatch_helper(node_ids):
    config = load_cloudwatch_example_config_file()
    config["cluster_name"] = DEFAULT_CLUSTER_NAME
    return CloudwatchHelper(
        config["provider"],
        node_ids,
        config["cluster_name"],
    )


def get_ssm_param_name(cluster_name, config_type):
    ssm_config_param_name = "AmazonCloudWatch-" + "ray_{}_config_{}".format(
        config_type, cluster_name
    )
    return ssm_config_param_name