package gke

import (
	"context"
	"fmt"
	"sync"
	"time"

	gkev1beta1 "github.com/crossplane-contrib/provider-gcp/apis/container/v1beta1"
	gkev1beta2 "github.com/crossplane-contrib/provider-gcp/apis/container/v1beta2"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/pointer"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	"github.com/tidbcloud/eks-provider/api/v1alpha1"
	. "github.com/tidbcloud/eks-provider/controllers/k8sversion/common"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/applier"
	"github.com/tidbcloud/infra-provider/pkg/applier/mutatefuncs"
	"github.com/tidbcloud/infra-provider/pkg/cloud/gcp"
	"github.com/tidbcloud/infra-provider/pkg/config"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/pkg/nodepool"
)

type GKEUpgrader struct {
	cfg           *config.RootConfig
	infraCli      client.Client
	gcpClient     gcp.Client
	eventRecorder record.EventRecorder

	// key: gke cluster CR's UID (do not use the name, since the CR may be recreated with the same name)
	// value: gke client
	gkeClientMap sync.Map

	// Used to apply the node pool CRs.
	// Should align with the nodepool task in the cluster controller.
	npApplier applier.Applier[*gkev1beta1.NodePool, gkev1beta1.NodePool]
}

var _ Upgrader = &GKEUpgrader{}

func NewUpgrader(cfg *config.RootConfig, infraCli client.Client, gcpClient gcp.Client, eventRecorder record.EventRecorder) *GKEUpgrader {
	return &GKEUpgrader{
		infraCli:      infraCli,
		cfg:           cfg,
		gcpClient:     gcpClient,
		npApplier:     applier.NewApplier(infraCli, mutatefuncs.MutateGKENodePool),
		eventRecorder: eventRecorder,
	}
}

// UpgradeControlPlane upgrades the control plane of a GKE cluster.
// Since the crossplane does not support upgrading the control plane via CR, we need to upgrade the control plane via gcloud sdk.
// See: https://github.com/crossplane-contrib/provider-gcp/issues/479
func (u *GKEUpgrader) UpgradeControlPlane(ctx context.Context, uo *cloudv1.UpgradeOperation, network *cloudv1.Network) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	desiredVersion := uo.Status.CurrentStep.Version
	gke, err := getGKECR(ctx, u.infraCli, network)
	if err != nil {
		logger.Error(err, "unable to get the gke cluster CR")
		return ctrl.Result{}, err
	}
	logger.Info("get gke cluster CR", "gke", gke.Name, "current version", gke.Status.AtProvider.CurrentMasterVersion, "status", gke.Status.AtProvider.Status)

	switch gke.Status.AtProvider.Status {
	case gkev1beta2.ClusterStateReconciling:
		LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, "The GKE cluster is reconciling, wait for the next reconcile", nil)
		return ctrl.Result{RequeueAfter: 1 * time.Minute}, nil

	case gkev1beta2.ClusterStateRunning:
		currentCPVersion := formatGKEVersion(gke.Status.AtProvider.CurrentMasterVersion)

		if currentCPVersion == desiredVersion {
			LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, "The GKE cluster's control plane is already upgraded", nil)
			if err := v1alpha1.UpdateNetworkStatus(ctx, u.infraCli, network, func(status *v1alpha1.CommonNetworkStatus) error {
				status.K8sClusterVersion = desiredVersion
				return nil
			}); err != nil {
				logger.Error(err, "unable to update network status")
				return ctrl.Result{}, err
			}

			return ctrl.Result{}, FinishUpgrading(ctx, u.infraCli, uo)
		} else if ExpectToBeUpgraded(cloudv1.ControlPlane, desiredVersion, currentCPVersion) {
			gkeNetworkSpec, err := v1alpha1.GKENetworkSpecForNetwork(network)
			if err != nil {
				logger.Error(err, "unable to get gke network spec from network CR")
				return ctrl.Result{}, err
			}

			selectedVersion, err := u.gcpClient.UpgradeCluster(ctx, gkeNetworkSpec.ProjectID, network.Spec.Region, gke.Name, desiredVersion)
			if err != nil {
				logger.Error(err, "unable to upgrade gke cluster")
				return ctrl.Result{}, err
			}

			// It will take about 25 minutes to upgrade the control plane.
			LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, fmt.Sprintf("Start to upgrade the GKE cluster's control plane to %s", selectedVersion), nil)
			return ctrl.Result{RequeueAfter: 20 * time.Minute}, nil
		} else {
			LogAndEmitEvent(logger, u.eventRecorder, uo, WarningEvent, fmt.Sprintf("The GKE cluster is in unexpected version: %s", gke.Status.AtProvider.CurrentMasterVersion), nil)
		}

	default:
		LogAndEmitEvent(logger, u.eventRecorder, uo, WarningEvent, fmt.Sprintf("The GKE cluster is in unexpected status: %s", gke.Status.AtProvider.Status), nil)
	}

	return ctrl.Result{}, nil
}

func (u *GKEUpgrader) UpgradeAdminNodes(ctx context.Context, uo *cloudv1.UpgradeOperation, network *cloudv1.Network) (ctrl.Result, error) {
	var np gkev1beta1.NodePool
	adminNPName := namer.UniqueIdentifier(network.ObjectMeta, namer.WithPrefix(consts.TiDBCloudPrefix), namer.WithLimit(40))
	if err := u.infraCli.Get(ctx, client.ObjectKey{Name: adminNPName}, &np); err != nil {
		if apierrors.IsNotFound(err) {
			return ctrl.Result{}, fmt.Errorf("admin node pool %s not found", adminNPName)
		}
		return ctrl.Result{}, fmt.Errorf("failed to get admin node pool: %w", err)
	}

	return u.upgradeNodePools(ctx, uo, network, &gkev1beta1.NodePoolList{Items: []gkev1beta1.NodePool{np}}, cloudv1.AdminNodes)
}

func (u *GKEUpgrader) UpgradeCDCNodes(ctx context.Context, uo *cloudv1.UpgradeOperation, network *cloudv1.Network, cluster *cloudv1.Cluster) (ctrl.Result, error) {
	var list gkev1beta1.NodePoolList
	if err := u.infraCli.List(ctx, &list, client.MatchingLabelsSelector{
		Selector: labels.Set{
			"cluster":   cluster.Name,
			"component": "ticdc",
		}.AsSelector(),
	}); err != nil && !apierrors.IsNotFound(err) {
		return ctrl.Result{}, fmt.Errorf("failed to list ticdc node pools: %w", err)
	}

	return u.upgradeNodePools(ctx, uo, network, &list, cloudv1.CDCNodes)
}

func (u *GKEUpgrader) UpgradeDMNodes(ctx context.Context, uo *cloudv1.UpgradeOperation, network *cloudv1.Network, cluster *cloudv1.Cluster) (ctrl.Result, error) {
	var list gkev1beta1.NodePoolList
	if err := u.infraCli.List(ctx, &list, client.MatchingLabelsSelector{
		Selector: labels.Set{
			"cluster":   cluster.Name,
			"component": "tiflow",
		}.AsSelector(),
	}); err != nil && !apierrors.IsNotFound(err) {
		return ctrl.Result{}, fmt.Errorf("failed to list ticdc node pools: %w", err)
	}

	return u.upgradeNodePools(ctx, uo, network, &list, cloudv1.DMNodes)
}

func (u *GKEUpgrader) UpgradeTiDBNodes(ctx context.Context, uo *cloudv1.UpgradeOperation, network *cloudv1.Network, cluster *cloudv1.Cluster) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	// Fetch necessary information.
	desiredVersion := uo.Status.CurrentStep.Version
	gke, err := getGKECR(ctx, u.infraCli, network)
	if err != nil {
		logger.Error(err, "unable to get the gke cluster CR")
		return ctrl.Result{}, err
	}
	gkeCli, err := u.getGKEClient(ctx, network, gke)
	if err != nil {
		logger.Error(err, "failed to get gke client")
		return ctrl.Result{}, err
	}

	// Check if the control plane has been upgraded to the desired version.
	if !isControlPlaneUpgraded(gke, desiredVersion) {
		LogAndEmitEvent(logger, u.eventRecorder, uo, WarningEvent, ControlPlaneNotUpgraded, nil)
		return ctrl.Result{}, uo.Fail(ctx, u.infraCli, ControlPlaneNotUpgraded)
	}

	currentNodeVersion, err := GetTiDBNodeVersion(cluster, network.Spec.Provider)
	if err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to get TiDB node versions from the cluster CR: %w", err)
	}
	logger.Info("current node version", "version", currentNodeVersion)

	// Check if the tidb nodes have been upgraded.
	if currentNodeVersion == desiredVersion {
		logger.Info("the tidb nodes are already upgraded")

		if err = nodepool.DeleteOverProvisioningDeployments(ctx, gkeCli, consts.ClusterPrefix, cluster.PrefixedId()); err != nil {
			logger.Error(err, "failed to delete over provisioning deployments")
			return ctrl.Result{}, err
		}
		LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, "Delete the old version nodes", nil)
		return ctrl.Result{}, FinishUpgrading(ctx, u.infraCli, uo)
	}

	// If all tidb clusters under the network are suspended, we can upgrade all node pools directly.
	if allTiDBClustersSuspended(ctx, u.infraCli, network) {
		logger.Info("All tidb clusters in the gke are suspended, will upgrade all tidb node pools directly")

		allNPs, err := getAllNodePoolsUnderNetwork(ctx, u.infraCli, network)
		if err != nil {
			return ctrl.Result{}, fmt.Errorf("failed to get all node pools: %w", err)
		}
		return u.upgradeNodePools(ctx, uo, network, toNodePoolList(allNPs), cloudv1.TiDBNodes)
	}

	// If the cluster is suspended, and there are other non-suspended clusters in the gke,
	// we can directly update the node version in the cluster CR's status,
	// so that when it's resumed, it will be added node selector to use new version nodes.
	// This check MUST be after the `allTiDBClustersSuspended` check.
	if cluster.Status.Phase == cloudv1.ClusterPhaseSuspend {
		logger.Info("will only update the node version in cluster CR's status", "desired version", desiredVersion)
		return ctrl.Result{Requeue: true}, UpdateClusterStatus(ctx, u.infraCli, cluster, desiredVersion)
	}

	// Get the old and new version node pools.
	oldNPs, newNPs, needCreate, err := getTiDBNodePoolsByCluster(ctx, u.infraCli, cluster, desiredVersion)
	if err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to get tidb node pools: %w", err)
	}
	logger.Info("tidb node pools", "old", GetCRNameList(oldNPs), "new", GetCRNameList(newNPs), "needCreate", needCreate)

	// Pause the network and cluster CR, start to upgrade the tidb node pools.
	scaleOutExtraTiDB := uo.GetCurrentStepOption(ScaleOutExtraTiDBAnnotation) != "false"
	logger.Info("start to upgrade the tidb node pools", "scaleOutExtraTiDB", scaleOutExtraTiDB)

	// Create new version node pools
	if needCreate {
		if err = createNewVersionNodePoolsForCluster(ctx, u.infraCli, network, cluster, oldNPs, desiredVersion, gke.Status.AtProvider.CurrentMasterVersion); err != nil {
			logger.Error(err, "failed to create new version tidb node pools")
		}

		// If there is no errors, need requeue to get new version node pools in the next reconcile.
		// Otherwise, `Requeue` will be ignored.
		return ctrl.Result{Requeue: err == nil}, err
	}

	// Add resource labels to the new version node pools.
	// We had better do this before the new version node pools scale out from zero,
	// otherwise, all new version nodes need to be recreated to add resource labels.
	if err := addResourceLabelsToNodePools(ctx, u.gcpClient, newNPs, getResourceLabels(network)); err != nil {
		logger.Error(err, "failed to add resource labels to the new version node pools")

		// If the gke cluster is reconciling, we should wait for next reconcile.
		if gke.Status.AtProvider.Status == gkev1beta2.ClusterStateReconciling {
			logger.Info("gke cluster is reconciling, wait for next reconcile")
			return ctrl.Result{RequeueAfter: 2 * time.Minute}, nil
		}

		return ctrl.Result{}, err
	}

	// Wait for new version node pools are ready
	ready := true
	for _, np := range newNPs {
		if np.Status.AtProvider.Status != gkev1beta1.NodePoolStateRunning {
			ready = false
			break
		}
	}
	if !ready {
		logger.Info("the new version node pools are not ready")
		return ctrl.Result{RequeueAfter: 1 * time.Minute}, nil
	}

	// Create over provisioning pods to scale out new version nodes first.
	// Otherwise, each pod will keep pending for about 70 seconds to be scheduled to a new version node.
	if err = nodepool.CreateOverProvisioningPods(ctx, gkeCli, cluster, consts.ClusterPrefix, cluster.PrefixedId(), desiredVersion, u.cfg, false); err != nil {
		logger.Error(err, "failed to create over provisioning pods")
		return ctrl.Result{}, err
	}
	// Wait for new version nodes ready
	if !nodepool.AllOverProvisioningPodsReady(ctx, gkeCli, consts.ClusterPrefix, cluster.PrefixedId()) {
		logger.Info("the new version nodes are not ready")
		return ctrl.Result{RequeueAfter: 2 * time.Minute}, nil
	}
	LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, "All new version nodes are ready", nil)

	// Update tc to restart the tidb cluster gracefully.
	var tc opv1alpha1.TidbCluster
	if err = gkeCli.Get(ctx, types.NamespacedName{Name: consts.ClusterPrefix, Namespace: cluster.PrefixedId()}, &tc); err != nil {
		logger.Error(err, "unable to get the tidbcluster CR")
		return ctrl.Result{}, err
	}
	if MutateTCForRestarting(&tc, desiredVersion) {
		logger.Info("update tidbcluster CR to restart the tidb cluster")

		nodepool.OverlayNodeSelectorWithK8sVersion(&tc, desiredVersion)
		if scaleOutExtraTiDB {
			tc.Spec.TiDB.Replicas = int32(cluster.Spec.Cluster.TiDB.Replicas + 1)
		}
		if err = k8s.Apply(ctx, gkeCli, &tc); err != nil {
			logger.Error(err, "unable to update the tidbcluster CR")
			return ctrl.Result{}, err
		}
		LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, "Restarting the TiDB Cluster", nil)
		return ctrl.Result{RequeueAfter: 5 * time.Minute}, nil
	}

	// In some case, tidb's asts may be not updated yet to scale out an extra tidb though tc is ready.
	// So we also check tidb's ready replicas to ensure tidb cluster is ready.
	if !IsTCReady(&tc) || tc.Spec.TiDB.Replicas != tc.Status.TiDB.StatefulSet.ReadyReplicas {
		LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, "Waiting for tidb cluster to be ready", nil)
		logger.Info("waiting for tidb cluster to be ready", "spec.tidb.replicas", tc.Spec.TiDB.Replicas, "status.tidb.readyReplicas", tc.Status.TiDB.StatefulSet.ReadyReplicas)
		return ctrl.Result{RequeueAfter: 1 * time.Minute}, nil
	}
	logger.Info("tidb cluster restarted and is ready")

	// Scale in the extra tidb
	if scaleOutExtraTiDB {
		tc.Spec.TiDB.Replicas = int32(cluster.Spec.Cluster.TiDB.Replicas)
		if err := k8s.Apply(ctx, gkeCli, &tc); err != nil {
			logger.Error(err, "unable to update the tidbcluster CR")
			return ctrl.Result{}, err
		}
		LogAndEmitEvent(logger, u.eventRecorder, uo, NormalEvent, "Scale in the extra tidb", nil)
	}

	logger.Info("updated the node version in cluster CR's status", "nodeVersion", desiredVersion)
	return ctrl.Result{Requeue: true}, UpdateClusterStatus(ctx, u.infraCli, cluster, desiredVersion)
}

// UpgradeEmptyNodePools upgrades node pools those do not have nodes.
func (u *GKEUpgrader) UpgradeEmptyNodePools(ctx context.Context, uo *cloudv1.UpgradeOperation, network *cloudv1.Network) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	// Check all node pools those do not have nodes, upgrade them directly.
	allNPs, err := getAllNodePoolsUnderNetwork(ctx, u.infraCli, network)
	if err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to get all node pools: %w", err)
	}

	var canUpgradeNPs []*gkev1beta1.NodePool
	for _, np := range allNPs {
		if size, err := getNodeCount(ctx, u.gcpClient, np); err == nil && size == 0 {
			canUpgradeNPs = append(canUpgradeNPs, np)
		}
	}
	logger.Info("node pools those do not have nodes", "node pools", GetCRNameList(canUpgradeNPs))
	return u.upgradeNodePools(ctx, uo, network, toNodePoolList(canUpgradeNPs), cloudv1.EmptyNodePools)
}

// upgradeNodePools upgrades node pools in a GKE cluster.
// It firstly checks if the control plane is upgraded to the desired version.
// If yes, it will upgrade node pools to the same version as the current control plane's version.
// Note: this function will NOT create new node pools.
func (u *GKEUpgrader) upgradeNodePools(ctx context.Context, uo *cloudv1.UpgradeOperation, network *cloudv1.Network, npList *gkev1beta1.NodePoolList, nodeType cloudv1.UpgradeOperationTarget) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	desiredVersion := uo.Status.CurrentStep.Version
	gke, err := getGKECR(ctx, u.infraCli, network)
	if err != nil {
		logger.Error(err, "unable to get the gke cluster CR")
		return ctrl.Result{}, err
	}
	if !isControlPlaneUpgraded(gke, desiredVersion) {
		LogAndEmitEvent(logger, u.eventRecorder, uo, WarningEvent, ControlPlaneNotUpgraded, nil)
		return ctrl.Result{}, uo.Fail(ctx, u.infraCli, ControlPlaneNotUpgraded)
	}

	// If the gke cluster is reconciling, we should wait for next reconcile.
	if gke.Status.AtProvider.Status == gkev1beta2.ClusterStateReconciling {
		logger.Info("gke cluster is reconciling, wait for next reconcile")
		return ctrl.Result{RequeueAfter: 2 * time.Minute}, nil
	}
	if gke.Status.AtProvider.Status != gkev1beta2.ClusterStateRunning {
		return ctrl.Result{}, fmt.Errorf("gke cluster %s is in unexpected status %s", gke.Name, gke.Status.AtProvider.Status)
	}

	// npSet is a set of node pool names to be upgraded.
	npSet := make(map[string]*gkev1beta1.NodePool, len(npList.Items))
	for i := range npList.Items {
		np := &npList.Items[i]
		npSet[np.Name] = np
	}
	toBeUpgradedNPs := make([]string, 0, len(npSet))

	// gke CR's status stores the current status of node pools.
	for _, npStatus := range gke.Status.AtProvider.NodePools {
		if _, ok := npSet[npStatus.Name]; ok {
			if npStatus.Version == gke.Status.AtProvider.CurrentMasterVersion && npStatus.Status == gkev1beta1.NodePoolStateRunning {
				logger.Info("node pool is already upgraded", "node pool", npStatus.Name, "version", npStatus.Version)
			} else {
				toBeUpgradedNPs = append(toBeUpgradedNPs, npStatus.Name)
			}
		}
	}
	// When there are no node pools left to be upgraded
	if len(toBeUpgradedNPs) == 0 {
		logger.Info(fmt.Sprintf("all node pools of %s are upgraded", nodeType))
		return ctrl.Result{}, FinishUpgrading(ctx, u.infraCli, uo)
	}

	logger.Info("node pools to be upgraded", "nodeType", nodeType, "node pools", toBeUpgradedNPs)
	gkeNetworkSpec, err := v1alpha1.GKENetworkSpecForNetwork(network)
	if err != nil {
		logger.Error(err, "unable to get gke network spec from network CR")
		return ctrl.Result{}, err
	}

	for _, npName := range toBeUpgradedNPs {
		np := npSet[npName]
		logger.Info("start to upgrade the node pool", "nodeType", nodeType, "name", np.Name, "spec version", np.Spec.ForProvider.Version)

		// Upgrade the node pool via gcloud sdk.
		// Note: we can not upgrade the node pool via CR.
		if err := u.gcpClient.UpgradeNodePool(ctx, gkeNetworkSpec.ProjectID, network.Spec.Region, gke.Name, np.Name, gke.Status.AtProvider.CurrentMasterVersion); err != nil {
			logger.Error(err, "unable to upgrade the node pool via sdk", "nodeType", nodeType, "name", np.Name)
			continue
		}

		// We also need to update the node pool CR's version in spec.
		np.Spec.ForProvider.Version = pointer.String(gke.Status.AtProvider.CurrentMasterVersion)
		if err = k8s.Apply(ctx, u.infraCli, np); err != nil {
			logger.Error(err, "unable to update the node pool CR", "name", np.Name)
			continue
		}
	}

	return ctrl.Result{RequeueAfter: 2 * time.Minute}, nil
}

func (u *GKEUpgrader) getGKEClient(ctx context.Context, network *cloudv1.Network, gke *gkev1beta2.Cluster) (client.Client, error) {
	c, ok := u.gkeClientMap.Load(string(gke.GetUID()))
	if ok {
		return c.(client.Client), nil
	}

	scheme, err := k8s.TiDBOperatorScheme()
	if err != nil {
		return nil, err
	}

	g, err := k8s.NetworkClientGetter(ctx, u.infraCli, network, scheme)
	if err != nil {
		return nil, err
	}

	gkeClient, err := g.GetClient(ctx)
	if err != nil {
		return nil, err
	}

	u.gkeClientMap.Store(string(gke.GetUID()), gkeClient)
	return gkeClient, nil
}