import click
import subprocess
from tensorkube.constants import get_cluster_name, get_image_registry_id
from tensorkube.services.aws_service import get_availability_zones, get_session_region

def apply(test: bool = False):
    try:
        cluster_name = get_cluster_name()
        region = get_session_region()
        release_name = "karpenter-configuration"
        chart_version = "0.1.4"

        click.echo(click.style("Updating Karpenter configuration to add p4de support for A100 GPUs", bold=True))

        # Build helm upgrade command
        update_karpenter_config_command = [
            "helm", "upgrade", "--install", release_name,
            f"oci://public.ecr.aws/{get_image_registry_id(test)}/tensorfuse/helm-charts/tk-karpenter-config",
            "--version", chart_version,
            "--set", f"clusterName={cluster_name}"
        ]

        # Add availability zones
        availability_zones = get_availability_zones(region=region)
        for (i, az) in enumerate(availability_zones):
            update_karpenter_config_command.append("--set")
            update_karpenter_config_command.append(f"karpenterTopologyZoneValues[{i}]={az}")

        # Run helm upgrade
        subprocess.run(update_karpenter_config_command, check=True)

        click.echo(click.style("✓ Successfully added p4de support to GPU nodepools", bold=True, fg="green"))
        click.echo(click.style("  • New A100 deployments will use p4de instances (80GB A100)", fg="yellow"))
        click.echo(click.style("  • Existing deployments continue on p4d until redeployed", fg="yellow"))
        click.echo(click.style("  • Both p4d and p4de instances are now supported", fg="yellow"))

    except subprocess.CalledProcessError as e:
        click.echo(click.style(f"Failed to update Karpenter configuration: {e}", bold=True, fg="red"))
        raise e
    except Exception as e:
        click.echo(click.style(f"Unexpected error during migration: {e}", bold=True, fg="red"))
        raise e
