package selfmanaged

import (
	"context"
	"encoding/base64"
	"fmt"
	"sort"
	"strings"

	"github.com/crossplane-contrib/provider-aws/apis/autoscaling/v1beta1"
	ec2v1alpha1 "github.com/crossplane-contrib/provider-aws/apis/ec2/v1alpha1"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

// Both x86 and arm machines will use bottlerocket for security.
const (
	ssmNameTemplate = "/aws/service/bottlerocket/aws-k8s-%s/%s/latest/image_id"

	K8sVersionAnnotationKey = "tidbcloud.com/self-managed-ng-k8s-version"
)

func GetSSMName(k8sVersion, machineArch string) string {
	switch machineArch {
	case consts.ARM64ArchName:
		return fmt.Sprintf(ssmNameTemplate, k8sVersion, "arm64")
	default:
		return fmt.Sprintf(ssmNameTemplate, k8sVersion, "x86_64")
	}
}

func GetBottlerocketUserData(eksName, endpoint, encodedCA string, nodeLabels map[string]string, taints []corev1.Taint, systemConfiguration string) *string {
	userDataString := getK8sSettings(eksName, endpoint, encodedCA) + getLabelsString(nodeLabels) + getTaintsString(taints) + systemConfiguration
	return pointer.String(base64.StdEncoding.EncodeToString([]byte(userDataString)))
}

func OverwriteNodeLabels(userData string, newNodeLabels map[string]string) *string {
	// First, decode the base64 encoded userdata string
	decodedUserData, err := base64.StdEncoding.DecodeString(userData)
	if err != nil {
		return nil
	}

	// Convert decoded userdata back into a string
	decodedString := string(decodedUserData)

	// Extract the k8s settings and taints, assuming the userData string always has a fixed structure
	k8sSettingsIdx := strings.Index(decodedString, "[settings.kubernetes.node-labels]")
	taintsIdx := strings.Index(decodedString, "[settings.kubernetes.node-taints]")

	k8sSettingsStr := decodedString[:k8sSettingsIdx]

	// Generate new labels string and concatenate with the existing userdata
	newLabelsStr := getLabelsString(newNodeLabels)
	newTaintsStr := decodedString[taintsIdx:]

	// Build new userdata string
	newUserDataString := k8sSettingsStr + newLabelsStr + newTaintsStr

	// Re-encode the userdata string in base64
	encodedUserData := base64.StdEncoding.EncodeToString([]byte(newUserDataString))

	return &encodedUserData
}

// getK8sSettings generates the Kubernetes settings string.
func getK8sSettings(eksName, endpoint, encodedCA string) string {
	return fmt.Sprintf(`[settings.kubernetes]
"cluster-name" = "%s"
"api-server" = "%s"
"cluster-certificate" = "%s"
"image-gc-low-threshold-percent" = 70
"image-gc-high-threshold-percent" = 75
`, eksName, endpoint, encodedCA)
}

// getLabelsString generates the node labels string.
func getLabelsString(nodeLabels map[string]string) string {
	var labelsBuilder strings.Builder
	labelsBuilder.WriteString("[settings.kubernetes.node-labels]\n")

	sortedLabels := make([]string, 0, len(nodeLabels))
	for k, v := range nodeLabels {
		sortedLabels = append(sortedLabels, fmt.Sprintf(`"%s" = "%s"`+"\n", k, v))
	}
	sort.Strings(sortedLabels)
	for _, label := range sortedLabels {
		labelsBuilder.WriteString(label)
	}

	return labelsBuilder.String()
}

// getTaintsString generates the node taints string.
func getTaintsString(taints []corev1.Taint) string {
	var taintsBuilder strings.Builder
	taintsBuilder.WriteString("[settings.kubernetes.node-taints]\n")

	sortedTaints := make([]string, 0, len(taints))
	for _, t := range taints {
		sortedTaints = append(sortedTaints, fmt.Sprintf(`"%s" = "%s:%s"`+"\n", t.Key, t.Value, t.Effect))
	}
	sort.Strings(sortedTaints)
	for _, t := range sortedTaints {
		taintsBuilder.WriteString(t)
	}

	return taintsBuilder.String()
}

// getEKSNameFromARN extracts EKS cluster name from a given Amazon EKS Resource Name (ARN) string.
// The function assumes the input ARN is for an EKS cluster and in the correct format.
//
// Example of an input ARN:
//
//	arn:aws:eks:us-west-2:123456789012:cluster/my-cluster
//
// Example of a returned EKS cluster name:
//
//	"my-cluster"
func getEKSNameFromARN(arn string) string {
	arnParts := strings.Split(arn, ":")
	name := strings.Split(arnParts[len(arnParts)-1], "/")[1]
	return name
}

// MutateLaunchTemplate mutates the LaunchTemplate to be used for the self-managed nodegroup.
func MutateLaunchTemplate(lt *ec2v1alpha1.LaunchTemplate, instanceProfileName, amiID, instanceType string, userData *string, sgIDs []*string) {
	if lt != nil && lt.Spec.ForProvider.LaunchTemplateData != nil {
		data := lt.Spec.ForProvider.LaunchTemplateData

		// Configure EC2 Instance Metadata Service (IMDS) to allow both v1 and v2.
		// We explicitly set HTTPTokens to "optional" to ensure IMDSv1 is available because:
		// 1. Some older versions of TiKV only support IMDSv1 for fetching instance metadata, which is required by BR.
		// 2. Without this explicit setting, certain AMIs default to "required" (IMDSv2 only), which would break compatibility
		// 3. When IMDSv1 becomes deprecated or unavailable in the future, this option should be removed
		//
		// AWS IMDS documentation:
		// - IMDSv2: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html
		// - HTTPTokens parameter: https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_LaunchTemplateInstanceMetadataOptionsRequest.html
		data.MetadataOptions = &ec2v1alpha1.LaunchTemplateInstanceMetadataOptionsRequest{
			HTTPTokens: pointer.String("optional"),
		}

		data.IAMInstanceProfile = &ec2v1alpha1.LaunchTemplateIAMInstanceProfileSpecificationRequest{
			Name: pointer.String(instanceProfileName),
		}

		data.ImageID = pointer.String(amiID)
		data.UserData = userData
		data.InstanceType = pointer.String(instanceType)

		data.NetworkInterfaces = []*ec2v1alpha1.LaunchTemplateInstanceNetworkInterfaceSpecificationRequest{
			{
				DeviceIndex:              pointer.Int64(0),
				AssociatePublicIPAddress: pointer.Bool(false),
				Groups:                   sgIDs,
			},
		}

		// Ref: https://github.com/hashicorp/terraform-provider-aws/issues/4570#issuecomment-612201179
		data.SecurityGroups = nil
		data.SecurityGroupIDs = nil
	}
}

// GetAutoScalingGroup returns an AutoScalingGroup(ASG) object for tidb cluster nodes.
//
// Parameters:
// - name: The name of the ASG
// - region: The AWS Region
// - providerConfig: The Crossplane ProviderConfig CR name
// - k8sVersion: The Kubernetes version of the AMI
// - labels: Kubernetes labels to add to the ASG object
// - tags: AWS tags to apply to the ASG
// - subnetIDs: AWS Subnet IDs for the ASG
// - minSize: Minimum number of instances in the ASG
// - maxSize: Maximum number of instances in the ASG
func GetAutoScalingGroup(
	name, region, providerConfig, k8sVersion string,
	labels map[string]string,
	tags []*v1beta1.Tag,
	subnetIDs []string,
	minSize, maxSize int32,
) *v1beta1.AutoScalingGroup {
	return &v1beta1.AutoScalingGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: labels,
			Annotations: map[string]string{
				K8sVersionAnnotationKey: k8sVersion,
			},
		},
		Spec: v1beta1.AutoScalingGroupSpec{
			ForProvider: v1beta1.AutoScalingGroupParameters{
				Region: region,
				LaunchTemplate: &v1beta1.LaunchTemplateSpecification{
					LaunchTemplateName: pointer.String(name),
				},
				// When we paused the cluster, the `DesiredCapacity` will be set to 0.
				// We need to set it back explicitly when resume the cluster,
				// otherwise will get the error: "Desired capacity:0 must be between the specified min size:1 and max size:1"
				DesiredCapacity:   pointer.Int64(int64(minSize)),
				MinSize:           pointer.Int64(int64(minSize)),
				MaxSize:           pointer.Int64(int64(maxSize)),
				VPCZoneIdentifier: pointer.String(strings.Join(subnetIDs, ",")),
				Tags:              tags,

				// The following parameters align with the managed nodegroup's ASG.
				// Better not to change them unless we have a good reason.
				CapacityRebalance: pointer.Bool(true),
				TerminationPolicies: []*string{
					pointer.String("AllocationStrategy"),
					pointer.String("OldestLaunchTemplate"),
					pointer.String("OldestInstance"),
				},
				HealthCheckGracePeriod: pointer.Int64(15),

				// We have a controller to add/remove the node protection.
				NewInstancesProtectedFromScaleIn: pointer.Bool(false),

				// Add a termination hook so that nodes can handle the termination gracefully.
				// See https://github.com/tidbcloud/infra-provider/pull/2087 for more details.
				LifecycleHookSpecificationList: []*v1beta1.LifecycleHookSpecification{
					{
						LifecycleHookName:   pointer.String("lifecycle-hook-termination"),
						DefaultResult:       pointer.String("CONTINUE"),
						HeartbeatTimeout:    pointer.Int64(300),
						LifecycleTransition: pointer.String("autoscaling:EC2_INSTANCE_TERMINATING"),
					},
				},
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: providerConfig},
			},
		},
	}
}

func GetInstanceType(ctx context.Context, cli client.Client, ng *v1beta1.AutoScalingGroup) *string {
	launchTemplateSpec := ng.Spec.ForProvider.LaunchTemplate

	if launchTemplateSpec != nil && launchTemplateSpec.LaunchTemplateName != nil {
		var launchTemplate ec2v1alpha1.LaunchTemplate
		if err := cli.Get(ctx, types.NamespacedName{Name: *launchTemplateSpec.LaunchTemplateName}, &launchTemplate); err == nil {
			if launchTemplate.Spec.ForProvider.LaunchTemplateData != nil {
				return launchTemplate.Spec.ForProvider.LaunchTemplateData.InstanceType
			}
		}
	}

	return nil
}

func GetInstanceTypeFromStatus(ng *v1beta1.AutoScalingGroup) string {
	for _, instance := range ng.Status.AtProvider.Instances {
		if instance.InstanceType != nil {
			return *instance.InstanceType
		}
	}
	return ""
}

// GetTiDBClusterSelfManagedNodeGroups returns the cluster's self-managed nodegroups that are currently in use.
func GetTiDBClusterSelfManagedNodeGroups(ctx context.Context, infraCli client.Client, selector client.MatchingLabelsSelector, cluster *cloudv1.Cluster) (asgs []*v1beta1.AutoScalingGroup, err error) {
	var asgList v1beta1.AutoScalingGroupList
	if err = infraCli.List(ctx, &asgList, selector); err != nil {
		return nil, fmt.Errorf("failed to list AutoScalingGroups: %w", err)
	}

	for i := range asgList.Items {
		asg := &asgList.Items[i]
		var specInstanceType string

		// Match the current AutoScalingGroup's component with the cluster's instance types in spec.
		switch asg.Labels[namer.ComponentSelectorKey] {
		case string(cloudv1.ComponentPD):
			specInstanceType = cluster.Spec.Cluster.PD.InstanceType
		case string(cloudv1.ComponentTiKV):
			specInstanceType = cluster.Spec.Cluster.TiKV.InstanceType
		case string(cloudv1.ComponentTiFlash):
			specInstanceType = cluster.Spec.Cluster.TiFlash.InstanceType
		case string(cloudv1.ComponentTiDB):
			specInstanceType = cluster.Spec.Cluster.TiDB.InstanceType
		case string(cloudv1.ComponentTiProxy):
			specInstanceType = cluster.Spec.Cluster.TiProxy.InstanceType
		default:
			continue
		}

		// Get the current AutoScalingGroup's instance type.
		if it := GetInstanceType(ctx, infraCli, asg); it != nil && specInstanceType == *it {
			asgs = append(asgs, asg)
		}
	}

	return asgs, nil
}

// TODO(ideascf): To refactor GetTiDBClusterSelfManagedNodeGroups to use this function.
// GetTiDBClusterSelfManagedNodeGroupsV2 returns the cluster's self-managed nodegroups that are currently in use.
func GetTiDBClusterSelfManagedNodeGroupsV2(ctx context.Context, infraCli client.Client, selector client.MatchingLabelsSelector, components []cloudv1.ClusterComponent) (asgs []*v1beta1.AutoScalingGroup, err error) {
	logger := log.FromContext(ctx)
	var asgList v1beta1.AutoScalingGroupList
	if err = infraCli.List(ctx, &asgList, selector); err != nil {
		return nil, fmt.Errorf("failed to list AutoScalingGroups: %w", err)
	}

	tp2component := make(map[cloudv1.ComponentType]cloudv1.ClusterComponent)
	for _, c := range components {
		tp2component[c.GetComponentType()] = c
	}

	for i := range asgList.Items {
		asg := &asgList.Items[i]

		componentType := cloudv1.ComponentType(asg.Labels[namer.ComponentSelectorKey])
		component, ok := tp2component[componentType]
		if !ok {
			logger.Info("component not found", "component_type", string(componentType), "selector", selector)
			continue
		}
		specInstanceType := component.GetInstanceType()

		// Get the current AutoScalingGroup's instance type.
		if it := GetInstanceType(ctx, infraCli, asg); it != nil && specInstanceType == *it {
			asgs = append(asgs, asg)
		}
	}

	return asgs, nil
}