package eks

import (
	"context"
	"fmt"

	"github.com/aws/aws-sdk-go-v2/service/ssm"
	"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"
	"github.com/hashicorp/go-multierror"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/util/retry"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/eks-provider/controllers/upgradeoperation/common"
	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/gitops"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/nodegroup/selfmanaged"
	"github.com/tidbcloud/infra-provider/pkg/util"
	crossplaneutil "github.com/tidbcloud/infra-provider/pkg/util/crossplane"
)

// getTiDBClusterSelfManagedNodeGroups returns the tidb cluster's self-managed node groups those are currently in use and match the desired version.
// It also returns a map from ASG name to extra group name, so callers can correctly distinguish
// between default and extra tidb groups when generating new ASG names.
func getTiDBClusterSelfManagedNodeGroups(ctx context.Context, infraCli client.Client, info *upgradeInfo, cluster *cloudv1.Cluster) (
	oldNGs []*v1beta1.AutoScalingGroup, newNGs []*v1beta1.AutoScalingGroup, asgExtraGroupMap map[string]string, err error,
) {
	return selfmanaged.ClassifyTiDBSelfManagedNGsForK8sVersionUpgrade(ctx, infraCli, cluster, info.desiredK8sVersion)
}

func createNewSelfManagedNodeGroups(ctx context.Context, infraCli client.Client, ssmClient *ssm.Client, provider gitops.Provider, info *upgradeInfo, oldSelfManagedNGs []*v1beta1.AutoScalingGroup, asgExtraGroupMap map[string]string) error {
	var errs *multierror.Error
	// amiIDs stores the ami id for each architecture.
	amiIDs := make(map[string]string, 2)

	// Create new version self-managed node groups based on the old version self-managed node groups.
	for i := range oldSelfManagedNGs {
		// Only create new self-managed node groups for those with desired capacity > 0.
		if oldSelfManagedNGs[i].Spec.ForProvider.DesiredCapacity != nil && *oldSelfManagedNGs[i].Spec.ForProvider.DesiredCapacity > 0 {
			it, err := getInstanceTypeFromSelfManagedNodeGroup(oldSelfManagedNGs[i], provider)
			if err != nil {
				return fmt.Errorf("failed to find instance type for asg %s: %w", oldSelfManagedNGs[i].Name, err)
			}

			amiID, ok := amiIDs[it.Arch]
			if !ok {
				amiIDPtr, err := aws.GetAMIID(ctx, ssmClient, selfmanaged.GetSSMName(info.desiredK8sVersion, it.Arch))
				if err != nil {
					return fmt.Errorf("unable to get the new ami id: %w", err)
				}
				amiID = *amiIDPtr
				amiIDs[it.Arch] = amiID
			}

			extraGroupName := asgExtraGroupMap[oldSelfManagedNGs[i].Name]
			errs = multierror.Append(errs, fromSelfManagedToSelfManaged(ctx, infraCli, info, oldSelfManagedNGs[i], amiID, extraGroupName))
		}
	}

	return errs.ErrorOrNil()
}

func getInstanceTypeFromSelfManagedNodeGroup(asg *v1beta1.AutoScalingGroup, provider gitops.Provider) (cloudv1.InstanceType, error) {
	if asg == nil {
		return cloudv1.InstanceType{}, fmt.Errorf("the self-managed node group is nil")
	}
	if len(asg.Status.AtProvider.Instances) == 0 {
		return cloudv1.InstanceType{}, fmt.Errorf("the self-managed node group %s has no instances", asg.Name)
	}

	return provider.FindInstanceType(*asg.Status.AtProvider.Instances[0].InstanceType)
}

func fromSelfManagedToSelfManaged(ctx context.Context, infraCli client.Client, info *upgradeInfo, ng *v1beta1.AutoScalingGroup, amiID string, extraGroupName string) error {
	var oldLT ec2v1alpha1.LaunchTemplate
	if err := infraCli.Get(ctx, types.NamespacedName{Name: ng.Name}, &oldLT); err != nil {
		return fmt.Errorf("unable to get the old launch template: %w", err)
	}

	instanceType := oldLT.Spec.ForProvider.LaunchTemplateData.InstanceType
	if instanceType == nil {
		return fmt.Errorf("instance type not found for launch template %s", oldLT.Name)
	}

	asgName := selfmanaged.NewVersionTiDBSelfManagedASGName(info.cluster, ng, *instanceType, info.desiredK8sVersion, info.envProvider.GetRegion(), extraGroupName)

	newLT, newNG := selfmanaged.FromSelfManagedToSelfManaged(&oldLT, ng, info.envProvider.GetClusterName(), asgName, amiID, info.desiredK8sVersion)
	if err := k8s.Apply(ctx, infraCli, newLT); err != nil {
		return fmt.Errorf("unable to create new launch template: %w", err)
	}

	return k8s.Apply(ctx, infraCli, newNG)
}

// countReadyNodeGroups returns how many node groups are upgraded to desiredVersion and are Ready+Synced.
func countReadyNodeGroups(asgs []*v1beta1.AutoScalingGroup, desiredVersion string) (ready, total int) {
	total = len(asgs)
	for _, asg := range asgs {
		if util.GetAnnotation(asg, selfmanaged.K8sVersionAnnotationKey) == desiredVersion && crossplaneutil.IsReadyAndSynced(asg.Status.Conditions) {
			ready++
		}
	}
	return ready, total
}

// allTiDBNodesUpgraded checks if all TiDB nodes are upgraded to the desired version.
func allTiDBNodesUpgraded(info *upgradeInfo) bool {
	if info.cluster.Status.Phase == cloudv1.ClusterPhaseSuspend {
		status, err := cloudv1.EKSClusterStatusFromCluster(info.cluster)
		return err == nil && status.NodeVersion == info.desiredK8sVersion
	} else {
		// If the cluster is not suspended, compare control plane and node versions to determine if nodes are upgraded.
		return info.currentNodeVersion != "" && info.desiredK8sVersion == info.currentNodeVersion
	}
}

func deleteNodeGroupAndLaunchTemplate(ctx context.Context, infraCli client.Client, ng client.Object, ngName string) error {
	var errs *multierror.Error

	if err := k8s.DeleteObject(ctx, infraCli, ng); err != nil {
		errs = multierror.Append(errs, fmt.Errorf("failed to delete the old nodegroup %s: %w", ngName, err))
	}

	var oldLaunchTemplate ec2v1alpha1.LaunchTemplate
	if err := infraCli.Get(ctx, client.ObjectKey{Name: ngName}, &oldLaunchTemplate); err == nil {
		if e := k8s.DeleteObject(ctx, infraCli, &oldLaunchTemplate); e != nil {
			errs = multierror.Append(errs, fmt.Errorf("failed to delete the old launch template %s: %w", oldLaunchTemplate.Name, e))
		}
	}

	return errs.ErrorOrNil()
}

// updateNodeGroupVersion updates the node group CR's version and release version.
// Since the release version may be modified back, we need to check it.
func updateNodeGroupVersion(ctx context.Context, infraCli client.Client, ng *manualv1alpha1.NodeGroup, desiredVersion string) error {
	ng.Spec.ForProvider.Version = pointer.String(desiredVersion)
	ng.Spec.ForProvider.ReleaseVersion = nil
	return infraCli.Update(ctx, ng)
}

func cordonNodes(ctx context.Context, cli client.Client, nodes corev1.NodeList) error {
	var errs *multierror.Error

	for _, node := range nodes.Items {
		if !node.Spec.Unschedulable {
			if err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
				var currentNode corev1.Node
				if err := cli.Get(ctx, types.NamespacedName{Name: node.Name}, &currentNode); err != nil {
					return err
				}
				currentNode.Spec.Unschedulable = true
				return cli.Update(ctx, &currentNode)
			}); err != nil {
				errs = multierror.Append(errs, err)
			}
		}
	}

	return errs.ErrorOrNil()
}

// cordonTiDBNodesByNodeGroups cordon tidb nodes that belong to the given self-managed node groups.
// So that after restarting pods, the pods will be scheduled to the new version nodes.
func cordonTiDBNodesByNodeGroups(ctx context.Context, cli client.Client, oldASGs []*v1beta1.AutoScalingGroup) error {
	var errs *multierror.Error

	for _, asg := range oldASGs {
		var nodeList corev1.NodeList
		if err := cli.List(ctx, &nodeList, client.MatchingLabelsSelector{
			Selector: labels.Set{
				// Each Auto Scaling Group (ASG) name is composed using the following pattern:
				//   {component}-{machineType}-{ns-cluster}-{hash}-{zones}-{version}
				// Since the ASG name already includes the cluster information (ns-cluster),
				// we can safely use "Name: asg.Name" here to match and filter nodes that belong
				// to this specific ASG, instead of relying on a separate cluster selector label.
				//
				// Example:
				//   tikv-m6g-2xlarge-nightly-ms-3200472-10-79d44d7b-us-west-2b-1-30
				// From this name, we can identify that the node belongs to cluster "tidb10852843712412958566".
				// namer.ClusterSelectorKey: clusterID,
				"Name": asg.Name,
			}.AsSelector(),
		}); err != nil {
			return err
		}

		if err := cordonNodes(ctx, cli, nodeList); err != nil {
			errs = multierror.Append(errs, err)
		}
	}

	return errs.ErrorOrNil()
}

// countTiDBPodsOnNewNodeGroups returns how many active tidb component pods are on new node groups.
func countTiDBPodsOnNewNodeGroups(ctx context.Context, cli client.Client, namespace string, newASGs []*v1beta1.AutoScalingGroup, extraTiDBNodegroupName string) (newPods, totalPods int, ngNamePodMap map[string][]string, err error) {
	logger := log.FromContext(ctx)

	// Build a set of new ASG names for O(1) lookup.
	newASGSet := map[string]struct{}{}
	for _, asg := range newASGs {
		if asg.Spec.ForProvider.DesiredCapacity != nil && *asg.Spec.ForProvider.DesiredCapacity > 0 {
			newASGSet[asg.Name] = struct{}{}
		}
	}

	// List all pods that are managed by tidb-operator and in the given namespace.
	var podList corev1.PodList
	if err = cli.List(ctx, &podList,
		client.InNamespace(namespace),
		client.MatchingLabelsSelector{
			Selector: labels.SelectorFromSet(map[string]string{
				consts.ManagedByLabelKey: "tidb-operator",
			}),
		},
	); err != nil || len(podList.Items) == 0 {
		return 0, 0, nil, fmt.Errorf("unable to list the tidb pods: %v; the number of pods: %d", err, len(podList.Items))
	}

	var errs *multierror.Error
	ngNamePodMap = map[string][]string{}
	for _, pod := range podList.Items {
		_, exists := common.ComponentSet[pod.Labels[consts.ComponentLabelKey]]
		if !exists || !pod.DeletionTimestamp.IsZero() {
			continue
		}
		totalPods++

		if pod.Spec.NodeName == "" {
			logger.Info("component pod has not been scheduled", "pod", pod.Name)
			continue
		}

		var node corev1.Node
		if err = cli.Get(ctx, types.NamespacedName{Name: pod.Spec.NodeName}, &node); err != nil {
			logger.Error(err, "unable to get the node for pod", "node", pod.Spec.NodeName, "pod", pod.Name, "namespace", pod.Namespace)
			errs = multierror.Append(errs, err)
			continue
		}

		ngName := node.Labels["Name"]
		if ngName == "" {
			ngName = node.Labels["eks.amazonaws.com/nodegroup"]
		}
		if ngName == "" {
			logger.Error(nil, "component node not found nodegroup name", "pod", pod.Name, "node", pod.Spec.NodeName)
			continue
		}
		ngNamePodMap[ngName] = append(ngNamePodMap[ngName], pod.Name)

		// A pod is correctly placed if it is on a new self-managed ASG OR on the extra tidb NodeGroup
		// (whose nodegroup name is extraTiDBNodegroupName).
		_, isOnNewASG := newASGSet[ngName]
		isOnExtraTiDBNodeGroup := ngName == extraTiDBNodegroupName
		if !isOnNewASG && !isOnExtraTiDBNodeGroup {
			logger.Info("tidb pod is still on old nodegroup", "pod", pod.Name, "nodegroup", ngName)
			continue
		}
		newPods++
	}
	if err = errs.ErrorOrNil(); err != nil {
		return 0, 0, nil, err
	}
	if totalPods == 0 {
		logger.Info("no active tidb component pods found when checking nodegroup placement", "namespace", namespace)
	}

	return newPods, totalPods, ngNamePodMap, nil
}

func countComponentPodsOnNewNodeGroups(ctx context.Context, cli client.Client, namespace, component string, newASGs []*v1beta1.AutoScalingGroup, tc *opv1alpha1.TidbCluster) (onNew, total int, err error) {
	logger := log.FromContext(ctx)

	targetNGs := map[string]struct{}{}
	for _, asg := range newASGs {
		if asg.Spec.ForProvider.DesiredCapacity != nil && *asg.Spec.ForProvider.DesiredCapacity > 0 {
			targetNGs[asg.Name] = struct{}{}
		}
	}
	if len(targetNGs) == 0 {
		return 0, 0, fmt.Errorf("no active new node groups found")
	}

	var podList corev1.PodList
	if err = cli.List(ctx, &podList,
		client.InNamespace(namespace),
		client.MatchingLabelsSelector{
			Selector: labels.SelectorFromSet(map[string]string{
				consts.ManagedByLabelKey: "tidb-operator",
				consts.ComponentLabelKey: component,
				consts.InstanceLabelKey:  tc.Name,
			}),
		},
	); err != nil {
		return 0, 0, fmt.Errorf("unable to list %s pods: %w", component, err)
	}

	for _, pod := range podList.Items {
		if pod.DeletionTimestamp != nil {
			continue
		}
		total++

		if pod.Spec.NodeName == "" {
			logger.Info("component pod has not been scheduled", "component", component, "pod", pod.Name)
			continue
		}

		var node corev1.Node
		if err = cli.Get(ctx, types.NamespacedName{Name: pod.Spec.NodeName}, &node); err != nil {
			return onNew, total, fmt.Errorf("unable to get node %s for pod %s: %w", pod.Spec.NodeName, pod.Name, err)
		}

		ngName := node.Labels["Name"]
		if ngName == "" {
			ngName = node.Labels["eks.amazonaws.com/nodegroup"]
		}
		if _, ok := targetNGs[ngName]; !ok {
			logger.Info("component pod is still on old nodegroup", "component", component, "pod", pod.Name, "nodegroup", ngName)
			continue
		}
		onNew++
	}

	if total == 0 {
		logger.Info("no active component pods found", "component", component)
	}
	return onNew, total, nil
}