package eks

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

	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
	elbv2 "github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2"
	"github.com/aws/aws-sdk-go/aws"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	awsutil "github.com/tidbcloud/infra-provider/regional-server/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/consts"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/loadbalancer/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/deploy"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/k8s"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	targetGroupAttributesKey = "service.beta.kubernetes.io/aws-load-balancer-target-group-attributes"
	subnetsAnnotationKey     = "service.beta.kubernetes.io/aws-load-balancer-subnets"
)

const gracefulUnhealthyDrainingIntervalSeconds = 0

// Service create a service for the LB
type Service struct {
	*task.BaseTask
	awsCli awsutil.Client
}

func NewService(ctx *task.Context, awsCli awsutil.Client) deploy.Task {
	return &Service{
		BaseTask: &task.BaseTask{Context: ctx},
		awsCli:   awsCli,
	}
}

func (t *Service) Name() string {
	return "Service"
}

func (t *Service) Create(ctx context.Context) deploy.Result {
	t.BeginAt = util.Pointer(time.Now())
	remoteCli, err := t.GetRemoteClient(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	remoteReader, err := t.GetRemoteAPIReader(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	svc := t.BuildBasicService()

	var oldSvc *corev1.Service
	found := &corev1.Service{}
	if err := remoteCli.Get(ctx, client.ObjectKeyFromObject(svc), found); err != nil {
		if !apierrors.IsNotFound(err) {
			return deploy.Error(fmt.Errorf("failed to get service %s: %w", svc.Name, err))
		}
	} else {
		oldSvc = found
	}

	if t.LoadBalancer.Spec.Type == v1alpha1.PublicLoadBalancerType {
		t.overlayPublicSvc(svc, oldSvc)
	} else {
		t.overlayPrivateSvc(svc, oldSvc)
	}
	t.Logger.Info("apply load balancer service desired state",
		serviceApplyDesiredStateLogValues(t.LoadBalancer, svc)...,
	)
	// Use the uncached reader for the compare step so rapid allowlist changes
	// are matched against the latest Service from the apiserver, not a stale
	// cache snapshot.
	result, err := k8s.ApplyWithResultUsingReader(ctx, remoteReader, remoteCli, svc)
	if err != nil {
		return deploy.PanicOrDone(err)
	}
	t.Logger.Info("apply load balancer service result",
		"loadBalancerName", t.LoadBalancer.Name,
		"serviceName", svc.Name,
		"serviceNamespace", svc.Namespace,
		"applyResult", result,
	)
	actual := &corev1.Service{}
	if err := remoteReader.Get(ctx, client.ObjectKeyFromObject(svc), actual); err != nil {
		t.Logger.Error(err, "failed to read back load balancer service after apply",
			"loadBalancerName", t.LoadBalancer.Name,
			"serviceName", svc.Name,
			"serviceNamespace", svc.Namespace,
		)
		return deploy.Done()
	}
	t.Logger.Info("load balancer service actual state after apply",
		serviceApplyActualStateLogValues(t.LoadBalancer.Name, actual)...,
	)
	return deploy.Done()
}

func serviceApplyDesiredStateLogValues(lb *v1alpha1.LoadBalancer, svc *corev1.Service) []any {
	return []any{
		"loadBalancerName", lb.Name,
		"loadBalancerResourceVersion", lb.ResourceVersion,
		"serviceName", svc.Name,
		"serviceNamespace", svc.Namespace,
		"loadBalancerIPWhiteListCount", len(lb.Spec.IpWhiteList),
		"desiredLoadBalancerSourceRangeCount", len(svc.Spec.LoadBalancerSourceRanges),
	}
}

func serviceApplyActualStateLogValues(loadBalancerName string, svc *corev1.Service) []any {
	return []any{
		"loadBalancerName", loadBalancerName,
		"serviceName", svc.Name,
		"serviceNamespace", svc.Namespace,
		"actualLoadBalancerSourceRangeCount", len(svc.Spec.LoadBalancerSourceRanges),
	}
}

func (t *Service) PostCreate(ctx context.Context) deploy.Result {
	if t.VPC == nil {
		return deploy.Error(fmt.Errorf("vpc is not found"))
	}
	if t.privateIPsAreAllFound(t.LoadBalancer.Status.PrivateIPs) {
		return deploy.Done()
	}
	role := &awsutil.AssumedRole{
		RoleARN: getAWSAssumedRoleForSDK(t.VPC),
	}
	svc := t.BuildServiceMeta(t.LoadBalancer)
	elbCli := t.awsCli.ELBV2(role)
	ec2Cli := t.awsCli.EC2(role)
	lbRep, err := elbCli.DescribeLoadBalancers(ctx, &elbv2.DescribeLoadBalancersInput{
		Names: []string{t.GetNLBName(t.Config.ResourceNamePrefix, svc.ObjectMeta)},
	})
	if err != nil {
		return deploy.Error(fmt.Errorf("describe load balancer %s: %w", svc.Name, err))
	}
	if len(lbRep.LoadBalancers) == 0 {
		return deploy.Error(fmt.Errorf("load balancer %s not found", svc.Name))
	}
	lb := lbRep.LoadBalancers[0]
	lbARN := *lb.LoadBalancerArn

	remoteCli, err := t.GetRemoteClient(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	if err := k8s.PatchObject(ctx, remoteCli, svc, newARNPatch(lbARN)); err != nil {
		return deploy.Error(fmt.Errorf("patch service %s: %w", svc.Name, err))
	}
	t.LoadBalancer.Status.CloudResourceId = lbARN

	eniResp, err := ec2Cli.DescribeNetworkInterfaces(ctx, &ec2.DescribeNetworkInterfacesInput{
		Filters: []types.Filter{
			{
				Name:   aws.String("description"),
				Values: []string{fmt.Sprintf("ELB net/%s/*", *lb.LoadBalancerName)},
			},
		},
	})
	if err != nil {
		return deploy.Error(fmt.Errorf("describe network interfaces for load balancer %s: %w", svc.Name, err))
	}
	var ips []string
	for _, eni := range eniResp.NetworkInterfaces {
		if eni.PrivateIpAddress != nil {
			ips = append(ips, *eni.PrivateIpAddress)
		}
	}
	if !t.privateIPsAreAllFound(ips) {
		return deploy.Error(fmt.Errorf("not all private ips are found, current %v", len(ips)))
	}
	// update only when all ips are found
	t.LoadBalancer.Status.PrivateIPs = ips

	return deploy.Done()
}

func (t *Service) privateIPsAreAllFound(ips []string) bool {
	if len(ips) == 0 {
		return false
	}
	total := len(t.VPC.Status.PrivateSubnets)
	if t.LoadBalancer.Spec.Type == v1alpha1.PublicLoadBalancerType {
		total = len(t.VPC.Status.PublicSubnets)
	}

	return len(ips) == total
}

func (t *Service) Delete(ctx context.Context) deploy.Result {
	remoteCli, err := t.GetRemoteClient(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	svc := t.BuildServiceMeta(t.LoadBalancer)
	if err := remoteCli.Get(ctx, client.ObjectKeyFromObject(svc), svc); err != nil {
		if client.IgnoreNotFound(err) != nil {
			return deploy.Error(fmt.Errorf("get service %s: %w", svc.Name, err))
		}
		return deploy.Done()
	}
	if _, ok := svc.Annotations["service.beta.kubernetes.io/aws-load-balancer-attributes"]; ok {
		origin := svc.DeepCopy()
		svc.Annotations["service.beta.kubernetes.io/aws-load-balancer-attributes"] = "load_balancing.cross_zone.enabled=true,deletion_protection.enabled=false"
		if err := remoteCli.Patch(ctx, svc, client.MergeFrom(origin)); err != nil {
			return deploy.Error(fmt.Errorf("patch service %s: %w", svc.Name, err))
		}
	}
	return deploy.PanicOrDone(k8s.DeleteObject(ctx, remoteCli, svc))
}

func (t *Service) overlayPrivateSvc(svc *corev1.Service, oldSvc *corev1.Service) {
	tags := t.EffectiveCloudTags(consts.ProviderTypeEKS)

	// See: https://docs.google.com/document/d/1Hn-smPLaeJOKCeZueQmLBadI4nSyT1nFmjyoVesc5vg/edit?usp=sharing
	targetAnnotations := computeTargetGroupAttributes(
		oldSvc,
		v1alpha1.ResolvePreserveClientIP(t.LoadBalancer.Spec.Type, t.LoadBalancer.Spec.PreserveClientIP),
		t.Config.DisableProxyProtocol,
		t.LoadBalancer.Spec.GracefulUnhealthyEnabled,
	)

	// See https://kubernetes-sigs.github.io/aws-load-balancer-controller/v2.4/guide/service/annotations
	annotations := map[string]string{
		"service.beta.kubernetes.io/aws-load-balancer-name":                     t.GetNLBName(t.Config.ResourceNamePrefix, svc.ObjectMeta),
		"service.beta.kubernetes.io/aws-load-balancer-type":                     "external",
		"service.beta.kubernetes.io/aws-load-balancer-nlb-target-type":          "ip",
		"service.beta.kubernetes.io/aws-load-balancer-scheme":                   "internal",
		"service.beta.kubernetes.io/aws-load-balancer-additional-resource-tags": cloudv1.TagStringForAWS(tags),
		"service.beta.kubernetes.io/aws-load-balancer-target-group-attributes":  targetAnnotations,
		"service.beta.kubernetes.io/aws-load-balancer-attributes":               "load_balancing.cross_zone.enabled=true,deletion_protection.enabled=true",

		"service.beta.kubernetes.io/aws-load-balancer-inbound-sg-rules-on-private-link-traffic": "off",
	}
	if subnets := t.nodePrivateSubnetIDs(); len(subnets) > 0 {
		annotations[subnetsAnnotationKey] = strings.Join(subnets, ",")
	}
	for k, v := range annotations {
		svc.Annotations[k] = v
	}
	// Caller-provided annotations win: merged last so they override the
	// controller defaults above (e.g. a custom target-group-attributes
	// string or healthcheck tuning).
	for k, v := range t.LoadBalancer.Spec.Annotations {
		svc.Annotations[k] = v
	}
}

func (t *Service) nodePrivateSubnetIDs() []string {
	if t.VPC == nil {
		return nil
	}
	nodeZones := t.VPC.Spec.NodeZones
	if len(nodeZones) == 0 || len(t.VPC.Status.PrivateSubnets) == 0 {
		return nil
	}

	zoneToSubnetID := make(map[string]string, len(t.VPC.Status.PrivateSubnets))
	for _, subnet := range t.VPC.Status.PrivateSubnets {
		if subnet.AvailabilityZone == "" || subnet.SubnetID == "" {
			continue
		}
		zoneToSubnetID[subnet.AvailabilityZone] = subnet.SubnetID
	}

	subnetIDs := make([]string, 0, len(nodeZones))
	for _, zone := range nodeZones {
		subnetID, ok := zoneToSubnetID[zone]
		if !ok {
			t.Logger.Info("skip setting load balancer subnets annotation due to incomplete node subnet mapping", "zone", zone)
			return nil
		}
		subnetIDs = append(subnetIDs, subnetID)
	}

	return subnetIDs
}

func (t *Service) overlayPublicSvc(svc *corev1.Service, oldSvc *corev1.Service) {
	// AWS specific annotations
	// tags use the shared tags with cluster
	tags := t.EffectiveCloudTags(consts.ProviderTypeEKS)

	targetAnnotations := computeTargetGroupAttributes(
		oldSvc,
		v1alpha1.ResolvePreserveClientIP(t.LoadBalancer.Spec.Type, t.LoadBalancer.Spec.PreserveClientIP),
		t.Config.DisableProxyProtocol,
		t.LoadBalancer.Spec.GracefulUnhealthyEnabled,
	)

	annotations := map[string]string{
		"service.beta.kubernetes.io/aws-load-balancer-name":                     t.GetNLBName(t.Config.ResourceNamePrefix, svc.ObjectMeta),
		"service.beta.kubernetes.io/aws-load-balancer-type":                     "external",
		"service.beta.kubernetes.io/aws-load-balancer-nlb-target-type":          "ip",
		"service.beta.kubernetes.io/aws-load-balancer-scheme":                   "internet-facing",
		"service.beta.kubernetes.io/aws-load-balancer-additional-resource-tags": cloudv1.TagStringForAWS(tags),
		"service.beta.kubernetes.io/aws-load-balancer-target-group-attributes":  targetAnnotations,
		"service.beta.kubernetes.io/aws-load-balancer-attributes":               "load_balancing.cross_zone.enabled=true,deletion_protection.enabled=true",
	}
	if subnets := t.nodePublicSubnetIDs(); len(subnets) > 0 {
		annotations[subnetsAnnotationKey] = strings.Join(subnets, ",")
	}
	for k, v := range annotations {
		svc.Annotations[k] = v
	}
	// Caller-provided annotations win: merged last so they override the
	// controller defaults above.
	for k, v := range t.LoadBalancer.Spec.Annotations {
		svc.Annotations[k] = v
	}
}

func (t *Service) nodePublicSubnetIDs() []string {
	if t.VPC == nil {
		return nil
	}
	nodeZones := t.VPC.Spec.NodeZones
	if len(nodeZones) == 0 || len(t.VPC.Status.PublicSubnets) == 0 {
		return nil
	}

	zoneToSubnetID := make(map[string]string, len(t.VPC.Status.PublicSubnets))
	for _, subnet := range t.VPC.Status.PublicSubnets {
		if subnet.AvailabilityZone == "" || subnet.SubnetID == "" {
			continue
		}
		zoneToSubnetID[subnet.AvailabilityZone] = subnet.SubnetID
	}

	subnetIDs := make([]string, 0, len(nodeZones))
	for _, zone := range nodeZones {
		subnetID, ok := zoneToSubnetID[zone]
		if !ok {
			t.Logger.Info("skip setting public load balancer subnets annotation due to incomplete node subnet mapping", "zone", zone)
			return nil
		}
		subnetIDs = append(subnetIDs, subnetID)
	}

	return subnetIDs
}

func computeTargetGroupAttributes(oldSvc *corev1.Service, preserveClientIP, disableProxyProtocol, gracefulUnhealthyEnabled bool) string {
	attributes := map[string]string{
		"preserve_client_ip.enabled": fmt.Sprintf("%t", preserveClientIP),
	}

	if oldSvc == nil {
		if !disableProxyProtocol {
			attributes["proxy_protocol_v2.enabled"] = "true"
			// This attribute can be set ONLY after contacting AWS Support to have it authorized at the REGION scope.
			// For more details, see: https://pingcap.feishu.cn/wiki/SI9EwimDXijOH6k4bhrcb45VnSb#share-ZhQOdUdrEogldVxl5tTcHDU0nrf
			attributes["proxy_protocol_v2.client_to_server.header_placement"] = "on_first_ack"
		}
	} else if v, ok := oldSvc.Annotations[targetGroupAttributesKey]; ok {
		attributes = parseTargetGroupAttributes(v)
	}

	delete(attributes, "target_health_state.unhealthy.draining_interval_seconds")
	delete(attributes, "target_health_state.unhealthy.connection_termination.enabled")
	if gracefulUnhealthyEnabled {
		attributes["target_health_state.unhealthy.draining_interval_seconds"] = fmt.Sprintf("%d", gracefulUnhealthyDrainingIntervalSeconds)
		attributes["target_health_state.unhealthy.connection_termination.enabled"] = "false"
	}
	return formatTargetGroupAttributes(attributes)
}

func parseTargetGroupAttributes(attrs string) map[string]string {
	attributes := map[string]string{}
	for _, attr := range strings.Split(attrs, ",") {
		key, value, ok := strings.Cut(attr, "=")
		key = strings.TrimSpace(key)
		if !ok || key == "" {
			continue
		}
		attributes[key] = strings.TrimSpace(value)
	}
	return attributes
}

func formatTargetGroupAttributes(attributes map[string]string) string {
	keys := make([]string, 0, len(attributes))
	for key := range attributes {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	attrs := make([]string, 0, len(keys))
	for _, key := range keys {
		attrs = append(attrs, key+"="+attributes[key])
	}
	return strings.Join(attrs, ",")
}

func getAWSAssumedRoleForSDK(vpc *v1alpha1.VPC) string {
	if vpc.Spec.AccountRef != nil {
		return "arn:aws:iam::" + vpc.Spec.AccountRef.AccountID + ":role/" + vpc.Spec.AccountRef.Role
	}
	return ""
}