package aws

import (
	"context"
	"fmt"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"

	eksv1alpha1 "github.com/tidbcloud/eks-provider/api/v1alpha1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

// GetRoleFromAccount returns the role name for the given account.
// If the account is invalid, it returns nil.
func GetRoleFromAccount(account *cloudv1.Account) (ar *AssumedRole) {
	if account != nil && account.Spec.Provider == cloudv1.ProviderAWS && account.Spec.AWS != nil {
		ar = &AssumedRole{
			RoleARN:    account.Spec.AWS.AssumeRoleARN,
			ExternalID: account.Spec.AWS.ExternalID,
		}
	}

	return
}

// GetRoleFromEKSNetworkSpec returns the role name for the given EKS network spec.
// If the account is invalid, it returns nil.
func GetRoleFromEKSNetworkSpec(spec *eksv1alpha1.EKSNetworkSpec) (ar *AssumedRole) {
	if spec != nil {
		ar = &AssumedRole{
			RoleARN:    spec.AssumeRoleArn,
			ExternalID: awssdk.String(spec.AssumeExternalId),
		}
	}

	return
}

// FindSubnetIdsByZones returns the subnet IDs for the given zones.
func FindSubnetIdsByZones(subnets []ec2types.Subnet, zones []string) []string {
	var exceptSubnetIds []string
	for _, subnet := range subnets {
		for _, zone := range zones {
			if *subnet.AvailabilityZone == zone {
				exceptSubnetIds = append(exceptSubnetIds, *subnet.SubnetId)
			}
		}
	}
	return exceptSubnetIds
}

// PickSubnetsByEndpointService picks subnets that have same availability zones with the given endpoint service.
// Suppose there are three subnets in a EKS dataplane, and they are in three different availability zones:
//   - subnet-1: us-west-2a
//   - subnet-2: us-west-2b
//   - subnet-3: us-west-2c
//
// If the endpoint service is in us-west-2a, then we can only pass subnet-1 to the VPCEndpoint's spec.
func PickSubnetsByEndpointService(ctx context.Context, cli *ec2.Client, svcName string, subnetIDs []string) ([]string, error) {
	endpointServiceOut, err := cli.DescribeVpcEndpointServices(ctx, &ec2.DescribeVpcEndpointServicesInput{
		ServiceNames: []string{svcName},
	})
	if err != nil {
		return nil, fmt.Errorf("failed to describe vpc endpoint service %s: %v", svcName, err)
	}
	if len(endpointServiceOut.ServiceDetails) != 1 || len(endpointServiceOut.ServiceDetails[0].AvailabilityZones) == 0 {
		return nil, fmt.Errorf("no availability zones for vpc endpoint service %s", svcName)
	}

	subnetOut, err := cli.DescribeSubnets(ctx, &ec2.DescribeSubnetsInput{
		SubnetIds: subnetIDs,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to describe subnets %v: %v", subnetIDs, err)
	}

	res := FindSubnetIdsByZones(subnetOut.Subnets, endpointServiceOut.ServiceDetails[0].AvailabilityZones)
	if len(res) == 0 {
		return nil, fmt.Errorf("no available subnets for the vpc endpoint service %s", svcName)
	}
	return res, nil
}