package extratidb

import (
	"context"
	"fmt"
	"sort"

	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/crossplane-contrib/provider-aws/apis/autoscaling/v1beta1"
	ec2v1alpha1 "github.com/crossplane-contrib/provider-aws/apis/ec2/v1alpha1"
	"github.com/crossplane-contrib/provider-aws/apis/eks/manualv1alpha1"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"

	eksv1alpha1 "github.com/tidbcloud/eks-provider/api/v1alpha1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/pkg/nodegroup"
	"github.com/tidbcloud/infra-provider/pkg/nodegroup/selfmanaged"
	"github.com/tidbcloud/infra-provider/pkg/util"
	"github.com/tidbcloud/infra-provider/pkg/util/crossplane"
)

type AWSHandler struct {
	infraCli client.Client
	awsCli   aws.Client
}

var _ Handler = &AWSHandler{}

func NewAWSHandler(infraCli client.Client, awsCli aws.Client) *AWSHandler {
	return &AWSHandler{
		infraCli: infraCli,
		awsCli:   awsCli,
	}
}

func (h *AWSHandler) PrepareNodes(ctx context.Context, eksCli client.Client, provider gitops.Provider, network *cloudv1.Network, cluster *cloudv1.Cluster) (ngName string, err error) {
	var list manualv1alpha1.NodeGroupList
	if err = h.infraCli.List(ctx, &list, client.MatchingLabelsSelector{
		Selector: GetNodeCRLabelSelector(cluster.Name),
	}); err != nil && !apierrors.IsNotFound(err) {
		return "", fmt.Errorf("failed to list nodegroups: %w", err)
	}

	if len(list.Items) > 0 {
		// The extra tidb node group already exists, do nothing.
		return list.Items[0].Name, nil
	}

	var (
		refLTName string
		refLT     ec2v1alpha1.LaunchTemplate
		refNG     *manualv1alpha1.NodeGroup
	)

	if util.GetAnnotation(cluster, consts.UseSelfManagedNodegroupAnnotation) == "true" {
		// If the cluster is using self-managed nodegroups, we only get the launch template name.
		names, err := h.getCurrentSelfManagedNodeGroupNames(ctx, cluster)
		if err != nil {
			return "", fmt.Errorf("failed to get current self-managed nodegroups' names: %w", err)
		}
		refLTName = names[0]
	} else {
		// If the cluster is using managed nodegroups, we get the launch template name and the nodegroup object.
		currentNgs, err := h.getCurrentNodeGroups(ctx, cluster)
		if err != nil {
			return "", fmt.Errorf("failed to get current nodegroups: %w", err)
		}
		refNG = currentNgs[0].DeepCopy()
		refLTName = refNG.Name
	}

	// Get the launch template object, will create the extra tidb's based on it.
	if err := h.infraCli.Get(ctx, client.ObjectKey{Name: refLTName}, &refLT); err != nil {
		return "", fmt.Errorf("failed to get launch template: %w", err)
	}

	// According to the current nodegroups' zones, get the extra tidb's zone.
	eksNetworkSpec, err := eksv1alpha1.EKSNetworkSpecForNetwork(network)
	if err != nil {
		return "", fmt.Errorf("failed to get eks network spec: %w", err)
	}
	zones, err := provider.FindRegionZones(network.Spec.Region)
	if err != nil {
		return "", fmt.Errorf("failed to find region zones: %w", err)
	}
	extraTiDBZone, err := GetExtraTiDBZone(ctx, eksCli, cluster.Name, zones)
	if err != nil {
		return "", fmt.Errorf("failed to get extra tidb's zone: %w", err)
	}

	// Get the extra tidb's subnets according to the zone.
	subnetOut, err := h.awsCli.EC2(aws.GetRoleFromEKSNetworkSpec(eksNetworkSpec), aws.WithRegion(network.Spec.Region)).DescribeSubnets(ctx, &ec2.DescribeSubnetsInput{
		SubnetIds: eksNetworkSpec.SubnetIds,
	})
	if err != nil {
		return "", fmt.Errorf("failed to describe subnets: %w", err)
	}
	extraTiDBSubnets := aws.FindSubnetIdsByZones(subnetOut.Subnets, []string{extraTiDBZone})

	return h.createExtraTiDBNodeGroup(ctx, provider, network, eksNetworkSpec, cluster, extraTiDBSubnets, extraTiDBZone, &refLT, refNG)
}

func (h *AWSHandler) PrepareNodesV2(ctx context.Context, eksCli client.Client, provider gitops.Provider, envProvider dataplane.EnvProvider, cluster *cloudv1.Cluster) (ngName string, err error) {
	var list manualv1alpha1.NodeGroupList
	if err = h.infraCli.List(ctx, &list, client.MatchingLabelsSelector{
		Selector: GetNodeCRLabelSelector(cluster.Name),
	}); err != nil && !apierrors.IsNotFound(err) {
		return "", fmt.Errorf("failed to list nodegroups: %w", err)
	}

	if len(list.Items) > 0 {
		// The extra tidb node group already exists, do nothing.
		return list.Items[0].Name, nil
	}

	var (
		refLTName string
		refLT     ec2v1alpha1.LaunchTemplate
	)

	if util.GetAnnotation(cluster, consts.UseSelfManagedNodegroupAnnotation) == "true" {
		// If the cluster is using self-managed nodegroups, we only get the launch template name.
		names, err := h.getCurrentSelfManagedNodeGroupNames(ctx, cluster)
		if err != nil {
			return "", fmt.Errorf("failed to get current self-managed nodegroups' names: %w", err)
		}
		refLTName = names[0]
	} else {
		return "", fmt.Errorf("managed nodegroups are not supported now")
	}

	// Get the launch template object, will create the extra tidb's based on it.
	if err := h.infraCli.Get(ctx, client.ObjectKey{Name: refLTName}, &refLT); err != nil {
		return "", fmt.Errorf("failed to get launch template: %w", err)
	}

	// According to the current nodegroups' zones, get the extra tidb's zone.
	if err != nil {
		return "", fmt.Errorf("failed to get eks network spec: %w", err)
	}
	zones, err := h.getClusterZones(ctx, provider, envProvider, cluster)
	if err != nil {
		return "", fmt.Errorf("failed to find region zones: %w", err)
	}
	extraTiDBZone, err := GetExtraTiDBZone(ctx, eksCli, cluster.Name, zones)
	if err != nil {
		return "", fmt.Errorf("failed to get extra tidb's zone: %w", err)
	}

	extraTiDBSubnets := envProvider.AWS().GetNodeSubnets(extraTiDBZone)
	return h.createExtraTiDBNodeGroupV2(ctx, provider, envProvider, cluster, extraTiDBSubnets, extraTiDBZone, refLT.DeepCopy())
}

func (h *AWSHandler) AreAllNodesReady(ctx context.Context, _ client.Client, cluster *cloudv1.Cluster) (ready bool, err error) {
	var list manualv1alpha1.NodeGroupList
	if err = h.infraCli.List(ctx, &list, client.MatchingLabelsSelector{
		Selector: GetNodeCRLabelSelector(cluster.Name),
	}); err != nil && !apierrors.IsNotFound(err) {
		return false, fmt.Errorf("failed to list nodegroups: %w", err)
	}

	if len(list.Items) == 0 {
		return false, nil
	}

	for _, ng := range list.Items {
		if !isNodeGroupReady(ng) {
			return false, nil
		}
	}

	return true, nil
}

func (h *AWSHandler) DeleteNodes(ctx context.Context, eksCli client.Client, cluster *cloudv1.Cluster) error {
	// Get all tidb pods
	var podList corev1.PodList
	if err := eksCli.List(ctx, &podList,
		client.InNamespace(cluster.PrefixedId()),
		client.MatchingLabelsSelector{
			Selector: labels.SelectorFromSet(map[string]string{
				consts.ComponentLabelKey: "tidb",
			}),
		},
	); err != nil && !apierrors.IsNotFound(err) {
		return fmt.Errorf("failed to list pods: %w", err)
	}

	totalTiDBReplicas := cluster.Spec.Cluster.TiDB.Replicas
	for _, eg := range cluster.Spec.ExtraGroups {
		totalTiDBReplicas += eg.TiDB.Replicas
	}
	if totalTiDBReplicas == len(podList.Items) {
		// The extra tidb doesn't exist, can delete its launch template and nodegroup.
		if err := h.infraCli.DeleteAllOf(ctx, &manualv1alpha1.NodeGroup{}, client.MatchingLabelsSelector{
			Selector: labels.SelectorFromSet(map[string]string{
				"tidbcloud-cluster": cluster.Name,
				"tidb-component":    "extra-tidb",
			}),
		}); err != nil {
			return fmt.Errorf("failed to delete extra tidb's nodegroup: %w", err)
		}

		if err := h.infraCli.DeleteAllOf(ctx, &ec2v1alpha1.LaunchTemplate{}, client.MatchingLabelsSelector{
			Selector: labels.SelectorFromSet(map[string]string{
				"tidbcloud-cluster": cluster.Name,
				"tidb-component":    "extra-tidb",
			}),
		}); err != nil {
			return fmt.Errorf("failed to delete extra tidb's launch template: %w", err)
		}
	} else if len(podList.Items) > totalTiDBReplicas {
		return fmt.Errorf("the extra tidb pod still exists, scale in it first")
	}

	return nil
}

func (h *AWSHandler) getCurrentNodeGroups(ctx context.Context, cluster *cloudv1.Cluster) (ngs []*manualv1alpha1.NodeGroup, err error) {
	n := namer.NewNamer().
		WithClusterName(cluster.Name).
		WithTiDBNamespace(cluster.PrefixedId())
	list := &manualv1alpha1.NodeGroupList{}
	if err = h.infraCli.List(ctx, list, n.ClusterSelector()); err != nil {
		return nil, err
	}

	for i := range list.Items {
		var specInstanceType string
		com := list.Items[i].Labels[namer.ComponentSelectorKey]
		switch com {
		case string(cloudv1.ComponentTiDB):
			specInstanceType = cluster.Spec.Cluster.TiDB.InstanceType
		default:
			continue
		}

		for _, instanceType := range list.Items[i].Spec.ForProvider.InstanceTypes {
			// only include nodegroups matching the instance type for the current spec
			if instanceType == specInstanceType {
				ngs = append(ngs, &list.Items[i])
				break
			}
		}
	}

	return ngs, nil
}

func (h *AWSHandler) getCurrentSelfManagedNodeGroupNames(ctx context.Context, cluster *cloudv1.Cluster) (names []string, err error) {
	list := &v1beta1.AutoScalingGroupList{}
	n := namer.NewNamer().
		WithClusterName(cluster.Name).
		WithTiDBNamespace(cluster.PrefixedId())
	if err = h.infraCli.List(ctx, list, n.ClusterSelector()); err != nil {
		return nil, err
	}

	for i := range list.Items {
		ng := &list.Items[i]
		if ng.Labels[namer.ComponentSelectorKey] == string(cloudv1.ComponentTiDB) {
			if it := selfmanaged.GetInstanceType(ctx, h.infraCli, ng); it != nil && *it == cluster.Spec.Cluster.TiDB.InstanceType {
				names = append(names, ng.Name)
				break
			}
		}
	}

	if len(names) == 0 {
		return nil, fmt.Errorf("no nodegroups found")
	}

	return names, nil
}

// createExtraTiDBNodeGroup creates a launchtemplate and a managed nodegroup for the extra tidb.
func (h *AWSHandler) createExtraTiDBNodeGroup(
	ctx context.Context,
	provider gitops.Provider,
	network *cloudv1.Network,
	networkSpec *eksv1alpha1.EKSNetworkSpec,
	cluster *cloudv1.Cluster,
	subnets []string,
	zone string,
	refLT *ec2v1alpha1.LaunchTemplate,
	refNG *manualv1alpha1.NodeGroup,
) (ngName string, err error) {
	ngName = extraTiDBNodeGroupName(cluster.Name, zone)
	ngLabels := GetNodeCRLabels(cluster.Name, zone)
	objMeta := metav1.ObjectMeta{
		Name:   ngName,
		Labels: ngLabels,
	}
	providerConfigName := namer.UniqueIdentifier(network.ObjectMeta)

	lt := &ec2v1alpha1.LaunchTemplate{
		ObjectMeta: objMeta,
		Spec: ec2v1alpha1.LaunchTemplateSpec{
			ForProvider: ec2v1alpha1.LaunchTemplateParameters{
				Region:             network.Spec.Region,
				LaunchTemplateName: pointer.String(ngName),
				LaunchTemplateData: &ec2v1alpha1.RequestLaunchTemplateData{
					BlockDeviceMappings: refLT.Spec.ForProvider.LaunchTemplateData.BlockDeviceMappings,
					SecurityGroupIDs:    refLT.Spec.ForProvider.LaunchTemplateData.SecurityGroupIDs,
				},
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: providerConfigName},
			},
		},
	}
	if err = k8s.Apply(ctx, h.infraCli, lt); err != nil {
		return "", fmt.Errorf("cannot create launch template for extra tidb nodegroup: %w", err)
	}

	var ngSpec manualv1alpha1.NodeGroupSpec
	if refNG != nil {
		// If the cluster uses managed nodegroups, we can set the ngSpec based on them.
		ngSpec = refNG.Spec
	} else {
		ngSpec = manualv1alpha1.NodeGroupSpec{
			ForProvider: manualv1alpha1.NodeGroupParameters{
				AMIType:     getAMITypeByInstanceType(provider, cluster.Spec.Cluster.TiDB.InstanceType),
				ClusterName: namer.LegacyBuildClusterNameByNetwork(network),
				Region:      network.Spec.Region,
				Taints: []manualv1alpha1.Taint{
					{
						Effect: "NO_SCHEDULE",
						Key:    pointer.String("dedicated"),
						Value:  pointer.String("tidb-tidb"),
					},
				},
				InstanceTypes: []string{cluster.Spec.Cluster.TiDB.InstanceType},
				NodeRole:      networkSpec.NodeRoleArn,
				LaunchTemplate: &manualv1alpha1.LaunchTemplateSpecification{
					NameRef: &xpv1.Reference{
						Name: ngName,
					},
				},
				Tags: map[string]string{
					consts.TagKeyCrossplaneVendor: consts.TagValCrossplaneVendor,
				},
				CapacityType: pointer.String("ON_DEMAND"),
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: providerConfigName},
			},
		}
	}

	// Make the nodegroup's version align with the eks controlplane's version
	ngSpec.ForProvider.Version = nil
	ngSpec.ForProvider.ReleaseVersion = nil
	ngSpec.ForProvider.Subnets = subnets
	// The extra tidb node group should have only one instance.
	ngSpec.ForProvider.ScalingConfig = &manualv1alpha1.NodeGroupScalingConfig{
		MinSize: pointer.Int32(1),
		MaxSize: pointer.Int32(1),
	}
	ngSpec.ForProvider.LaunchTemplate = &manualv1alpha1.LaunchTemplateSpecification{
		NameRef: &xpv1.Reference{
			Name: lt.Name,
		},
	}
	// The labels need to be the same as the original's, which are used as the selector by tidb pods.
	ngSpec.ForProvider.Labels = namer.NewNamer().
		LegacyWithNetworkName(network.Name).
		LegacyWithNetworkNamespace(network.Namespace).
		WithClusterName(cluster.Name).
		WithTiDBNamespace(cluster.PrefixedId()).
		WithComponent(string(cloudv1.ComponentTiDB)).Labels()

	// We have a feature gate `split_csi_daemonset`, which will split the ebs csi daemonset into two daemonsets.
	// To ensure the extra tidb node has a CSI DaemonSet pod, we need to add this label.
	ngSpec.ForProvider.Labels["usedby"] = "customer-tidb"

	extraTiDBNg := &manualv1alpha1.NodeGroup{
		ObjectMeta: objMeta,
		Spec:       ngSpec,
	}
	if err = k8s.Apply(ctx, h.infraCli, extraTiDBNg); err != nil {
		return "", fmt.Errorf("failed to create extra tidb nodegroup: %w", err)
	}

	return ngName, nil
}

func (h *AWSHandler) createExtraTiDBNodeGroupV2(
	ctx context.Context,
	provider gitops.Provider,
	envProvider dataplane.EnvProvider,
	cluster *cloudv1.Cluster,
	subnets []string,
	zone string,
	refLT *ec2v1alpha1.LaunchTemplate,
) (ngName string, err error) {
	ngName = extraTiDBNodeGroupName(cluster.Name, zone)
	ngLabels := GetNodeCRLabels(cluster.Name, zone)
	objMeta := metav1.ObjectMeta{
		Name:   ngName,
		Labels: ngLabels,
	}

	lt := &ec2v1alpha1.LaunchTemplate{
		ObjectMeta: objMeta,
		Spec: ec2v1alpha1.LaunchTemplateSpec{
			ForProvider: ec2v1alpha1.LaunchTemplateParameters{
				Region:             envProvider.GetRegion(),
				LaunchTemplateName: pointer.String(ngName),
				LaunchTemplateData: &ec2v1alpha1.RequestLaunchTemplateData{
					BlockDeviceMappings: refLT.Spec.ForProvider.LaunchTemplateData.BlockDeviceMappings,
					SecurityGroupIDs:    refLT.Spec.ForProvider.LaunchTemplateData.SecurityGroupIDs,
				},
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: refLT.Spec.ProviderConfigReference.Name},
			},
		},
	}
	if err = k8s.Apply(ctx, h.infraCli, lt); err != nil {
		return "", fmt.Errorf("cannot create launch template for extra tidb nodegroup: %w", err)
	}

	var nodeRole string
	if cluster.Status.NodeIdentity != nil && cluster.Status.NodeIdentity.AWS != nil {
		nodeRole = cluster.Status.NodeIdentity.AWS.RoleARN
	} else {
		return "", fmt.Errorf("failed to get node role from launch template")
	}

	ngSpec := manualv1alpha1.NodeGroupSpec{
		ForProvider: manualv1alpha1.NodeGroupParameters{
			AMIType:     getAMITypeByInstanceType(provider, cluster.Spec.Cluster.TiDB.InstanceType),
			ClusterName: envProvider.GetClusterName(),
			Region:      envProvider.GetRegion(),
			Taints: []manualv1alpha1.Taint{
				{
					Effect: "NO_SCHEDULE",
					Key:    pointer.String("dedicated"),
					Value:  pointer.String("tidb-tidb"),
				},
			},
			InstanceTypes: []string{cluster.Spec.Cluster.TiDB.InstanceType},
			NodeRole:      nodeRole,
			LaunchTemplate: &manualv1alpha1.LaunchTemplateSpecification{
				NameRef: &xpv1.Reference{
					Name: ngName,
				},
			},
			Tags: map[string]string{
				consts.TagKeyCrossplaneVendor: consts.TagValCrossplaneVendor,
			},
			CapacityType: pointer.String("ON_DEMAND"),
		},
		ResourceSpec: xpv1.ResourceSpec{
			ProviderConfigReference: &xpv1.Reference{Name: refLT.Spec.ProviderConfigReference.Name},
		},
	}

	// Make the nodegroup's version align with the eks controlplane's version
	ngSpec.ForProvider.Version = nil
	ngSpec.ForProvider.ReleaseVersion = nil
	ngSpec.ForProvider.Subnets = subnets
	// The extra tidb node group should have only one instance.
	ngSpec.ForProvider.ScalingConfig = &manualv1alpha1.NodeGroupScalingConfig{
		MinSize: pointer.Int32(1),
		MaxSize: pointer.Int32(1),
	}
	ngSpec.ForProvider.LaunchTemplate = &manualv1alpha1.LaunchTemplateSpecification{
		NameRef: &xpv1.Reference{
			Name: lt.Name,
		},
	}
	// The labels need to be the same as the original's, which are used as the selector by tidb pods.
	ngSpec.ForProvider.Labels = namer.NewNamer().
		WithClusterName(cluster.Name).
		WithTiDBNamespace(cluster.PrefixedId()).
		WithComponent(string(cloudv1.ComponentTiDB)).Labels()
	if cluster.Spec.NetworkName != "" {
		ngSpec.ForProvider.Labels[namer.LegacyNetworkSelectorKey] = cluster.Spec.NetworkName
		ngSpec.ForProvider.Labels[namer.LegacyNetworkNamespaceSelectorKey] = cluster.Namespace
	}

	// We have a feature gate `split_csi_daemonset`, which will split the ebs csi daemonset into two daemonsets.
	// To ensure the extra tidb node has a CSI DaemonSet pod, we need to add this label.
	ngSpec.ForProvider.Labels["usedby"] = "customer-tidb"

	extraTiDBNg := &manualv1alpha1.NodeGroup{
		ObjectMeta: objMeta,
		Spec:       ngSpec,
	}
	if err = k8s.Apply(ctx, h.infraCli, extraTiDBNg); err != nil {
		return "", fmt.Errorf("failed to create extra tidb nodegroup: %w", err)
	}

	return ngName, nil
}

// getClusterZones gets the cluster zones. If the cluster spec has the zones info, it will return the zones in spec;
// otherwise, it will call provider.FindRegionZones to get the zones in the region by config.
func (h *AWSHandler) getClusterZones(_ context.Context, provider gitops.Provider, envProvider dataplane.EnvProvider, cluster *cloudv1.Cluster) ([]string, error) {

	if cloudaws := cluster.Spec.Kubernetes.Cloud.AWS; cloudaws != nil {
		zones := []string{}
		for _, group := range cloudaws.SubnetGroups {
			if group.Name == cloudv1.AWSSubnetGroupNodes {
				for _, subnet := range group.Subnets {
					zones = append(zones, subnet.Zone)
				}
			}
		}
		sort.Strings(zones)
		return zones, nil
	}

	zones, err := provider.FindRegionZones(envProvider.GetRegion())
	if err != nil {
		return nil, fmt.Errorf("failed to find region zones: %w", err)
	}
	return zones, nil
}

func getAMITypeByInstanceType(provider gitops.Provider, instanceType string) *string {
	instanceTypes, err := provider.GetInstanceTypes()
	if err != nil {
		return nil
	}
	ami, found := nodegroup.GetAMITypeByInstanceType(instanceTypes, instanceType)
	if found {
		return pointer.String(ami)
	}
	return nil
}

func isNodeGroupReady(ng manualv1alpha1.NodeGroup) bool {
	// If the node group is being deleted, we consider it as ready.
	if !ng.DeletionTimestamp.IsZero() {
		return true
	}

	createdNodes := ng.Status.AtProvider.ScalingConfig.DesiredSize != nil && *ng.Status.AtProvider.ScalingConfig.DesiredSize > 0
	active := ng.Status.AtProvider.Status == manualv1alpha1.NodeGroupStatusActive
	readyAndSynced := crossplane.IsReadyAndSynced(ng.Status.Conditions)

	return createdNodes && active && readyAndSynced
}