gen_kconfig_instance

#!/usr/bin/env python3
# ex: set filetype=python:

"""
Retrieve instance family and type information from AWS. Use it to
construct the "instance" Kconfig menu.
"""

import sys
import argparse

from aws_common import (
    AwsNotConfiguredError,
    get_default_region,
    get_all_instance_types,
    get_jinja2_environment,
    require_aws_credentials,
)


def parse_all_instance_families(instance_types, quiet=False):
    """
    Get all available instance families in the specified region.

    Args:
        region (str): AWS region to query

    Returns:
        dict: Dictionary with family info including count of instances per family
    """
    families = {}
    for instance_type in instance_types:
        instance_name = instance_type["InstanceType"]
        family = instance_name.split(".")[0]

        if family not in families:
            families[family] = {
                "family_name": family,
                "instance_count": 0,
                "has_gpu": False,
                "architectures": set(),
            }

        families[family]["instance_count"] += 1

        if "GpuInfo" in instance_type:
            families[family]["has_gpu"] = True
        cpu_architectures = instance_type.get("ProcessorInfo", {}).get(
            "SupportedArchitectures", []
        )
        families[family]["architectures"].update(cpu_architectures)

    # Convert architecture sets to sorted lists for JSON serialization
    for family in families.values():
        family["architectures"] = sorted(list(family["architectures"]))

    return families


def get_gpu_info(instance_type):
    """
    Extract GPU information from instance type data.

    Args:
        instance_type (dict): Instance type data from AWS API

    Returns:
        str: Formatted GPU information string
    """
    if "GpuInfo" not in instance_type:
        return "None"

    gpu_info = instance_type["GpuInfo"]
    gpus = gpu_info.get("Gpus", [])

    if not gpus:
        return "GPU present (details unavailable)"

    gpu_details = []
    for gpu in gpus:
        gpu_name = gpu.get("Name", "Unknown GPU")
        gpu_count = gpu.get("Count", 1)
        gpu_memory = gpu.get("MemoryInfo", {}).get("SizeInMiB", 0)

        if gpu_count > 1:
            detail = f"{gpu_count}x {gpu_name}"
        else:
            detail = gpu_name

        if gpu_memory > 0:
            detail += f" ({gpu_memory // 1024}GB)"

        gpu_details.append(detail)

    return ", ".join(gpu_details)


def get_instance_family_info(
    family_name, instance_types, region="us-east-1", quiet=False
):
    """
    Get instance types, pricing, and hardware info for an AWS instance family.

    Args:
        family_name (str): Instance family name (e.g., 'm5', 't3', 'c5')
        region (str): AWS region to query (default: us-east-1)
        quiet (bool): Suppress debug messages

    Returns:
        list: List of dictionaries containing instance information
    """

    family_instances = []
    for instance_type in instance_types:
        instance_name = instance_type["InstanceType"]
        if instance_name.startswith(family_name + "."):
            family_instances.append(instance_type)

    if not family_instances:
        if not quiet:
            print(f"No instances found starting with '{family_name}'.", file=sys.stderr)
        return []

    if not quiet:
        print(
            f"Found {len(family_instances)} instances in family '{family_name}'",
            file=sys.stderr,
        )

    instance_info = []

    for instance_type in family_instances:
        instance_name = instance_type["InstanceType"]

        cpu_architectures = instance_type.get("ProcessorInfo", {}).get(
            "SupportedArchitectures", ["Unknown"]
        )
        cpu_isa = ", ".join(cpu_architectures) if cpu_architectures else "Unknown"

        gpu_info = get_gpu_info(instance_type)

        hardware_info = {
            "instance_type": instance_name,
            "vcpus": instance_type["VCpuInfo"]["DefaultVCpus"],
            "memory_gb": instance_type["MemoryInfo"]["SizeInMiB"] / 1024,
            "cpu_isa": cpu_isa,
            "gpu": gpu_info,
            "network_performance": instance_type.get("NetworkInfo", {}).get(
                "NetworkPerformance", "Not specified"
            ),
            "storage": "EBS-only",
            "bare_metal": instance_type["BareMetal"],
            "free_tier": instance_type["FreeTierEligible"],
        }

        if "InstanceStorageInfo" in instance_type:
            storage_info = instance_type["InstanceStorageInfo"]
            total_storage = storage_info.get("TotalSizeInGB", 0)
            storage_type = storage_info.get("Disks", [{}])[0].get("Type", "Unknown")
            hardware_info["storage"] = f"{total_storage} GB {storage_type}"

        hardware_info["price_per_hour_usd"] = "Not available"
        instance_info.append(hardware_info)

    return sorted(instance_info, key=lambda x: x["memory_gb"])


def output_instance_kconfig(instance_types, region):
    """Output instance menu in Kconfg format."""
    families = parse_all_instance_families(instance_types)
    sorted_families = sorted(families.values(), key=lambda x: x["family_name"])
    environment = get_jinja2_environment()
    template = environment.get_template("families.j2")
    print(
        template.render(
            sorted_families=sorted_families,
        )
    )
    print()

    template = environment.get_template("family.j2")
    for family in sorted_families:
        instances = get_instance_family_info(
            family["family_name"], instance_types, region, True
        )
        if not instances:
            continue
        print(
            template.render(
                family_name=family["family_name"].upper().replace("-", "_"),
                instances=instances,
            )
        )
        print()


def output_families_kconfig(instance_types):
    """Output instance families in Kconfg format."""
    families = parse_all_instance_families(instance_types)
    sorted_families = sorted(families.values(), key=lambda x: x["family_name"])
    environment = get_jinja2_environment()
    template = environment.get_template("families.j2")
    print(
        template.render(
            sorted_families=sorted_families,
        )
    )


def output_families_raw(instance_types, region, quiet=False):
    """Output available instance families in table format."""
    families = parse_all_instance_families(instance_types)

    if not quiet:
        print(f"Available instance families in {region}:\n")

    print(f"{'Family':<10} {'Count':<6} {'GPU':<5} {'Architectures':<20}")
    print("-" * 45)

    sorted_families = sorted(families.values(), key=lambda x: x["family_name"])
    for family in sorted_families:
        gpu_indicator = "Yes" if family["has_gpu"] else "No"
        architectures = ", ".join(family["architectures"])

        print(
            f"{family['family_name']:<10} "
            f"{family['instance_count']:<6} "
            f"{gpu_indicator:<5} "
            f"{architectures:<20}"
        )

    if not quiet:
        print(f"\nFound {len(families)} instance families", file=sys.stderr)


def output_family_kconfig(family_name, instances):
    """Output results in kconfig format."""
    environment = get_jinja2_environment()
    template = environment.get_template("family.j2")
    print(
        template.render(
            family_name=family_name.upper().replace("-", "_"),
            instances=instances,
        )
    )


def output_family_raw(instances, quiet=False):
    """Output results in table format."""
    if not quiet:
        print(f"Found {len(instances)} instance types:\n")

    print(
        f"{'Instance Type':<15} {'vCPUs':<6} {'Memory (GB)':<12} {'CPU ISA':<10} {'GPU':<25} {'Storage':<20} {'Network':<15} {'Price/Hour':<12}"
    )
    print("-" * 130)

    for instance in instances:
        print(
            f"{instance['instance_type']:<15} "
            f"{instance['vcpus']:<6} "
            f"{instance['memory_gb']:<12.1f} "
            f"{instance['cpu_isa']:<10} "
            f"{instance['gpu']:<25} "
            f"{instance['storage']:<20} "
            f"{instance['network_performance']:<15} "
            f"{instance['price_per_hour_usd']:<12}"
        )


def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Get AWS EC2 instance family information including pricing and hardware specs",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python %(prog)s m5
  python %(prog)s t3 --region us-west-2
  python %(prog)s c5 --format kconfig
  python %(prog)s r5 --quiet
  python %(prog)s --families
  python %(prog)s --families --format raw
        """,
    )
    parser.add_argument(
        "family_name",
        nargs="?",  # Make family_name optional when using --families
        help="Instance family name (e.g., m5, t3, c5, r5)",
    )

    parser.add_argument(
        "--families",
        action="store_true",
        help="List all available instance families in the region",
    )
    parser.add_argument(
        "--format",
        "-f",
        choices=["raw", "kconfig"],
        default="kconfig",
        help="Output format (default: kconfig)",
    )
    parser.add_argument(
        "--quiet", "-q", action="store_true", help="Suppress informational messages"
    )
    parser.add_argument(
        "--region", "-r", help="AWS region (default: from ~/.aws/config or us-east-1)"
    )
    return parser.parse_args()


def main():
    """Main function to run the program."""
    args = parse_arguments()

    # Allow make dynconfig to succeed without AWS credentials
    try:
        require_aws_credentials()
    except AwsNotConfiguredError:
        if not args.quiet:
            print("AWS not configured - skipping (optional)", file=sys.stderr)
        sys.exit(0)

    if args.region:
        region = args.region
    else:
        region = get_default_region()

    instance_types = get_all_instance_types(region, args.quiet)
    if not instance_types:
        sys.exit(1)

    if args.families:
        if args.format == "kconfig":
            output_families_kconfig(instance_types)
        else:
            output_families_raw(instance_types, region, args.quiet)
        return

    if args.family_name:
        if not args.quiet:
            print(
                f"Fetching information for the {args.family_name} family in {region}...",
                file=sys.stderr,
            )

        instances = get_instance_family_info(
            args.family_name, instance_types, region, args.quiet
        )

        if not instances:
            print(
                f"No instances found for family '{args.family_name}'.", file=sys.stderr
            )
            print(
                "Try running with --families to see available instance families.",
                file=sys.stderr,
            )
            sys.exit(1)

        if args.format == "kconfig":
            output_family_kconfig(args.family_name, instances)
        else:
            output_family_raw(instances, args.quiet)
        return

    output_instance_kconfig(instance_types, region)


if __name__ == "__main__":
    main()