package extragroup

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	"github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
	ec2v1alpha1 "github.com/crossplane-contrib/provider-aws/apis/ec2/v1alpha1"
	route53v1alpha1 "github.com/crossplane-contrib/provider-aws/apis/route53/v1alpha1"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/crossplane/crossplane-runtime/pkg/meta"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/eks-provider/controllers/cluster/task"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/controllers/common"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

const (
	defaultDNSTTLSeconds = 300

	// This annotations mark whether allowed principals are set
	annoKeyPrivateLinkAllowedPrincipals    = "private-link.aws.cloud.pingcap.com/allowed-principals"
	annoValPrivateLinkAllowedPrincipalsAll = "*"
)

type PrivateLink struct {
	*task.BaseTask

	AWSCli aws.Client
}

func NewPrivateLink(ctx *task.Context, awsCli aws.Client) *PrivateLink {
	return &PrivateLink{BaseTask: &task.BaseTask{Context: ctx}, AWSCli: awsCli}
}

func (t *PrivateLink) Name() string {
	return "tidbgroup.PrivateLink"
}

func (t *PrivateLink) Create(ctx context.Context) deploy.Result {
	if t.IsRegionalManaged() {
		return deploy.Done()
	}
	if t.Config.DryRun.Enabled {
		return deploy.Done()
	}

	remote := t.EnvProvider.GetClient()

	errList := []error{}
	groupNames := []string{}

	// NOTE: length of spec and status will always be equal
	// because status have been normalized in Service task
	for i := range t.Cluster.Spec.ExtraGroups {
		group := &t.Cluster.Spec.ExtraGroups[i]
		status := &t.Cluster.Status.ExtraEndpoints[i]

		if group.Deleted() || group.SqlEndpoints.PrivateLinkService == nil || group.SqlEndpoints.PrivateLinkService.AWS == nil {
			status.SqlEndpointsStatus.PrivateLinkService = nil
			continue
		}
		groupNames = append(groupNames, group.Name)

		if err := t.createExtraGroup(ctx, remote, group, status); err != nil {
			errList = append(errList, err)
		}
	}

	if len(errList) != 0 {
		return deploy.Error(errors.NewAggregate(errList))
	}

	if err := t.deleteExtraGroups(ctx, remote, groupNames...); err != nil {
		return deploy.Error(err)
	}

	return deploy.Done()
}

func (t *PrivateLink) Delete(ctx context.Context) deploy.Result {
	if t.IsRegionalManaged() {
		return deploy.Done()
	}
	if t.Config.DryRun.Enabled {
		return deploy.Done()
	}

	remote := t.EnvProvider.GetClient()

	// delete all extra groups
	if err := t.deleteExtraGroups(ctx, remote); err != nil {
		return deploy.Error(err)
	}

	return deploy.Done()
}

func (t *PrivateLink) createExtraGroup(ctx context.Context, remote client.Client, group *cloudv1.ExtraGroupSpec, gStatus *cloudv1.ExtraGroupStatus) error {

	n := t.TidbGroupNamer(group)
	svc := corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      n.GetPrivateTiDBServiceName(),
			Namespace: n.GetTiDBNamespace(),
		},
	}

	plStatus := gStatus.SqlEndpointsStatus.PrivateLinkService
	if plStatus == nil {
		plStatus = &cloudv1.PrivateLinkServiceStatus{
			Phase: cloudv1.PrivateLinkServicePhaseCreating,
			AWS:   &cloudv1.AWSPrivateLinkServiceStatus{},
		}
		gStatus.SqlEndpointsStatus.PrivateLinkService = plStatus
	}

	if err := remote.Get(ctx, client.ObjectKeyFromObject(&svc), &svc); err != nil {
		return err
	}

	nlbARN, ok := svc.Annotations[annoKeyLoadBalancerARN]
	if !ok {
		return fmt.Errorf("wait for private svc %s ready", client.ObjectKeyFromObject(&svc))
	}

	vesc := t.NewVPCEndpointServiceConfiguration(group, n, nlbARN)
	if err := k8s.Apply(ctx, t.Local, vesc); err != nil {
		return err
	}

	plStatus.CreationTimestamp = vesc.CreationTimestamp

	sc := vesc.Status.AtProvider.ServiceConfiguration
	if sc == nil || sc.ServiceID == nil {
		return fmt.Errorf("wait for vpc endpoint service setup")
	}

	svcID := *sc.ServiceID

	plStatus.AWS.NLB = nlbARN
	plStatus.AWS.ServiceID = svcID
	plStatus.AWS.ServiceName = *sc.ServiceName
	plStatus.AWS.PrivateDNSName = n.GetPrivateLinkEndpointAddress()
	plStatus.AWS.Port = pointer.Int32(group.SqlEndpoints.GetPrivateSQLEndpointPort())

	// only update az ids if length of azs is not matched
	// NOTE: we get az ids to keep same as tasks/privatelink.go
	if len(plStatus.AWS.AvailabilityZones) != len(sc.AvailabilityZones) {
		azs := []string{}
		for _, az := range sc.AvailabilityZones {
			azs = append(azs, *az)
		}
		ac := t.AWSCli.EC2(t.GetAWSAssumedRoleForSDK())
		res, err := ac.DescribeAvailabilityZones(ctx, &ec2.DescribeAvailabilityZonesInput{
			ZoneNames: azs,
		})
		if err != nil {
			return fmt.Errorf("cannot describe zones: %w", err)
		}

		azIDs := []string{}
		for _, az := range res.AvailabilityZones {
			azIDs = append(azIDs, *az.ZoneId)
		}
		plStatus.AWS.AvailabilityZones = azIDs
	}

	if err := t.setAllowedPrincipals(ctx, vesc, svcID); err != nil {
		return err
	}

	if *sc.PrivateDNSNameConfiguration.State != string(ec2types.DnsNameStateVerified) {
		rrs := t.NewDNSRecord(group.Name, n, sc.PrivateDNSNameConfiguration)
		if err := k8s.Apply(ctx, t.Local, rrs); err != nil {
			return err
		}

		if rrs.Status.GetCondition(xpv1.TypeReady).Status != corev1.ConditionTrue {
			return fmt.Errorf("wait for dns record ready for verification")
		}

		if err := t.verifyDNSOwnership(ctx, svcID); err != nil {
			return err
		}

		return fmt.Errorf("wait for dns ownership verified")
	}

	if sc.ServiceState == nil || *sc.ServiceState != string(ec2types.ServiceStateAvailable) {
		return fmt.Errorf("wait for vpc endpoint service available")
	}

	plStatus.Phase = cloudv1.PrivateLinkServicePhaseActive

	conns, err := t.getAllVPCEndpointConns(ctx, svcID)
	if err != nil {
		return err
	}

	allowed := group.SqlEndpoints.PrivateLinkService.AWS.AcceptedEndpoints
	connsStatus, unrejected := t.calculateConnsStatus(allowed, conns)

	unaccepted := getUnacceptedVPCEndpoints(connsStatus)

	plStatus.AWS.Connections = connsStatus

	if err := t.acceptVPCEndpoints(ctx, svcID, unaccepted...); err != nil {
		return err
	}

	// reject all untracked endpoints
	if err := t.rejectVPCEndpoints(ctx, svcID, unrejected...); err != nil {
		return err
	}

	if len(unaccepted)+len(unrejected) != 0 {
		return fmt.Errorf("wait for vpc endpoint connections synced")
	}

	return nil
}

func (t *PrivateLink) deleteExtraGroups(ctx context.Context, remote client.Client, keep ...string) error {
	dnsList := route53v1alpha1.ResourceRecordSetList{}
	if err := t.Local.List(ctx, &dnsList, client.MatchingLabelsSelector{
		Selector: t.childrenSelector(""),
	}); err != nil {
		return err
	}
	for _, record := range dnsList.Items {
		group, ok := record.GetLabels()[labelKeyExtraGroup]
		if !ok {
			continue
		}

		if groupIsContained(group, keep...) {
			continue
		}

		t.Logger.Info("delete dns record", "group", group)

		if err := t.Local.Delete(ctx, &record); err != nil {
			return err
		}
	}

	vescList := ec2v1alpha1.VPCEndpointServiceConfigurationList{}
	if err := t.Local.List(ctx, &vescList,
		client.MatchingLabelsSelector{
			Selector: t.childrenSelector(""),
		},
	); err != nil {
		return err
	}
	for _, vesc := range vescList.Items {
		group, ok := vesc.GetLabels()[labelKeyExtraGroup]
		if !ok {
			continue
		}

		if groupIsContained(group, keep...) {
			continue
		}

		t.Logger.Info("delete service configuration", "group", group)

		svcID := ""
		sc := vesc.Status.AtProvider.ServiceConfiguration
		if sc != nil && sc.ServiceID != nil {
			svcID = *sc.ServiceID
		}

		if err := t.Local.Delete(ctx, &vesc); err != nil {
			return err
		}

		if svcID != "" {
			if err := t.rejectAllVPCEndpoints(ctx, svcID); err != nil {
				return err
			}
		}
	}

	return nil
}

func (t *PrivateLink) verifyDNSOwnership(ctx context.Context, svcID string) error {
	c := t.AWSCli.EC2(t.GetAWSAssumedRoleForSDK())
	res, err := c.
		StartVpcEndpointServicePrivateDnsVerification(
			ctx, &ec2.StartVpcEndpointServicePrivateDnsVerificationInput{
				ServiceId: &svcID,
				DryRun:    &t.Config.DryRun.Enabled,
			}, func(opts *ec2.Options) {
				opts.Region = t.EnvProvider.GetRegion()
			})
	if err != nil {
		return err
	}
	if res.ReturnValue == nil || !*res.ReturnValue {
		return fmt.Errorf("verify domain ownership for %s failed", svcID)
	}

	return nil
}

func getUnacceptedVPCEndpoints(connsStatus []*cloudv1.AWSPrivateLinkServiceConnection) []string {
	eps := []string{}
	for _, conn := range connsStatus {
		if conn.Phase == cloudv1.PrivateLinkServiceConnectionPhasePending {
			eps = append(eps, conn.Endpoint)
		}
	}

	return eps
}

func (t *PrivateLink) setAllowedPrincipals(ctx context.Context, vesc *ec2v1alpha1.VPCEndpointServiceConfiguration, svcID string) error {
	val, ok := vesc.Annotations[annoKeyPrivateLinkAllowedPrincipals]
	if ok && val == annoValPrivateLinkAllowedPrincipalsAll {
		return nil
	}

	c := t.AWSCli.EC2(t.GetAWSAssumedRoleForSDK())
	res, err := c.ModifyVpcEndpointServicePermissions(ctx,
		&ec2.ModifyVpcEndpointServicePermissionsInput{
			ServiceId:            &svcID,
			AddAllowedPrincipals: []string{"*"},
			DryRun:               &t.Config.DryRun.Enabled,
		}, func(opts *ec2.Options) {
			opts.Region = t.EnvProvider.GetRegion()
		})
	if err != nil {
		return err
	}
	if res.ReturnValue == nil || !*res.ReturnValue {
		return fmt.Errorf("modify vpc endpoint service permission to '*' failed")
	}

	if err := k8s.PatchObject(ctx, t.Local, vesc, newAllowedPrincipalsPatch()); err != nil {
		return err
	}

	return nil
}

func newAllowedPrincipalsPatch() k8s.ObjectPatch {
	return newAnnoPatch(annoKeyPrivateLinkAllowedPrincipals, annoValPrivateLinkAllowedPrincipalsAll)
}

func (t *PrivateLink) getAllVPCEndpointConns(ctx context.Context, svcID string) ([]ec2types.VpcEndpointConnection, error) {
	c := t.AWSCli.EC2(t.GetAWSAssumedRoleForSDK())
	res, err := c.DescribeVpcEndpointConnections(ctx,
		&ec2.DescribeVpcEndpointConnectionsInput{
			Filters: []ec2types.Filter{
				{
					Name: pointer.String("service-id"),
					Values: []string{
						svcID,
					},
				},
			},
		}, func(opts *ec2.Options) {
			opts.Region = t.EnvProvider.GetRegion()
		})
	if err != nil {
		return nil, err
	}

	return res.VpcEndpointConnections, nil
}

func (t *PrivateLink) acceptVPCEndpoints(ctx context.Context, svcID string, acceptList ...string) error {
	if len(acceptList) == 0 {
		return nil
	}
	c := t.AWSCli.EC2(t.GetAWSAssumedRoleForSDK())
	res, err := c.AcceptVpcEndpointConnections(ctx,
		&ec2.AcceptVpcEndpointConnectionsInput{
			ServiceId:      &svcID,
			VpcEndpointIds: acceptList,
			DryRun:         &t.Config.DryRun.Enabled,
		}, func(opts *ec2.Options) {
			opts.Region = t.EnvProvider.GetRegion()
		})
	if err != nil {
		return err
	}

	if len(res.Unsuccessful) != 0 {
		return fmt.Errorf("accept some vpc endpoints(%v) failed", res.Unsuccessful)
	}

	return nil
}

func (t *PrivateLink) rejectVPCEndpoints(ctx context.Context, svcID string, rejectList ...string) error {
	if len(rejectList) == 0 {
		return nil
	}
	c := t.AWSCli.EC2(t.GetAWSAssumedRoleForSDK())
	res, err := c.RejectVpcEndpointConnections(ctx,
		&ec2.RejectVpcEndpointConnectionsInput{
			ServiceId:      &svcID,
			VpcEndpointIds: rejectList,
			DryRun:         &t.Config.DryRun.Enabled,
		}, func(opts *ec2.Options) {
			opts.Region = t.EnvProvider.GetRegion()
		})
	if err != nil {
		return err
	}

	if len(res.Unsuccessful) != 0 {
		return fmt.Errorf("reject some vpc endpoints(%v) failed", res.Unsuccessful)
	}

	return nil

}

func (t *PrivateLink) rejectAllVPCEndpoints(ctx context.Context, svcID string) error {
	conns, err := t.getAllVPCEndpointConns(ctx, svcID)
	if err != nil {
		return err
	}
	rejectList := []string{}
	for _, conn := range conns {
		if isUnrejected(&conn, false) {
			rejectList = append(rejectList, *conn.VpcEndpointId)
		}
	}

	if err := t.rejectVPCEndpoints(ctx, svcID, rejectList...); err != nil {
		return err
	}

	return nil
}

const (
	epStatePendingAcceptance = "pendingacceptance"
	epStatePending           = "pending"
	epStateAvailable         = "available"
	epStateDeleting          = "deleting"
	epStateDeleted           = "deleted"
	epStateRejected          = "rejected"
	epStateFailed            = "failed"
	epStateExpired           = "expired"
)

// Only connections in allowed list will display the state in status
// State mapping is shown as below:
// - VPC Endpoint should exist before accepting
//   - "not found"       => Failed
//   - deleting          => Failed
//   - deleted           => Failed
//
// - VPC Endpoint should be ready for accepting
//   - pendingAcceptance => Pending
//   - pending           => Pending
//
// - VPC Endpoint is available
//   - available         => Active
//
// - Rejected VPC Endpoint can be re-accepted, so just display Pending state
//   - rejected          => Pending
//
// - NOTE: These state are hard to produce and test, so just set them to Failed
//   - failed            => Failed
//   - expired           => Failed
func (t *PrivateLink) calculateConnsStatus(allowed []string, conns []ec2types.VpcEndpointConnection) ([]*cloudv1.AWSPrivateLinkServiceConnection, []string) {
	connsMap := map[string]*ec2types.VpcEndpointConnection{}
	allowedMap := map[string]struct{}{}
	for i := range conns {
		connsMap[*conns[i].VpcEndpointId] = &conns[i]
	}

	connsStatus := []*cloudv1.AWSPrivateLinkServiceConnection{}
	unrejected := []string{}

	for _, item := range allowed {
		status := cloudv1.AWSPrivateLinkServiceConnection{
			Endpoint: item,
		}

		conn, ok := connsMap[item]
		if !ok {
			// all conns that are not in allowed list will be ignored in status
			status.Phase = cloudv1.PrivateLinkServiceConnectionPhaseFailed
			status.Message = fmt.Sprintf("The Vpc Endpoint Id '%s' does not exist", item)
			connsStatus = append(connsStatus, &status)
			continue
		}

		state := strings.ToLower(string(conn.VpcEndpointState))
		switch state {
		case epStateAvailable:
			status.Phase = cloudv1.PrivateLinkServiceConnectionPhaseActive
		case epStatePending, epStatePendingAcceptance, epStateRejected:
			status.Phase = cloudv1.PrivateLinkServiceConnectionPhasePending
		case epStateDeleting, epStateDeleted:
			status.Phase = cloudv1.PrivateLinkServiceConnectionPhaseFailed
			status.Message = fmt.Sprintf("The Vpc Endpoint '%s' is %s", status.Endpoint, state)
		case epStateFailed, epStateExpired:
			status.Phase = cloudv1.PrivateLinkServiceConnectionPhaseFailed
			status.Message = fmt.Sprintf("The Vpc Endpoint '%s' is failed to be active", status.Endpoint)
		default:
			status.Phase = cloudv1.PrivateLinkServiceConnectionPhaseFailed
			status.Message = fmt.Sprintf("The Vpc Endpoint '%s' is in unknown state %s", status.Endpoint, state)
		}

		allowedMap[item] = struct{}{}
		connsStatus = append(connsStatus, &status)
	}

	for _, conn := range conns {
		if _, ok := allowedMap[*conn.VpcEndpointId]; !ok {
			if isUnrejected(&conn, true) {
				unrejected = append(unrejected, *conn.VpcEndpointId)
			}
		}
	}

	return connsStatus, unrejected
}

func isUnrejected(conn *ec2types.VpcEndpointConnection, ignorePending bool) bool {
	state := strings.ToLower(string(conn.VpcEndpointState))
	switch state {
	case epStateAvailable:
		return true
	case epStatePending, epStatePendingAcceptance:
		if ignorePending {
			return false
		}
		return true
	case epStateDeleting, epStateDeleted, epStateRejected, epStateFailed, epStateExpired:
	}

	return false
}

func (t *PrivateLink) NewVPCEndpointServiceConfiguration(group *cloudv1.ExtraGroupSpec, n *namer.Namer, nlbARN string) *ec2v1alpha1.VPCEndpointServiceConfiguration {
	dnsName := n.GetPrivateLinkEndpointAddress()

	tags := util.GetEC2V1AlphaTags(common.ConcatenateTags(cloudv1.TagsFromLabels(t.Cluster.Labels), map[string]string{
		common.Cluster:   t.Cluster.Name,
		common.Component: common.Loadbalance,
		common.UsedBy:    common.CustomerTiDB,
	}))

	vesc := &ec2v1alpha1.VPCEndpointServiceConfiguration{
		ObjectMeta: metav1.ObjectMeta{
			// use dns name
			Name:   dnsName,
			Labels: t.childrenLabels(group.Name),
		},
		Spec: ec2v1alpha1.VPCEndpointServiceConfigurationSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.GetAWSProviderConfigName(),
				},
			},
			ForProvider: ec2v1alpha1.VPCEndpointServiceConfigurationParameters{
				Region:             t.EnvProvider.GetRegion(),
				AcceptanceRequired: pointer.Bool(true),
				PrivateDNSName:     &dnsName,
				CustomVPCEndpointServiceConfigurationParameters: ec2v1alpha1.
					CustomVPCEndpointServiceConfigurationParameters{
					NetworkLoadBalancerARNs: []*string{
						&nlbARN,
					},
				},
				TagSpecifications: []*ec2v1alpha1.TagSpecification{
					{
						ResourceType: pointer.String("vpc-endpoint-service"),
						Tags:         tags,
					},
				},
			},
		},
	}

	return vesc
}

func (t *PrivateLink) NewDNSRecord(groupName string, n *namer.Namer, dnsConfig *ec2v1alpha1.PrivateDNSNameConfiguration) *route53v1alpha1.ResourceRecordSet {
	dnsName := *dnsConfig.Name + "." + t.Config.FullManage.PublicHostedZoneDomain
	rrs := &route53v1alpha1.ResourceRecordSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:   n.GetPrivateLinkEndpointAddress(),
			Labels: t.childrenLabels(groupName),
			Annotations: map[string]string{
				meta.AnnotationKeyExternalName: dnsName,
			},
		},
		Spec: route53v1alpha1.ResourceRecordSetSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.Config.ControlplaneAWSProviderConfig,
				},
			},
			ForProvider: route53v1alpha1.ResourceRecordSetParameters{
				Type: *dnsConfig.Type,
				TTL:  pointer.Int64(defaultDNSTTLSeconds),
				ResourceRecords: []route53v1alpha1.ResourceRecord{
					{
						// TXT record value should be enclosed in quotation marks,
						// for example, "Sample text entry".
						Value: strconv.Quote(*dnsConfig.Value),
					},
				},
				ZoneID: pointer.String(t.Config.FullManage.PublicHostedZoneID),
			},
		},
	}

	return rrs
}

func (t *PrivateLink) childrenLabels(groupName string) map[string]string {
	return childrenLabels(t.Cluster, labelValOwnedByPrivateLink, groupName)
}

func (t *PrivateLink) childrenSelector(groupName string) labels.Selector {
	return childrenSelector(t.Cluster, labelValOwnedByPrivateLink, groupName)
}