package utils

import (
	"context"
	"fmt"
	"slices"

	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/resources/armsubscriptions"
	"github.com/juju/errors"
)

func GetAzureAvailabilityZonesN(region string, n int) ([]string, error) {
	// 创建 Azure 认证客户端
	cred, err := azidentity.NewDefaultAzureCredential(nil)
	if err != nil {
		return nil, errors.Annotate(err, "Failed to get azure credentials")
	}

	// 创建订阅客户端
	client, err := armsubscriptions.NewClient(cred, nil)
	if err != nil {
		return nil, errors.Annotate(err, "Failed to create subscription client")
	}

	// 创建订阅客户端
	subscriptionID, err := RunCommandAndGetOutput("az", "account", "show", "--query=id", "--output=tsv")
	if err != nil {
		return nil, errors.Annotate(err, "Failed to get default subscription")
	}

	pager := client.NewListLocationsPager(subscriptionID, &armsubscriptions.ClientListLocationsOptions{
		IncludeExtendedLocations: Ptr(false),
	})

	zones := []string{}
	for pager.More() {
		resp, err := pager.NextPage(context.TODO())
		if err != nil {
			return nil, errors.Annotate(err, "Failed to get next page")
		}
		for _, loc := range resp.LocationListResult.Value {
			if *loc.Name == region {
				for _, z := range loc.AvailabilityZoneMappings {
					zones = append(zones, fmt.Sprintf("%s-%s", region, *z.LogicalZone))
				}
			}
		}
	}
	slices.Sort(zones)
	if n > 0 {
		if len(zones) < n {
			return nil, errors.Errorf("AvailabilityZones count %d is less than required %d", len(zones), n)
		}
		zones = zones[0:n]
	}
	return zones, nil
}