#!/usr/bin/env python
"""
Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
SPDX-License-Identifier: MIT-0

Updates our dynamic patches from the pricing API
This script requires Boto3 and Credentials to call the Pricing API
"""

import json
import logging

import boto3
from botocore.client import Config

LOGGER = logging.getLogger("cfnlint")


region_map = {
    "Any": "all",
    "Africa (Cape Town)": "af-south-1",
    "Asia Pacific (Hong Kong)": "ap-east-1",
    "Asia Pacific (Hyderabad)": "ap-south-2",
    "Asia Pacific (Jakarta)": "ap-southeast-3",
    "Asia Pacific (Malaysia)": "ap-southeast-5",
    "Asia Pacific (Melbourne)": "ap-southeast-4",
    "Asia Pacific (Mumbai)": "ap-south-1",
    "Asia Pacific (New Zealand)": "ap-southeast-6",
    "Asia Pacific (Osaka-Local)": "ap-northeast-3",
    "Asia Pacific (Osaka)": "ap-northeast-3",
    "Asia Pacific (Seoul)": "ap-northeast-2",
    "Asia Pacific (Singapore)": "ap-southeast-1",
    "Asia Pacific (Sydney)": "ap-southeast-2",
    "Asia Pacific (Taipei)": "ap-east-2",
    "Asia Pacific (Thailand)": "ap-southeast-7",
    "Asia Pacific (Tokyo)": "ap-northeast-1",
    "AWS GovCloud (US-East)": "us-gov-east-1",
    "AWS GovCloud (US-West)": "us-gov-west-1",
    "Canada (Central)": "ca-central-1",
    "Canada West (Calgary)": "ca-west-1",
    "China (Beijing)": "cn-north-1",
    "China (Ningxia)": "cn-northwest-1",
    "EU (Frankfurt)": "eu-central-1",
    "EU (Ireland)": "eu-west-1",
    "EU (London)": "eu-west-2",
    "EU (Milan)": "eu-south-1",
    "EU (Paris)": "eu-west-3",
    "EU (Stockholm)": "eu-north-1",
    "Europe (Spain)": "eu-south-2",
    "Europe (Zurich)": "eu-central-2",
    "Israel (Tel Aviv)": "il-central-1",
    "Mexico (Central)": "mx-central-1",
    "Middle East (Bahrain)": "me-south-1",
    "Middle East (UAE)": "me-central-1",
    "South America (Sao Paulo)": "sa-east-1",
    "US East (N. Virginia)": "us-east-1",
    "US East (Ohio)": "us-east-2",
    "US West (Los Angeles)": "us-west-2",
    "US West (N. California)": "us-west-1",
    "US West (Oregon)": "us-west-2",
}

session = boto3.session.Session()
config = Config(retries={"max_attempts": 10})
client = session.client("pricing", region_name="us-east-1", config=config)

_UNSUPPORTED_REGIONS = set()


def configure_logging():
    """Setup Logging"""
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)

    LOGGER.setLevel(logging.INFO)
    log_formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    ch.setFormatter(log_formatter)

    # make sure all other log handlers are removed before adding it back
    for handler in LOGGER.handlers:
        LOGGER.removeHandler(handler)
    LOGGER.addHandler(ch)


def get_paginator(service):
    LOGGER.info("Get " + service + " pricing paginator")
    return client.get_paginator("get_products").paginate(
        ServiceCode=service,
        FormatVersion="aws_v1",
    )


def get_dax_pricing():
    results = {}
    for page in get_paginator("AmazonDAX"):
        for price_item in page.get("PriceList", []):
            products = json.loads(price_item)
            product = products.get("product", {})
            if product:
                if product.get("productFamily") in ["DAX"]:
                    location = region_map.get(product.get("attributes").get("location"))
                    if not location:
                        _UNSUPPORTED_REGIONS.add(
                            product.get("attributes").get("location")
                        )
                        continue
                    if not results.get(location):
                        results[location] = set()
                    results[location].add(
                        product.get("attributes").get("usagetype").split(":")[1]
                    )
    return results


def get_mq_pricing():
    remap = {"mq.m5.2xl": "mq.m5.2xlarge", "mq.m5.4xl": "mq.m5.4xlarge"}
    results = {}
    for page in get_paginator("AmazonMQ"):
        for price_item in page.get("PriceList", []):
            products = json.loads(price_item)
            product = products.get("product", {})
            if product:
                if product.get("productFamily") in ["Broker Instances"]:
                    location = region_map.get(product.get("attributes").get("location"))
                    if not location:
                        _UNSUPPORTED_REGIONS.add(
                            product.get("attributes").get("location")
                        )
                        continue
                    if not results.get(location):
                        results[location] = set()
                    usage_type = (
                        product.get("attributes").get("usagetype").split(":")[1]
                    )
                    results[location].add(remap.get(usage_type, usage_type))
    return results


def get_rds_pricing():
    product_map = {
        "2": ["mysql"],
        "3": ["oracle-se1"],
        "4": ["oracle-se"],
        "5": ["oracle-ee"],
        "6": ["oracle-se1"],
        "8": ["sqlserver-se"],
        "9": ["sqlserver-ee"],
        "10": ["sqlserver-ex"],
        "11": ["sqlserver-web"],
        "12": ["sqlserver-se"],
        "14": ["postgres"],
        "15": ["sqlserver-ee"],
        "16": ["aurora-mysql", "aurora"],
        "18": ["mariadb"],
        "19": ["oracle-se2"],
        "20": ["oracle-se2"],
        "21": ["aurora-postgresql"],
        "28": ["db2-ae", "db2-se"],
    }

    license_map = {
        "License included": "license-included",
        "Bring your own license": "bring-your-own-license",
        "No license required": "general-public-license",
    }

    rds_details = {}

    results = {}
    for page in get_paginator("AmazonRDS"):
        for price_item in page.get("PriceList", []):
            products = json.loads(price_item)
            product = products.get("product", {})
            if product:
                if product.get("productFamily") in ["Database Instance"]:
                    if product.get("attributes").get("locationType") == "AWS Outposts":
                        continue
                    # Get overall instance types
                    location = region_map.get(product.get("attributes").get("location"))
                    if not location:
                        _UNSUPPORTED_REGIONS.add(
                            product.get("attributes").get("location")
                        )
                        continue
                    if not results.get(location):
                        results[location] = set(["db.serverless"])
                    instance_type = product.get("attributes").get("instanceType")
                    if instance_type:
                        results[location].add(instance_type)
                    # Rds Instance Size spec
                    product_names = product_map.get(
                        product.get("attributes").get("engineCode"), []
                    )
                    license_name = license_map.get(
                        product.get("attributes").get("licenseModel")
                    )
                    deployment_option = product.get("attributes").get(
                        "deploymentOption"
                    )

                    if deployment_option not in [
                        "Single-AZ",
                        "Multi-AZ",
                        "Multi-AZ (readable standbys)",
                    ]:
                        continue

                    instance_type = product.get("attributes").get("instanceType")
                    for product_name in product_names:
                        if not rds_details.get(location):
                            rds_details[location] = {}
                        if not rds_details.get(location).get(deployment_option):
                            rds_details[location][deployment_option] = {}
                        if (
                            not rds_details.get(location)
                            .get(deployment_option)
                            .get(license_name)
                        ):
                            rds_details[location][deployment_option][license_name] = {}
                        if (
                            not rds_details.get(location)
                            .get(deployment_option)
                            .get(license_name)
                            .get(product_name)
                        ):
                            if (
                                license_name == "general-public-license"
                                and product_name
                                in ["aurora-mysql", "aurora-postgresql"]
                            ):
                                rds_details[location][deployment_option][license_name][
                                    product_name
                                ] = set(["db.serverless"])
                            else:
                                rds_details[location][deployment_option][license_name][
                                    product_name
                                ] = set()
                        rds_details[location][deployment_option][license_name][
                            product_name
                        ].add(instance_type)
    specs = {}
    cluster_specs = {}
    for product_region, product_values in rds_details.items():
        for deployment_option, deployment_values in product_values.items():
            if deployment_option in ["Single-AZ", "Multi-AZ"]:
                for license_name, license_values in deployment_values.items():
                    for product_name, instance_types in license_values.items():
                        if product_region not in specs:
                            specs[product_region] = {"allOf": []}
                        if license_name == "general-public-license":
                            specs[product_region]["allOf"].append(
                                {
                                    "if": {
                                        "properties": {
                                            "Engine": {
                                                "const": product_name,
                                            },
                                            "LicenseModel": False,
                                            "DBInstanceClass": {"type": "string"},
                                        },
                                        "required": ["Engine", "DBInstanceClass"],
                                    },
                                    "then": {
                                        "properties": {
                                            "DBInstanceClass": {
                                                "enum": list(sorted(instance_types)),
                                            },
                                        }
                                    },
                                }
                            )
                        specs[product_region]["allOf"].append(
                            {
                                "if": {
                                    "properties": {
                                        "Engine": {
                                            "const": product_name,
                                        },
                                        "LicenseModel": {
                                            "const": license_name,
                                        },
                                        "DBInstanceClass": {"type": "string"},
                                    },
                                    "required": [
                                        "Engine",
                                        "LicenseModel",
                                        "DBInstanceClass",
                                    ],
                                },
                                "then": {
                                    "properties": {
                                        "DBInstanceClass": {
                                            "enum": list(sorted(instance_types)),
                                        },
                                    }
                                },
                            }
                        )
            else:
                for license_name, license_values in deployment_values.items():
                    for product_name, instance_types in license_values.items():
                        if product_region not in cluster_specs:
                            cluster_specs[product_region] = {"allOf": []}
                        cluster_specs[product_region]["allOf"].append(
                            {
                                "if": {
                                    "properties": {
                                        "Engine": {
                                            "const": product_name,
                                        },
                                        "DBClusterInstanceClass": {"type": "string"},
                                    },
                                    "required": ["Engine", "DBClusterInstanceClass"],
                                },
                                "then": {
                                    "properties": {
                                        "DBClusterInstanceClass": {
                                            "enum": list(sorted(instance_types)),
                                        },
                                    }
                                },
                            }
                        )

    LOGGER.info("Updating RDS Spec files")
    filename = (
        "src/cfnlint/data/schemas/extensions/"
        "aws_rds_dbinstance/dbinstanceclass_enum.json"
    )
    with open(filename, "w+", encoding="utf-8") as f:
        json.dump(specs, f, indent=1, sort_keys=True, separators=(",", ": "))
        f.write("\n")

    filename = (
        "src/cfnlint/data/schemas/extensions/"
        "aws_rds_dbcluster/dbclusterinstanceclass_enum.json"
    )
    with open(filename, "w+", encoding="utf-8") as f:
        json.dump(cluster_specs, f, indent=1, sort_keys=True, separators=(",", ": "))
        f.write("\n")


def get_sagemaker_pricing():
    results = {}
    for page in get_paginator("AmazonSageMaker"):
        for price_item in page.get("PriceList", []):
            products = json.loads(price_item)
            product = products.get("product", {})
            if not product:
                continue
            if product.get("productFamily") != "ML Instance":
                continue
            attrs = product.get("attributes", {})
            if attrs.get("locationType") != "AWS Region":
                continue
            location = region_map.get(attrs.get("location"))
            if not location:
                _UNSUPPORTED_REGIONS.add(attrs.get("location"))
                continue
            instance_name = attrs.get("instanceName", "")
            if not instance_name:
                continue
            usagetype = attrs.get("usagetype", "")
            if ":" not in usagetype:
                continue
            prefix = usagetype.split(":")[0]
            if "-" in prefix:
                prefix = prefix.split("-", 1)[-1]

            if location not in results:
                results[location] = {}
            if prefix not in results[location]:
                results[location][prefix] = set()
            results[location][prefix].add(instance_name)

    return results


def get_results(service, product_families, default=None):
    if default is None:
        default = set()
    results = {}
    for page in get_paginator(service):
        for price_item in page.get("PriceList", []):
            products = json.loads(price_item)
            product = products.get("product", {})
            if product:
                if (
                    product.get("productFamily") in product_families
                    and product.get("attributes").get("locationType") == "AWS Region"
                ):
                    location = region_map.get(product.get("attributes").get("location"))
                    if not location:
                        _UNSUPPORTED_REGIONS.add(
                            product.get("attributes").get("location")
                        )
                        continue
                    if not results.get(location):
                        results[location] = default
                    instance_type = product.get("attributes").get("instanceType")
                    if instance_type:
                        results[location].add(instance_type)
    return results


def write_output(resource, filename, obj):
    filename = f"src/cfnlint/data/schemas/extensions/{resource}/{filename}.json"
    output = {
        "description": "Automatically updated using update_specs_from_pricing",
    }
    for region, values in obj.items():
        output[region] = {"enum": sorted(list(values))}

    with open(filename, "w+", encoding="utf-8") as f:
        json.dump(output, f, indent=1, sort_keys=True, separators=(",", ": "))
        f.write("\n")


def main():
    """main function"""
    configure_logging()

    write_output(
        "aws_ec2_instance",
        "instancetype_enum",
        get_results("AmazonEC2", ["Compute Instance", "Compute Instance (bare metal)"]),
    )
    write_output("aws_amazonmq_broker", "instancetype_enum", get_mq_pricing())
    get_rds_pricing()
    write_output(
        "aws_redshift_cluster",
        "nodetype_enum",
        get_results("AmazonRedshift", ["Compute Instance"]),
    )
    write_output("aws_dax_cluster", "nodetype_enum", get_dax_pricing())
    write_output(
        "aws_docdb_dbinstance",
        "dbinstanceclass_enum",
        get_results("AmazonDocDB", ["Database Instance"]),
    )
    write_output(
        "aws_neptune_dbinstance",
        "dbinstanceclass_enum",
        get_results(
            "AmazonNeptune", ["Database Instance"], default=set(["db.serverless"])
        ),
    )
    write_output(
        "aws_elasticache_cachecluster",
        "cachenodetype_enum",
        get_results("AmazonElastiCache", ["Cache Instance"]),
    )
    _opensearch_results = get_results(
        "AmazonES", ["Amazon OpenSearch Service Instance"]
    )
    write_output(
        "aws_elasticsearch_domain",
        "elasticsearchclusterconfig_instancetype_enum",
        _opensearch_results,
    )
    write_output(
        "aws_opensearchservice_domain",
        "clusterconfig_instancetype_enum",
        _opensearch_results,
    )
    write_output(
        "aws_emr_cluster",
        "instancetypeconfig_instancetype_enum",
        get_results("ElasticMapReduce", ["Elastic Map Reduce Instance"]),
    )
    write_output(
        "aws_managedblockchain_node",
        "nodeconfiguration_instancetype_enum",
        get_results("AmazonManagedBlockchain", ["Blockchain Instance"]),
    )
    write_output(
        "aws_gamelift_fleet",
        "ec2instancetype_enum",
        get_results("AmazonGameLift", ["GameLift EC2 Instance"]),
    )
    write_output(
        "aws_appstream_fleet",
        "instancetype_enum",
        get_results("AmazonAppStream", ["Streaming Instance"]),
    )

    _sagemaker_results = get_sagemaker_pricing()
    _sagemaker_processing = {
        region: prefixes.get("Processing", set())
        for region, prefixes in _sagemaker_results.items()
        if prefixes.get("Processing")
    }
    _sagemaker_transform = {
        region: prefixes.get("Tsform", set())
        for region, prefixes in _sagemaker_results.items()
        if prefixes.get("Tsform")
    }
    _sagemaker_hosting = {
        region: prefixes.get("Host", set()) | prefixes.get("AsyncInf", set())
        for region, prefixes in _sagemaker_results.items()
        if prefixes.get("Host") or prefixes.get("AsyncInf")
    }
    _sagemaker_cluster = {
        region: prefixes.get("Cluster", set())
        for region, prefixes in _sagemaker_results.items()
        if prefixes.get("Cluster")
    }
    write_output(
        "aws_sagemaker_cluster",
        "instancetype_enum",
        _sagemaker_cluster,
    )
    write_output(
        "aws_sagemaker_processing",
        "instancetype_enum",
        _sagemaker_processing,
    )
    write_output(
        "aws_sagemaker_transform",
        "instancetype_enum",
        _sagemaker_transform,
    )
    write_output(
        "aws_sagemaker_hosting",
        "instancetype_enum",
        _sagemaker_hosting,
    )
    for region in _UNSUPPORTED_REGIONS:
        LOGGER.warning(f"Region {region!r} is not supported")


if __name__ == "__main__":
    try:
        main()
    except (ValueError, TypeError) as e:
        LOGGER.error(e)
