import copy
from datetime import datetime

import ray
from ray.autoscaler.tags import (
    NODE_KIND_HEAD,
    TAG_RAY_LAUNCH_CONFIG,
    TAG_RAY_NODE_KIND,
    TAG_RAY_USER_NODE_TYPE,
)

# Override global constants used in AWS autoscaler config artifact names.
# This helps ensure that any unmocked test doesn't alter non-test artifacts.
ray.autoscaler._private.aws.config.RAY = "ray-autoscaler-aws-test"
ray.autoscaler._private.aws.config.DEFAULT_RAY_INSTANCE_PROFILE = (
    ray.autoscaler._private.aws.config.RAY + "-v1"
)
ray.autoscaler._private.aws.config.DEFAULT_RAY_IAM_ROLE = (
    ray.autoscaler._private.aws.config.RAY + "-v1"
)
ray.autoscaler._private.aws.config.SECURITY_GROUP_TEMPLATE = (
    ray.autoscaler._private.aws.config.RAY + "-{}"
)

# Default IAM instance profile to expose to tests.
DEFAULT_INSTANCE_PROFILE = {
    "Arn": "arn:aws:iam::336924118301:instance-profile/ExampleInstanceProfile",
    "CreateDate": datetime(2013, 6, 12, 23, 52, 2, 2),
    "InstanceProfileId": "AIPA0000000000EXAMPLE",
    "InstanceProfileName": "ExampleInstanceProfile",
    "Path": "/",
    "Roles": [
        {
            "Arn": "arn:aws:iam::123456789012:role/Test-Role",
            "AssumeRolePolicyDocument": "ExampleAssumeRolePolicyDocument",
            "CreateDate": datetime(2013, 1, 9, 6, 33, 26, 2),
            "Path": "/",
            "RoleId": "AROA0000000000EXAMPLE",
            "RoleName": "Test-Role",
        },
    ],
}

# Default EC2 key pair to expose to tests.
DEFAULT_KEY_PAIR = {
    "KeyFingerprint": "00:11:22:33:44:55:66:77:88:99:AA:BB:CC:DD:EE:FF:00",
    "KeyName": ray.autoscaler._private.aws.config.RAY + "_us-west-2",
}

# Primary EC2 subnet to expose to tests.
DEFAULT_SUBNET = {
    "AvailabilityZone": "us-west-2a",
    "AvailableIpAddressCount": 251,
    "CidrBlock": "10.0.1.0/24",
    "DefaultForAz": False,
    "MapPublicIpOnLaunch": True,
    "State": "available",
    "SubnetId": "subnet-0000000",
    "VpcId": "vpc-0000000",
}


def subnet_in_vpc(vpc_num):
    """Returns a copy of DEFAULT_SUBNET whose VpcId ends with the digits
    of vpc_num."""
    subnet = copy.copy(DEFAULT_SUBNET)
    subnet["VpcId"] = f"vpc-{vpc_num:07d}"
    return subnet


A_THOUSAND_SUBNETS_IN_DIFFERENT_VPCS = [
    subnet_in_vpc(vpc_num) for vpc_num in range(1, 1000)
] + [DEFAULT_SUBNET]


def subnet_in_az(idx):
    azs = ["a", "b", "c", "d"]
    subnet = copy.copy(DEFAULT_SUBNET)
    subnet["AvailabilityZone"] = "us-west-2" + azs[idx % 4]
    subnet["SubnetId"] = f"subnet-{idx:07d}"
    return subnet


TWENTY_SUBNETS_IN_DIFFERENT_AZS = [subnet_in_az(i) for i in range(20)]

# Secondary EC2 subnet to expose to tests as required.
AUX_SUBNET = {
    "AvailabilityZone": "us-west-2a",
    "AvailableIpAddressCount": 251,
    "CidrBlock": "192.168.1.0/24",
    "DefaultForAz": False,
    "MapPublicIpOnLaunch": True,
    "State": "available",
    "SubnetId": "subnet-fffffff",
    "VpcId": "vpc-fffffff",
}

# Default cluster name to expose to tests.
DEFAULT_CLUSTER_NAME = "test-cluster-name"

# Default security group settings immediately after creation
# (prior to inbound rule configuration).
DEFAULT_SG = {
    "Description": "Auto-created security group for Ray workers",
    "GroupName": ray.autoscaler._private.aws.config.RAY + "-" + DEFAULT_CLUSTER_NAME,
    "OwnerId": "test-owner",
    "GroupId": "sg-1234abcd",
    "VpcId": DEFAULT_SUBNET["VpcId"],
    "IpPermissions": [],
    "IpPermissionsEgress": [
        {
            "FromPort": -1,
            "ToPort": -1,
            "IpProtocol": "-1",
            "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
        }
    ],
    "Tags": [],
}

# Secondary security group settings after creation
# (prior to inbound rule configuration).
AUX_SG = copy.deepcopy(DEFAULT_SG)
AUX_SG["GroupName"] += "-aux"
AUX_SG["GroupId"] = "sg-dcba4321"

# Default security group settings immediately after creation on aux subnet
# (prior to inbound rule configuration).
DEFAULT_SG_AUX_SUBNET = copy.deepcopy(DEFAULT_SG)
DEFAULT_SG_AUX_SUBNET["VpcId"] = AUX_SUBNET["VpcId"]
DEFAULT_SG_AUX_SUBNET["GroupId"] = AUX_SG["GroupId"]

DEFAULT_IN_BOUND_RULES = [
    {
        "FromPort": -1,
        "ToPort": -1,
        "IpProtocol": "-1",
        "UserIdGroupPairs": [{"GroupId": DEFAULT_SG["GroupId"]}],
    },
    {
        "FromPort": 22,
        "ToPort": 22,
        "IpProtocol": "tcp",
        "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
    },
]
# Default security group settings once default inbound rules are applied
# (if used by both head and worker nodes)
DEFAULT_SG_WITH_RULES = copy.deepcopy(DEFAULT_SG)
DEFAULT_SG_WITH_RULES["IpPermissions"] = DEFAULT_IN_BOUND_RULES

# Default security group once default inbound rules are applied
# (if using separate security groups for head and worker nodes).
DEFAULT_SG_DUAL_GROUP_RULES = copy.deepcopy(DEFAULT_SG_WITH_RULES)
DEFAULT_SG_DUAL_GROUP_RULES["IpPermissions"][0]["UserIdGroupPairs"].append(
    {"GroupId": AUX_SG["GroupId"]}
)

# Default security group on aux subnet once default inbound rules are applied.
DEFAULT_SG_WITH_RULES_AUX_SUBNET = copy.deepcopy(DEFAULT_SG_DUAL_GROUP_RULES)
DEFAULT_SG_WITH_RULES_AUX_SUBNET["VpcId"] = AUX_SUBNET["VpcId"]
DEFAULT_SG_WITH_RULES_AUX_SUBNET["GroupId"] = AUX_SG["GroupId"]

# Default security group with custom name
DEFAULT_SG_WITH_NAME = copy.deepcopy(DEFAULT_SG)
DEFAULT_SG_WITH_NAME["GroupName"] = "test_security_group_name"

CUSTOM_IN_BOUND_RULES = [
    {
        "FromPort": 443,
        "ToPort": 443,
        "IpProtocol": "TCP",
        "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
    },
    {
        "FromPort": 8265,
        "ToPort": 8265,
        "IpProtocol": "TCP",
        "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
    },
]

# Default security group with custom name once...
# default and custom in bound rules are applied
DEFAULT_SG_WITH_NAME_AND_RULES = copy.deepcopy(DEFAULT_SG_WITH_NAME)
DEFAULT_SG_WITH_NAME_AND_RULES["IpPermissions"] = (
    DEFAULT_IN_BOUND_RULES + CUSTOM_IN_BOUND_RULES
)

# Default launch template to expose to tests.
DEFAULT_LT = {
    "LaunchTemplateId": "lt-00000000000000000",
    "LaunchTemplateName": "ExampleLaunchTemplate",
    "VersionNumber": 2,
    "CreateTime": datetime(2020, 8, 17, 23, 30, 3),
    "CreatedBy": DEFAULT_INSTANCE_PROFILE["Roles"][0]["Arn"],
    "DefaultVersion": True,
    "LaunchTemplateData": {
        "EbsOptimized": False,
        "IamInstanceProfile": {"Arn": DEFAULT_INSTANCE_PROFILE["Arn"]},
        "NetworkInterfaces": [
            {
                "DeviceIndex": 0,
                "Groups": [DEFAULT_SG["GroupId"]],
                "SubnetId": DEFAULT_SUBNET["SubnetId"],
            }
        ],
        "ImageId": "ami-00000000000000000",
        "InstanceType": "m5.large",
        "TagSpecifications": [
            {
                "ResourceType": "instance",
                "Tags": [{"Key": "test-key", "Value": "test-value"}],
            },
            {
                "ResourceType": "volume",
                "Tags": [{"Key": "test-key", "Value": "test-value"}],
            },
        ],
    },
}

# Default node provider tags to expose to tests.
DEFAULT_NODE_PROVIDER_INSTANCE_TAGS = {
    TAG_RAY_NODE_KIND: NODE_KIND_HEAD,
    TAG_RAY_LAUNCH_CONFIG: "test-ray-launch-config",
    TAG_RAY_USER_NODE_TYPE: "ray.head.default",
}