package privatelink

import (
	"context"
	"fmt"
	"sort"
	"strings"
	"time"

	"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"
	"k8s.io/utils/ptr"

	regionalapiv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/dataplane-manager/v1"
	cloudaws "github.com/tidbcloud/infra-provider/regional-server/pkg/cloud/aws"
	bizerr "github.com/tidbcloud/infra-provider/regional-server/pkg/common/errors"
)

const (
	endpointServiceZonesCacheTTL   = 5 * time.Minute
	endpointServiceZonesAWSTimeout = 30 * time.Second
)

type endpointServiceEC2Client interface {
	DescribeVpcEndpointServices(ctx context.Context, params *ec2.DescribeVpcEndpointServicesInput, optFns ...func(*ec2.Options)) (*ec2.DescribeVpcEndpointServicesOutput, error)
	DescribeAvailabilityZones(ctx context.Context, params *ec2.DescribeAvailabilityZonesInput, optFns ...func(*ec2.Options)) (*ec2.DescribeAvailabilityZonesOutput, error)
}

type endpointServiceEC2Factory func(role *cloudaws.AssumedRole, region string) endpointServiceEC2Client

type privateLinkEndpointServiceZoneResolver struct {
	ec2Factory endpointServiceEC2Factory
	cache      *endpointServiceZonesCache
}

func newPrivateLinkEndpointServiceZoneResolver(awsCli cloudaws.Client) *privateLinkEndpointServiceZoneResolver {
	var factory endpointServiceEC2Factory
	if awsCli != nil {
		factory = func(role *cloudaws.AssumedRole, region string) endpointServiceEC2Client {
			return awsCli.EC2(role, cloudaws.WithRegion(region))
		}
	}
	return newPrivateLinkEndpointServiceZoneResolverWithFactory(factory, endpointServiceZonesCacheTTL, time.Now)
}

func newPrivateLinkEndpointServiceZoneResolverWithFactory(factory endpointServiceEC2Factory, ttl time.Duration, now func() time.Time) *privateLinkEndpointServiceZoneResolver {
	if now == nil {
		now = time.Now
	}
	return &privateLinkEndpointServiceZoneResolver{
		ec2Factory: factory,
		cache:      newEndpointServiceZonesCache(ttl, now),
	}
}

func (r *privateLinkEndpointServiceZoneResolver) Resolve(ctx context.Context, roleARN, externalID, serviceRegion, serviceName string) ([]*regionalapiv1.EndpointServiceAvailabilityZone, error) {
	roleARN = strings.TrimSpace(roleARN)
	externalID = strings.TrimSpace(externalID)
	serviceRegion = strings.TrimSpace(serviceRegion)
	serviceName = strings.TrimSpace(serviceName)
	if roleARN == "" || serviceRegion == "" || serviceName == "" {
		return nil, bizerr.New(bizerr.ParameterInvalid, "role arn, service region and service name are required")
	}
	if r.ec2Factory == nil {
		return nil, bizerr.New(bizerr.InternalServerError, "aws client is not configured")
	}

	key := privateLinkEndpointServiceZonesCacheKey(roleARN, serviceRegion, serviceName)
	return r.cache.resolve(ctx, key, func() ([]*regionalapiv1.EndpointServiceAvailabilityZone, error) {
		// Keep the shared AWS request independent from any single waiter so one canceled caller does not poison others.
		awsCtx, cancel := context.WithTimeout(context.Background(), endpointServiceZonesAWSTimeout)
		defer cancel()
		return r.describe(awsCtx, roleARN, externalID, serviceRegion, serviceName)
	})
}

func privateLinkEndpointServiceZonesCacheKey(roleARN, serviceRegion, serviceName string) string {
	// ExternalID is intentionally excluded: current dataplane config maps each role ARN to one external ID,
	// and AWS service visibility is determined by the assumed role in the target region.
	return roleARN + "\x00" + serviceRegion + "\x00" + serviceName
}

func (r *privateLinkEndpointServiceZoneResolver) describe(ctx context.Context, roleARN, externalID, serviceRegion, serviceName string) ([]*regionalapiv1.EndpointServiceAvailabilityZone, error) {
	assumedRole := &cloudaws.AssumedRole{RoleARN: roleARN}
	if externalID != "" {
		assumedRole.ExternalID = ptr.To(externalID)
	}
	ec2Cli := r.ec2Factory(assumedRole, serviceRegion)
	if ec2Cli == nil {
		return nil, bizerr.New(bizerr.InternalServerError, "ec2 client is not configured")
	}

	serviceOutput, err := ec2Cli.DescribeVpcEndpointServices(ctx, &ec2.DescribeVpcEndpointServicesInput{
		ServiceNames: []string{serviceName},
	})
	if err != nil {
		return nil, bizerr.Wrap(bizerr.AWSClientError, fmt.Sprintf("describe vpc endpoint service %s error", serviceName), err)
	}
	if len(serviceOutput.ServiceDetails) == 0 {
		return nil, bizerr.New(bizerr.ResourceNotFound, fmt.Sprintf("vpc endpoint service %s not found", serviceName))
	}

	serviceDetail := serviceOutput.ServiceDetails[0]
	if !isInterfaceEndpointService(serviceDetail.ServiceType) {
		return nil, bizerr.New(bizerr.ParameterInvalid, fmt.Sprintf("vpc endpoint service %s is not an Interface type service", serviceName))
	}

	zoneNames := serviceDetail.AvailabilityZones
	if len(zoneNames) == 0 {
		return nil, bizerr.New(bizerr.ResourceNotFound, fmt.Sprintf("vpc endpoint service %s has no availability zones", serviceName))
	}

	azOutput, err := ec2Cli.DescribeAvailabilityZones(ctx, &ec2.DescribeAvailabilityZonesInput{
		ZoneNames: zoneNames,
	})
	if err != nil {
		return nil, bizerr.Wrap(bizerr.AWSClientError, "describe availability zones error", err)
	}
	zoneIDByName := make(map[string]string, len(azOutput.AvailabilityZones))
	for _, zone := range azOutput.AvailabilityZones {
		name := aws.ToString(zone.ZoneName)
		id := aws.ToString(zone.ZoneId)
		if name != "" && id != "" {
			zoneIDByName[name] = id
		}
	}

	zones := make([]*regionalapiv1.EndpointServiceAvailabilityZone, 0, len(zoneNames))
	for _, zoneName := range zoneNames {
		zoneID, ok := zoneIDByName[zoneName]
		if !ok {
			return nil, bizerr.New(bizerr.ResourceNotFound, fmt.Sprintf("availability zone ID for %s was not found", zoneName))
		}
		zones = append(zones, &regionalapiv1.EndpointServiceAvailabilityZone{
			AvailabilityZone:   zoneName,
			AvailabilityZoneId: zoneID,
		})
	}
	sort.Slice(zones, func(i, j int) bool {
		if zones[i].AvailabilityZoneId == zones[j].AvailabilityZoneId {
			return zones[i].AvailabilityZone < zones[j].AvailabilityZone
		}
		return zones[i].AvailabilityZoneId < zones[j].AvailabilityZoneId
	})
	return zones, nil
}

func isInterfaceEndpointService(serviceTypes []ec2types.ServiceTypeDetail) bool {
	for _, serviceType := range serviceTypes {
		if serviceType.ServiceType == ec2types.ServiceTypeInterface {
			return true
		}
	}
	return false
}

func copyEndpointServiceAvailabilityZones(zones []*regionalapiv1.EndpointServiceAvailabilityZone) []*regionalapiv1.EndpointServiceAvailabilityZone {
	if zones == nil {
		return nil
	}
	ret := make([]*regionalapiv1.EndpointServiceAvailabilityZone, 0, len(zones))
	for _, zone := range zones {
		if zone == nil {
			ret = append(ret, nil)
			continue
		}
		ret = append(ret, &regionalapiv1.EndpointServiceAvailabilityZone{
			AvailabilityZone:   zone.AvailabilityZone,
			AvailabilityZoneId: zone.AvailabilityZoneId,
		})
	}
	return ret
}