package aws

import (
	"context"
	"fmt"
	"strings"

	ec2api "github.com/aws/aws-sdk-go-v2/service/ec2"
	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/pingcap/errors"
	ec2v1beta1 "github.com/upbound/provider-aws/apis/ec2/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/controllers/kubernetesagent/metadata"
	"github.com/tidbcloud/observability/agent-control-server/pkg/agentbuild"
	awsclient "github.com/tidbcloud/observability/agent-control-server/pkg/cloud/aws"
	"github.com/tidbcloud/observability/agent-control-server/pkg/cloud/aws/awsupbound"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/deploy"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

type PrivateLinkTask struct {
	*reconciler.BaseTask[*agentv2.KubernetesAgent]
}

func NewPrivateLinkTask(ctx *reconciler.Context[*agentv2.KubernetesAgent]) *PrivateLinkTask {
	return &PrivateLinkTask{
		BaseTask: &reconciler.BaseTask[*agentv2.KubernetesAgent]{Context: ctx},
	}
}

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

func (t *PrivateLinkTask) ClusterMeta() deploy.ClusterMeta {
	return metadata.ClusterMeta(t.Agent)
}

func (t *PrivateLinkTask) buildVPCEndpointTags(endpointName string) map[string]*string {
	return map[string]*string{
		consts.CrossplaneKeyVendor:    ptr.To(consts.CrossplaneValueVendor),
		consts.CrossplaneKeyName:      ptr.To(endpointName),
		consts.CrossplaneKeyComponent: ptr.To(consts.CrossplaneValueComponent),
		consts.CrossplaneKeyTenant:    ptr.To(util.GetLabel(t.Agent, agentv2.TagKeyTenant)),
	}
}

func (t *PrivateLinkTask) Create(ctx context.Context) deploy.Result {
	needUpdate := util.ShouldUpdateO11yVersion(t.Agent.GetAnnotations(), t.Config.O11yConfig.Version)
	endpoint := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
	if endpoint != "" && !needUpdate {
		return deploy.Done()
	}

	endpointServiceName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yEndpointServiceName)
	if endpointServiceName == "" {
		return deploy.Error(errors.New("endpoint service name not found wait for o11y cluster ready"))
	}

	if t.Config.IsServerlessMode() {
		return t.createWithNativeAPI(ctx, endpointServiceName)
	}
	return t.createWithUpbound(ctx, endpointServiceName)
}

func (t *PrivateLinkTask) createWithNativeAPI(ctx context.Context, endpointServiceName string) deploy.Result {
	cli := t.AWS.EC2(nil)
	endpointID, err := awsclient.FindReusableVPCEndpointID(ctx, cli, string(ec2types.VpcEndpointTypeInterface), endpointServiceName, t.Agent.Spec.EKS.VpcID)
	if err != nil {
		return deploy.Error(errors.Annotate(err, "failed to find reusable vpc endpoint"))
	}
	if endpointID != "" {
		t.L.Info(fmt.Sprintf("found an existing endpoint (%s), reusing it", endpointID))
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName, endpointID)
		return deploy.Done()
	}

	dnsName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yServiceDNSName)
	if dnsName == "" {
		return deploy.Error(errors.New("service dns name not found"))
	}

	endpointName := util.InterfaceEndpointName(t.Agent.ObjectMeta, dnsName, consts.O11yEndpointNamePrefixDataplaneNextGen)
	endpointID, err = t.createInterfaceEndpointNative(ctx, endpointServiceName, endpointName)
	if err != nil {
		return deploy.Error(errors.Annotate(err, "failed to create interface endpoint"))
	}

	// Use endpoint ID as annotation in native API, and endpoint name in upbound
	util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName, endpointID)
	return deploy.Done()
}

func (t *PrivateLinkTask) createInterfaceEndpointNative(ctx context.Context, endpointServiceName, endpointName string) (string, error) {
	t.L.Info("creating interface endpoint",
		"endpointName", endpointName,
		"vpcID", t.Agent.Spec.EKS.VpcID,
		"serviceName", endpointServiceName,
		"securityGroupID", t.Agent.Spec.EKS.SecurityGroupID,
	)

	// Select subnets for interface endpoint
	cli := t.AWS.EC2(nil)
	subnetIDs, err := t.pickSubnets(ctx, cli, endpointServiceName, t.Agent.Spec.EKS.NodeSubnetIDs)
	if err != nil {
		return "", errors.Annotate(err, "failed to pick subnets for interface endpoint")
	}
	if len(subnetIDs) == 0 {
		return "", errors.New("no subnets selected for interface endpoint")
	}

	// Build AWS tags from our tags map
	tags := t.buildVPCEndpointTags(endpointName)
	awsTags := make([]ec2types.Tag, 0, len(tags))
	for k, v := range tags {
		awsTags = append(awsTags, ec2types.Tag{Key: ptr.To(k), Value: v})
	}

	// Create VPC endpoint with the existing EKS security group
	subnetIDs = subnetIDs[:1]
	input := &ec2api.CreateVpcEndpointInput{
		VpcId:             ptr.To(t.Agent.Spec.EKS.VpcID),
		ServiceName:       ptr.To(endpointServiceName),
		VpcEndpointType:   ec2types.VpcEndpointTypeInterface,
		SubnetIds:         subnetIDs,
		SecurityGroupIds:  []string{t.Agent.Spec.EKS.SecurityGroupID},
		PrivateDnsEnabled: ptr.To(true),
		TagSpecifications: []ec2types.TagSpecification{{
			ResourceType: ec2types.ResourceTypeVpcEndpoint,
			Tags:         awsTags,
		}},
	}

	output, err := cli.CreateVpcEndpoint(ctx, input)
	if err != nil {
		return "", errors.Annotate(err, "failed to create interface endpoint")
	}

	if output.VpcEndpoint == nil || output.VpcEndpoint.VpcEndpointId == nil {
		return "", errors.New("vpc endpoint id not found in response")
	}

	endpointID := *output.VpcEndpoint.VpcEndpointId
	t.L.Info(fmt.Sprintf("create interface endpoint (%s) with sgID (%s) success", endpointID, t.Agent.Spec.EKS.SecurityGroupID))
	return endpointID, nil
}

func (t *PrivateLinkTask) createWithUpbound(ctx context.Context, endpointServiceName string) deploy.Result {
	endpointName, err := awsupbound.FindReusableVPCEndpointName(ctx, t.Local, string(ec2types.VpcEndpointTypeInterface), endpointServiceName, t.Agent.Spec.EKS.VpcID)
	if err != nil {
		return deploy.Error(errors.Annotate(err, "failed to find reusable vpc endpoint"))
	}
	if endpointName != "" {
		t.L.Info(fmt.Sprintf("found an existed endpoint (%s), do not need to create new one", endpointName))
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName, endpointName)
		return deploy.Done()
	}

	dnsName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yServiceDNSName)
	if dnsName == "" {
		return deploy.Error(errors.New("service dns name not found"))
	}

	endpointName = util.InterfaceEndpointName(t.Agent.ObjectMeta, dnsName, consts.O11yEndpointNamePrefixDataplaneNextGen)
	err = t.createInterfaceEndpointUpbound(ctx, endpointServiceName, endpointName)
	if err != nil {
		return deploy.Error(errors.Annotate(err, "failed to create interface endpoint"))
	}

	t.L.Info(fmt.Sprintf("create interface endpoint (%s) success", endpointName))
	util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName, endpointName)
	return deploy.Done()
}

func (t *PrivateLinkTask) createInterfaceEndpointUpbound(ctx context.Context, endpointServiceName, interfaceEndpointName string) error {
	t.L.Info("creating interface endpoint", "endpointServiceName", endpointServiceName, "interfaceEndpointName", interfaceEndpointName)
	region := t.Config.O11yConfig.Region
	if region == "" {
		return errors.New("AWS region is not configured")
	}

	// Create an interface endpoint
	endpoint := &ec2v1beta1.VPCEndpoint{
		ObjectMeta: metav1.ObjectMeta{
			Name: interfaceEndpointName,
		},
	}

	// Select subnets for interface endpoint
	cli := t.AWS.EC2(&awsclient.AssumedRole{RoleARN: t.Agent.Spec.EKS.AssumeRoleArn})
	subnetIDs, err := t.pickSubnets(ctx, cli, endpointServiceName, t.Agent.Spec.EKS.NodeSubnetIDs)
	if err != nil {
		return errors.Annotate(err, "failed to pick subnets for interface endpoint")
	}
	if len(subnetIDs) == 0 {
		return errors.New("no subnets selected for interface endpoint")
	}

	// Create subnet associations for the filtered subnet set.
	subnetIDs = subnetIDs[:1]
	for _, subnet := range subnetIDs {
		subnetAssociation := &ec2v1beta1.VPCEndpointSubnetAssociation{
			ObjectMeta: metav1.ObjectMeta{
				Name: fmt.Sprintf("%s-%s", interfaceEndpointName, subnet),
			},
		}
		_, err = controllerutil.CreateOrUpdate(ctx, t.Local, subnetAssociation, func() error {
			subnetAssociation.Spec.ForProvider.Region = ptr.To(region)
			subnetAssociation.Spec.ForProvider.SubnetID = ptr.To(subnet)
			subnetAssociation.Spec.ForProvider.VPCEndpointIDRef = &xpv1.Reference{
				Name: interfaceEndpointName,
			}
			subnetAssociation.Spec.ResourceSpec.ProviderConfigReference = t.Agent.Spec.EKS.ProviderConfigReference
			return nil
		})
		if err != nil {
			return errors.Annotate(err, "failed to create/update interface endpoint subnet association")
		}
	}

	// Create security group association
	sgAssociation := &ec2v1beta1.VPCEndpointSecurityGroupAssociation{
		ObjectMeta: metav1.ObjectMeta{
			Name: fmt.Sprintf("%s-%s", interfaceEndpointName, t.Agent.Spec.EKS.SecurityGroupID),
		},
	}
	_, err = controllerutil.CreateOrUpdate(ctx, t.Local, sgAssociation, func() error {
		sgAssociation.Spec.ForProvider.Region = ptr.To(region)
		sgAssociation.Spec.ForProvider.SecurityGroupID = ptr.To(t.Agent.Spec.EKS.SecurityGroupID)
		sgAssociation.Spec.ForProvider.VPCEndpointIDRef = &xpv1.Reference{
			Name: interfaceEndpointName,
		}
		sgAssociation.Spec.ResourceSpec.ProviderConfigReference = t.Agent.Spec.EKS.ProviderConfigReference
		return nil
	})
	if err != nil {
		return errors.Annotate(err, "failed to create/update interface endpoint security group association")
	}

	// Build tags using helper function
	tags := t.buildVPCEndpointTags(interfaceEndpointName)
	_, err = controllerutil.CreateOrUpdate(ctx, t.Local, endpoint, func() error {
		endpoint.Spec.ProviderConfigReference = t.Agent.Spec.EKS.ProviderConfigReference
		endpoint.Spec.ForProvider.Region = ptr.To(region)
		endpoint.Spec.ForProvider.ServiceName = ptr.To(endpointServiceName)
		endpoint.Spec.ForProvider.VPCID = ptr.To(t.Agent.Spec.EKS.VpcID)
		endpoint.Spec.ForProvider.VPCEndpointType = ptr.To(string(ec2types.VpcEndpointTypeInterface))
		endpoint.Spec.ForProvider.PrivateDNSEnabled = ptr.To(true)
		endpoint.Spec.ForProvider.Tags = tags

		if endpoint.Labels == nil {
			endpoint.Labels = make(map[string]string)
		}
		endpoint.Labels[consts.LabelKeyRelatedCluster] = t.Agent.Spec.KubernetesClusterName
		endpoint.Labels[agentv2.TagKeyTenant] = util.GetLabel(t.Agent, agentv2.TagKeyTenant)
		return nil
	})
	if err != nil {
		return errors.Annotate(err, "failed to create/update interface endpoint")
	}
	return nil
}

func (t *PrivateLinkTask) PostCreate(ctx context.Context) deploy.Result {
	return deploy.Done()
}

func (t *PrivateLinkTask) Delete(ctx context.Context) deploy.Result {
	if t.Config.IsServerlessMode() {
		if t.Agent.GetLabels()[consts.LabelKeyPreservePrivateLink] == "true" {
			return deploy.Done()
		}
		return t.deleteWithNativeAPI(ctx)
	}
	return t.deleteWithUpbound(ctx)
}

func (t *PrivateLinkTask) deleteWithNativeAPI(ctx context.Context) deploy.Result {
	endpointID := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName)
	if endpointID == "" {
		return deploy.Done()
	}

	cli := t.AWS.EC2(nil)
	_, err := cli.DeleteVpcEndpoints(ctx, &ec2api.DeleteVpcEndpointsInput{
		VpcEndpointIds: []string{endpointID},
	})
	if err != nil {
		return deploy.Error(errors.Annotate(err, "failed to delete vpc interface endpoint"))
	}

	t.L.Info(fmt.Sprintf("delete interface endpoint (%s) success", endpointID))
	util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySvcInterfaceEndpointName, "")
	return deploy.Done()
}

func (t *PrivateLinkTask) deleteWithUpbound(ctx context.Context) deploy.Result {
	dnsName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yServiceDNSName)
	if dnsName == "" {
		return deploy.Done()
	}

	endpointName := util.InterfaceEndpointName(t.Agent.ObjectMeta, dnsName, consts.O11yEndpointNamePrefixDataplaneNextGen)
	for _, subnet := range t.Agent.Spec.EKS.NodeSubnetIDs {
		subnetAssociation := &ec2v1beta1.VPCEndpointSubnetAssociation{
			ObjectMeta: metav1.ObjectMeta{
				Name: fmt.Sprintf("%s-%s", endpointName, subnet),
			},
		}
		if err := agentbuild.DeleteObject(ctx, t.Local, subnetAssociation); err != nil {
			return deploy.Error(errors.Annotate(err, "failed to delete vpc interface gateway endpoint subnet association"))
		}
	}

	sgAssociation := &ec2v1beta1.VPCEndpointSecurityGroupAssociation{
		ObjectMeta: metav1.ObjectMeta{
			Name: fmt.Sprintf("%s-%s", endpointName, t.Agent.Spec.EKS.SecurityGroupID),
		},
	}
	if err := agentbuild.DeleteObject(ctx, t.Local, sgAssociation); err != nil {
		return deploy.Error(errors.Annotate(err, "failed to delete vpc interface gateway endpoint security group association"))
	}

	endpoint := &ec2v1beta1.VPCEndpoint{
		ObjectMeta: metav1.ObjectMeta{
			Name: endpointName,
		},
	}
	if err := agentbuild.DeleteObject(ctx, t.Local, endpoint); err != nil {
		return deploy.Error(errors.Annotate(err, "failed to delete vpc interface gateway endpoint"))
	}
	return deploy.Done()
}

func (t *PrivateLinkTask) pickSubnets(ctx context.Context, cli *ec2api.Client, endpointServiceName string, subnetIDs []string) ([]string, error) {
	if len(subnetIDs) == 0 {
		return nil, errors.New("no subnet ids provided")
	}

	serviceOutput, err := cli.DescribeVpcEndpointServices(ctx, &ec2api.DescribeVpcEndpointServicesInput{
		ServiceNames: []string{endpointServiceName},
	})
	if err != nil {
		return nil, errors.Annotate(err, "failed to describe vpc endpoint services")
	}
	if len(serviceOutput.ServiceDetails) == 0 {
		return nil, errors.New("no service details found for the given endpoint service name")
	}

	allowedZones := make(map[string]struct{})
	supportedZones := make([]string, 0)
	for _, detail := range serviceOutput.ServiceDetails {
		for _, zone := range detail.AvailabilityZones {
			if _, exists := allowedZones[zone]; exists {
				continue
			}
			allowedZones[zone] = struct{}{}
			supportedZones = append(supportedZones, zone)
		}
	}
	if len(allowedZones) == 0 {
		return nil, errors.New("endpoint service does not expose availability zones")
	}

	subnetOutput, err := cli.DescribeSubnets(ctx, &ec2api.DescribeSubnetsInput{SubnetIds: subnetIDs})
	if err != nil {
		return nil, errors.Annotate(err, "failed to describe subnets")
	}

	selected := make([]string, 0, len(subnetOutput.Subnets))
	providedSubnets := make([]string, 0, len(subnetOutput.Subnets))
	for _, subnet := range subnetOutput.Subnets {
		var (
			subnetID = ""
			zone     = ""
		)
		if subnet.SubnetId != nil {
			subnetID = *subnet.SubnetId
		}
		if subnet.AvailabilityZone != nil {
			zone = *subnet.AvailabilityZone
		}
		providedSubnets = append(providedSubnets, fmt.Sprintf("%s(%s)", subnetID, zone))
		if subnet.SubnetId == nil || subnet.AvailabilityZone == nil {
			continue
		}
		if _, ok := allowedZones[zone]; ok {
			selected = append(selected, subnetID)
		}
	}

	if len(selected) == 0 {
		return nil, errors.Errorf(
			"no subnets match endpoint service availability zones. endpoint supports zones: [%s]; provided subnets: [%s]",
			strings.Join(supportedZones, ","),
			strings.Join(providedSubnets, ","),
		)
	}
	return selected, nil
}