

import boto3
import json

pricing_client = boto3.client("pricing", region_name="us-east-1")

PRICE_CACHE = {}

def get_pricing_location(region):

    REGION_MAP = {
        "us-east-1": "US East (N. Virginia)",
        "us-east-2": "US East (Ohio)",
        "us-west-1": "US West (N. California)",
        "us-west-2": "US West (Oregon)",
        "eu-west-1": "EU (Ireland)",
        "eu-west-2": "EU (London)",
        "eu-west-3": "EU (Paris)",
        "ap-south-1": "Asia Pacific (Mumbai)"
    }

    return REGION_MAP.get(region, "US East (N. Virginia)")

EC2_CLIENT_CACHE = {}

def get_ami_root_volume(ami_id, region):
    if region not in EC2_CLIENT_CACHE:
        EC2_CLIENT_CACHE[region] = boto3.client("ec2", region_name=region)
    ec2 = EC2_CLIENT_CACHE[region]
    response = ec2.describe_images(ImageIds=[ami_id])
    images = response.get("Images", [])
    if not images:
        return None, None
    image = images[0]
    mappings = image.get("BlockDeviceMappings", [])
    for mapping in mappings:
        ebs = mapping.get("Ebs")
        if ebs:
            size = ebs.get("VolumeSize")
            vtype = ebs.get("VolumeType")
            return size, vtype

    return None, None

def get_ec2_price(instance_type, region):

    cache_key = f"{instance_type}:{region}"

    if cache_key in PRICE_CACHE:
        return PRICE_CACHE[cache_key]

    # location = REGION_MAP.get(region, "US East (N. Virginia)")
    location = get_pricing_location(region)

    response = pricing_client.get_products(
        ServiceCode="AmazonEC2",
        Filters=[
            {"Type": "TERM_MATCH", "Field": "instanceType", "Value": instance_type},
            {"Type": "TERM_MATCH", "Field": "location", "Value": location},
            {"Type": "TERM_MATCH", "Field": "operatingSystem", "Value": "Linux"},
            {"Type": "TERM_MATCH", "Field": "tenancy", "Value": "Shared"},
            {"Type": "TERM_MATCH", "Field": "preInstalledSw", "Value": "NA"},
            {"Type": "TERM_MATCH", "Field": "capacitystatus", "Value": "Used"},
        ],
        MaxResults=1
    )

    if not response["PriceList"]:
        return 0

    price_item = json.loads(response["PriceList"][0])

    terms = price_item["terms"]["OnDemand"]

    for term in terms.values():
        for dim in term["priceDimensions"].values():
            price = float(dim["pricePerUnit"]["USD"])

            # store in cache
            PRICE_CACHE[cache_key] = price

            return price

    return 0


def get_ebs_price(volume_type, region):

    # location = REGION_MAP.get(region, "US East (N. Virginia)")
    location = get_pricing_location(region)

    response = pricing_client.get_products(
        ServiceCode="AmazonEC2",
        Filters=[
            {"Type": "TERM_MATCH", "Field": "volumeApiName", "Value": volume_type},
            {"Type": "TERM_MATCH", "Field": "location", "Value": location},
            {"Type": "TERM_MATCH", "Field": "productFamily", "Value": "Storage"}
        ],
        MaxResults=1
    )

    if not response["PriceList"]:
        return 0

    price_item = json.loads(response["PriceList"][0])

    terms = price_item["terms"]["OnDemand"]

    for term in terms.values():
        for dim in term["priceDimensions"].values():
            return float(dim["pricePerUnit"]["USD"])

    return 0

def get_instance_price(instance_type, region):
    """
    Generic instance pricing.
    Currently uses EC2 pricing API.
    Works for EC2 and similar instance-based services.
    """
    return get_ec2_price(instance_type, region)


def get_storage_price(volume_type, region):
    """
    Generic storage pricing.
    Currently uses EBS pricing API.
    """
    return get_ebs_price(volume_type, region)