package privatelinkgroup

import (
	"context"
	"sort"

	awssdk "github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	elbv2types "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types"
	"github.com/pingcap/errors"

	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
)

type LoadBalancerClient interface {
	DescribeLoadBalancers(context.Context, *elasticloadbalancingv2.DescribeLoadBalancersInput, ...func(*elasticloadbalancingv2.Options)) (*elasticloadbalancingv2.DescribeLoadBalancersOutput, error)
}

type NLBObservation struct {
	ARN         string
	VPCID       string
	State       elbv2types.LoadBalancerStateEnum
	AddressType v1alpha1.IPAddressType
	SubnetIDs   []string
}

func DescribeNLB(ctx context.Context, cli LoadBalancerClient, arn string) (NLBObservation, error) {
	if cli == nil {
		return NLBObservation{}, errors.New("AWS load balancer client is not initialized")
	}
	out, err := cli.DescribeLoadBalancers(ctx, &elasticloadbalancingv2.DescribeLoadBalancersInput{
		LoadBalancerArns: []string{arn},
	})
	if err != nil {
		return NLBObservation{}, err
	}
	if out == nil || len(out.LoadBalancers) != 1 {
		return NLBObservation{}, errors.Errorf("AWS did not return PrivateLink NLB %s", arn)
	}
	lb := out.LoadBalancers[0]
	var state elbv2types.LoadBalancerStateEnum
	if lb.State != nil {
		state = lb.State.Code
	}
	addressType, err := addressTypeFromAWS(lb.IpAddressType)
	if err != nil {
		return NLBObservation{}, errors.Wrapf(err, "PrivateLink NLB %s", arn)
	}
	subnetIDs := make([]string, 0, len(lb.AvailabilityZones))
	for _, zone := range lb.AvailabilityZones {
		if subnetID := awssdk.ToString(zone.SubnetId); subnetID != "" {
			subnetIDs = append(subnetIDs, subnetID)
		}
	}
	sort.Strings(subnetIDs)
	return NLBObservation{
		ARN: arn, VPCID: awssdk.ToString(lb.VpcId),
		State: state, AddressType: addressType, SubnetIDs: subnetIDs,
	}, nil
}

func addressTypeFromAWS(value elbv2types.IpAddressType) (v1alpha1.IPAddressType, error) {
	switch value {
	case elbv2types.IpAddressTypeIpv4:
		return v1alpha1.IPAddressTypeIPv4, nil
	case elbv2types.IpAddressTypeDualstack:
		return v1alpha1.IPAddressTypeDualstack, nil
	default:
		return "", errors.Errorf("unsupported IP address type %q", value)
	}
}