package utils

import (
	"context"
	"slices"
	"strings"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
	"github.com/juju/errors"

	interal_types "github.com/tidbcloud/auto-deploy/internal/types"
)

func GetAvailabilityZonesString(provider interal_types.CloudProviderType, region string) (string, error) {
	return GetAvailabilityZonesNString(provider, region, -1)
}

func GetAvailabilityZoneIndex(provider interal_types.CloudProviderType, region string, index int) (string, error) {
	if index < 0 {
		return "", errors.NotValidf("index")
	}

	n := index + 1
	azList, err := GetAvailabilityZonesN(provider, region, n)
	if err != nil {
		return "", err
	}
	return azList[index], nil
}

func GetAvailabilityZonesNString(provider interal_types.CloudProviderType, region string, n int) (string, error) {
	zones, err := GetAvailabilityZonesN(provider, region, n)
	if err != nil {
		return "", err
	}
	return strings.Join(zones, ","), nil
}

func GetAvailabilityZonesN(provider interal_types.CloudProviderType, region string, n int) ([]string, error) {
	switch provider {
	case interal_types.AWS:
		return GetAwsAvailabilityZonesN(region, n)
	case interal_types.Azure:
		return GetAzureAvailabilityZonesN(region, n)
	default:
		return []string{}, errors.NotImplementedf("GetAvailabilityZonesN for %s", provider)
	}
}
func GetAwsAvailabilityZonesN(region string, n int) ([]string, error) {
	ctx := context.Background()
	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		return nil, errors.Annotatef(err, "Failed to load default aws config")
	}

	cfg.Region = region
	client := ec2.NewFromConfig(cfg)
	result, err := client.DescribeAvailabilityZones(ctx, &ec2.DescribeAvailabilityZonesInput{
		AllAvailabilityZones: Ptr(true),
		Filters: []types.Filter{
			{
				Name:   Ptr("zone-type"),
				Values: []string{"availability-zone"},
			},
			{
				Name:   Ptr("region-name"),
				Values: []string{region},
			},
		},
	})
	if err != nil {
		return nil, errors.Annotatef(err, "Failed to list availability zones: %s", region)
	}

	azIds := Map(result.AvailabilityZones, func(zone types.AvailabilityZone) string { return *zone.ZoneId })
	slices.Sort(azIds)
	if n > 0 {
		if len(azIds) < n {
			return nil, errors.Errorf("AvailabilityZones count %d is less than required %d", len(azIds), n)
		}
		azIds = azIds[0:n]
	}
	return azIds, nil
}

func AwsAzIdToName(region, ids, profile string) (string, error) {
	if profile == "" {
		profile = "default"
	}
	ctx := context.Background()
	cfg, err := config.LoadDefaultConfig(ctx, config.WithSharedConfigProfile(profile))
	if err != nil {
		return "", errors.Annotatef(err, "Failed to load default aws config")
	}

	cfg.Region = region
	client := ec2.NewFromConfig(cfg)
	result, err := client.DescribeAvailabilityZones(ctx, &ec2.DescribeAvailabilityZonesInput{
		AllAvailabilityZones: Ptr(true),
		Filters: []types.Filter{
			{
				Name:   Ptr("zone-type"),
				Values: []string{"availability-zone"},
			},
			{
				Name:   Ptr("region-name"),
				Values: []string{region},
			},
		},
	})
	if err != nil {
		return "", errors.Annotatef(err, "Failed to list availability zones: %s", region)
	}
	azNamesMap := ToMap(result.AvailabilityZones, func(i int, t types.AvailabilityZone) string { return *t.ZoneId },
		func(i int, t types.AvailabilityZone) string { return *t.ZoneName })

	idsList := strings.Split(ids, ",")
	azNames := make([]string, len(idsList))
	for i, id := range idsList {
		name, ok := azNamesMap[id]
		if !ok {
			return "", errors.BadRequestf("Availability zone id %s is not valid, all availability zones %s", id, Map(result.AvailabilityZones, func(zone types.AvailabilityZone) string { return *zone.ZoneId }))
		}
		azNames[i] = name
	}
	return strings.Join(azNames, ","), nil
}