/*
Copyright 2024.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package clusterupgradeoperation

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"math/rand/v2"
	"strings"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	o11yv1 "github.com/tidbcloud/observability/api/regionalserver/resourceworker/v1"
	"github.com/tidbcloud/observability/pkg/common/constants"
	"github.com/tidbcloud/observability/pkg/resourceworker/controllers/clusterupgradeoperation/providers"
)

const (
	controlPlaneUpgradePollInterval = 2 * time.Minute
	controlPlaneUpgradeTimeout      = 45 * time.Minute

	addonUpgradePollInterval = 10 * time.Second
	addonUpgradeTimeout      = 30 * time.Minute

	nodeGroupUpgradeTimeout = 180 * time.Minute

	defaultVMClusterName = "vmcluster"
)

// processUpgradeControlPlane handles the control plane upgrade processor.
// All cloud providers start the upgrade asynchronously; we poll GetControlPlaneVersion
// until the version matches target or we hit the timeout.
func (r *ClusterUpgradeOperationReconciler) processUpgradeControlPlane(ctx context.Context, op *o11yv1.ClusterUpgradeOperation) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	logger.Info("Processing upgrade-controlplane step")

	// Parse options
	var opts o11yv1.ControlPlaneUpgradeOptions
	if op.Status.CurrentStep.Options.Raw != nil {
		if err := json.Unmarshal(op.Status.CurrentStep.Options.Raw, &opts); err != nil {
			op.FailOperation(fmt.Sprintf("failed to unmarshal controlplane options: %v", err))
			return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
		}
	}

	// Get cloud provider
	provider, err := r.cloudProviderRegistry.Get(op.Spec.CloudProvider)
	if err != nil {
		op.FailOperation(fmt.Sprintf("failed to get cloud provider: %v", err))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	clusterName := op.Spec.K8sClusterName
	targetVersion := opts.TargetVersion

	// 1. Check cluster status: if control plane is already upgrading, just requeue (no need to call UpgradeControlPlane again)
	status, err := provider.GetControlPlaneStatus(ctx, clusterName)
	if err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to get control plane status: %w", err)
	}
	if status == providers.ControlPlaneStatusUpdating {
		logger.Info("Control plane upgrade already in progress, will poll again", "cluster", clusterName)
		if err := r.updateStatusWithRetry(ctx, op); err != nil {
			return ctrl.Result{}, err
		}
		return ctrl.Result{RequeueAfter: controlPlaneUpgradePollInterval}, nil
	}

	// 2. Check current control plane version
	currentVersion, err := provider.GetControlPlaneVersion(ctx, clusterName)
	if err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to get control plane version: %w", err)
	}
	if status == providers.ControlPlaneStatusActive && k8sVersionsEqual(targetVersion, currentVersion) {
		// Already at target (e.g. upgrade completed in a previous reconcile)
		op.FinishCurrentStep()
		op.UpdateCurrentStep(time.Now())
		if err := r.updateStatusWithRetry(ctx, op); err != nil {
			return ctrl.Result{}, err
		}
		return requeue, nil
	}

	// 3. Enforce timeout (control plane upgrade is async and may take long)
	if op.Status.CurrentStep.BeginTime != nil &&
		time.Since(op.Status.CurrentStep.BeginTime.Time) > controlPlaneUpgradeTimeout {
		op.FailOperation(fmt.Sprintf("control plane upgrade did not complete within %v (current version: %q)", controlPlaneUpgradeTimeout, currentVersion))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	// 4. Start upgrade
	if err := provider.UpgradeControlPlane(ctx, clusterName, targetVersion); err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to upgrade control plane: %w", err)
	}

	// 5. Requeue to poll again
	if err := r.updateStatusWithRetry(ctx, op); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{RequeueAfter: controlPlaneUpgradePollInterval}, nil
}

// normalizeK8sVersion trims optional "v" prefix, spaces, and vendor suffix for comparison.
// e.g. "1.28" vs "1.28.0" kept as-is; "1.33.3-aliyun.1" becomes "1.33.3".
func normalizeK8sVersion(v string) string {
	v = strings.TrimSpace(v)
	if len(v) > 0 && v[0] == 'v' {
		v = v[1:]
	}
	if idx := strings.Index(v, "-"); idx > 0 {
		v = v[:idx]
	}
	return v
}

// k8sVersionsEqual compares current and target versions based on target format:
// - If target is x.y: compare first 2 parts (e.g. "1.28" matches "1.28.0")
// - If target is x.y.z: compare first 3 parts (e.g. "1.28.0" matches "1.28.0")
// - If target is x.y.z-xxx: compare full strings including suffix (e.g. "1.28.0-aliyun.1" must match exactly)
func k8sVersionsEqual(target, current string) bool {
	hasSuffix := strings.Index(target, "-") > 0
	if hasSuffix {
		// x.y.z-xxx: compare full strings (trim + remove v prefix only, keep suffix)
		tFull := strings.TrimSpace(target)
		cFull := strings.TrimSpace(current)
		if len(tFull) > 0 && tFull[0] == 'v' {
			tFull = tFull[1:]
		}
		if len(cFull) > 0 && cFull[0] == 'v' {
			cFull = cFull[1:]
		}
		return tFull == cFull
	}

	t := normalizeK8sVersion(target)
	c := normalizeK8sVersion(current)
	if t == "" || c == "" {
		return t == c
	}

	tParts := strings.Split(t, ".")
	cParts := strings.Split(c, ".")

	switch len(tParts) {
	case 2:
		// x.y: compare first 2 parts
		if len(cParts) < 2 {
			return false
		}
		return tParts[0] == cParts[0] && tParts[1] == cParts[1]
	case 3, 4, 5:
		// x.y.z: compare first 3 parts
		if len(cParts) < 3 {
			return false
		}
		return tParts[0] == cParts[0] && tParts[1] == cParts[1] && tParts[2] == cParts[2]
	default:
		return t == c
	}
}

func addonVersionsEqual(target, current string) bool {
	t := strings.TrimSpace(target)
	c := strings.TrimSpace(current)
	if len(t) > 0 && t[0] == 'v' {
		t = t[1:]
	}
	if len(c) > 0 && c[0] == 'v' {
		c = c[1:]
	}
	return t == c
}

// processUpgradeClusterAddons handles the cluster addons upgrade processor.
// Same pattern as control plane: trigger UpgradeClusterAddons (async), then poll GetClusterAddonsStatus and RequeueAfter.
func (r *ClusterUpgradeOperationReconciler) processUpgradeClusterAddons(ctx context.Context, op *o11yv1.ClusterUpgradeOperation) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	logger.Info("Processing upgrade-cluster-addons step")

	// Parse options
	var opts o11yv1.ClusterAddonsUpgradeOptions
	if op.Status.CurrentStep.Options.Raw != nil {
		if err := json.Unmarshal(op.Status.CurrentStep.Options.Raw, &opts); err != nil {
			op.FailOperation(fmt.Sprintf("failed to unmarshal cluster addons options: %v", err))
			return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
		}
	}

	// Get cloud provider
	provider, err := r.cloudProviderRegistry.Get(op.Spec.CloudProvider)
	if err != nil {
		op.FailOperation(fmt.Sprintf("failed to get cloud provider: %v", err))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	clusterName := op.Spec.K8sClusterName

	// 1. Check current addon status
	statuses, err := provider.GetClusterAddonsStatus(ctx, clusterName, opts.Addons)
	if err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to get cluster addons status: %w", err)
	}

	// 2. Fail on UPDATE_FAILED
	for _, s := range statuses {
		if s.Status == "UPDATE_FAILED" || s.Status == "CREATE_FAILED" || s.Status == "DELETE_FAILED" {
			op.FailOperation(fmt.Sprintf("addon %s in failure state: %s", s.Name, s.Status))
			return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
		}
	}

	// 3. All addons ACTIVE and at target version -> finish step
	allDone := true
	for i, s := range statuses {
		if s.Status != "ACTIVE" {
			allDone = false
			break
		}
		if !addonVersionsEqual(opts.Addons[i].Version, s.Version) {
			allDone = false
			break
		}
	}
	if allDone {
		op.FinishCurrentStep()
		op.UpdateCurrentStep(time.Now())
		if err := r.updateStatusWithRetry(ctx, op); err != nil {
			return ctrl.Result{}, err
		}
		return requeue, nil
	}

	// 4. Timeout check
	if op.Status.CurrentStep.BeginTime != nil &&
		time.Since(op.Status.CurrentStep.BeginTime.Time) > addonUpgradeTimeout {
		op.FailOperation(fmt.Sprintf("addon upgrade did not complete within %v", addonUpgradeTimeout))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	// 5. Any addon still UPDATING (or CREATING) -> just requeue (no need to call UpgradeClusterAddons again)
	for _, s := range statuses {
		if s.Status == "UPDATING" || s.Status == "CREATING" {
			logger.Info("Addon upgrade already in progress, will poll again", "addon", s.Name, "cluster", clusterName)
			if err := r.updateStatusWithRetry(ctx, op); err != nil {
				return ctrl.Result{}, err
			}
			return ctrl.Result{RequeueAfter: addonUpgradePollInterval}, nil
		}
	}

	// 6. Trigger upgrade (UpdateAddon is async)
	if err := provider.UpgradeClusterAddons(ctx, clusterName, opts.Addons); err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to upgrade cluster addons: %w", err)
	}

	// 7. Requeue to poll again
	if err := r.updateStatusWithRetry(ctx, op); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{RequeueAfter: addonUpgradePollInterval}, nil
}

// processUpgradeNodeGroup handles the nodegroup upgrade processor
// This consolidates: blue-green creation, VMCluster update, drain, and delete
func (r *ClusterUpgradeOperationReconciler) processUpgradeNodeGroup(ctx context.Context, op *o11yv1.ClusterUpgradeOperation) (ctrl.Result, error) {
	// Parse options
	var opts o11yv1.NodeGroupUpgradeOptions
	if op.Status.CurrentStep.Options.Raw != nil {
		if err := json.Unmarshal(op.Status.CurrentStep.Options.Raw, &opts); err != nil {
			op.FailOperation(fmt.Sprintf("failed to unmarshal nodegroup options: %v", err))
			return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
		}
	}

	logger := log.FromContext(ctx)
	logger.Info("Processing upgrade-node-group step", "cluster", op.Spec.K8sClusterName, "nodegroup", opts.NodeGroupName, "strategy", opts.Strategy)

	// Get cloud provider
	provider, err := r.cloudProviderRegistry.Get(op.Spec.CloudProvider)
	if err != nil {
		op.FailOperation(fmt.Sprintf("failed to get cloud provider: %v", err))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	// Get nodegroups by name
	nodeGroup, err := provider.GetNodeGroupByName(ctx, op.Spec.K8sClusterName, opts.NodeGroupName)
	if err != nil {
		if errors.Is(err, providers.ErrNodeGroupNotFound) {
			if opts.Strategy == o11yv1.NodeGroupUpgradeStrategyBlueGreen &&
				op.Status.CurrentStep.SubStep == o11yv1.BlueGreenSubStepDeleteOldNodeGroup {
				op.FinishCurrentStep()
				op.UpdateCurrentStep(time.Now())
				if err := r.updateStatusWithRetry(ctx, op); err != nil {
					return ctrl.Result{}, err
				}
				return requeue, nil
			}
		}
		return ctrl.Result{RequeueAfter: interval}, err
	}

	// Process based on strategy
	switch opts.Strategy {
	case o11yv1.NodeGroupUpgradeStrategyBlueGreen:
		return r.processBlueGreenNodeGroupUpgrade(ctx, op, provider, nodeGroup, opts)
	case o11yv1.NodeGroupUpgradeStrategyInPlace:
		return r.processInPlaceUpgrade(ctx, op, provider, nodeGroup, opts)
	default:
		op.FailOperation(fmt.Sprintf("unknown nodegroup upgrade strategy: %s", opts.Strategy))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}
}

// setBlueGreenSubStep sets CurrentStep.SubStep and SubStepBeginTime when the sub-step changes (for observability).
func setBlueGreenSubStep(ctx context.Context, op *o11yv1.ClusterUpgradeOperation, subStep string, now time.Time) {
	if op.Status.CurrentStep == nil {
		return
	}
	logger := log.FromContext(ctx)
	logger.Info("Setting blue green sub-step", "from", op.Status.CurrentStep.SubStep, "to", subStep)
	if op.Status.CurrentStep.SubStep != subStep {
		op.Status.CurrentStep.SubStep = subStep
		op.Status.CurrentStep.SubStepBeginTime = &metav1.Time{Time: now}
	}
}

// processBlueGreenNodeGroupUpgrade handles VMStorage nodegroup upgrade
// Fully stateless - derives all progress from actual resource state:
// - Old/new nodegroups determined by version check
// - Rolling update status determined by checking if pods are on new nodegroups and ready
// Flow: create all NGs -> wait all ready -> cordon old -> apply vmcluster -> wait pods on new nodes -> drain/delete old
// The function uses CurrentStep.SubStep to track progress and only executes the current sub-step.
func (r *ClusterUpgradeOperationReconciler) processBlueGreenNodeGroupUpgrade(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	provider providers.CloudProviderInterface,
	nodeGroup *providers.NodeGroup,
	opts o11yv1.NodeGroupUpgradeOptions,
) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	now := time.Now()

	// Get current sub-step, default to create-nodegroup if not set
	currentSubStep := ""
	if op.Status.CurrentStep != nil {
		currentSubStep = op.Status.CurrentStep.SubStep
	}
	if currentSubStep == "" {
		currentSubStep = o11yv1.BlueGreenSubStepCreateNodeGroup
	}

	// Get new nodegroup name
	var newNg *providers.NodeGroup = nil
	newGroupName := opts.NewNodeGroupName
	if newGroupName == "" {
		const letters = "abcdefghijklmnopqrstuvwxyz"
		b := make([]byte, 8)
		for i := range b {
			b[i] = letters[rand.IntN(len(letters))] // #nosec G404 -- nodegroup name does not require cryptographic randomness
		}
		newGroupName = string(b)
	}
	logger.Info("Getting new nodegroup", "newGroupName", newGroupName)
	if currentSubStep != o11yv1.BlueGreenSubStepCreateNodeGroup {
		var err error
		newNg, err = provider.GetNodeGroupByName(ctx, op.Spec.K8sClusterName, newGroupName)
		if err != nil {
			return ctrl.Result{RequeueAfter: interval}, fmt.Errorf("failed to get nodegroup by name: %w", err)
		}
	}
	// Execute sub-steps based on current state
	switch currentSubStep {
	case o11yv1.BlueGreenSubStepCreateNodeGroup:
		return r.handleCreateNodeGroup(ctx, op, provider, nodeGroup, newGroupName, now)

	case o11yv1.BlueGreenSubStepWaitReady:
		return r.handleWaitReady(ctx, op, newNg, now)

	case o11yv1.BlueGreenSubStepCordon:
		return r.handleCordon(ctx, op, provider, nodeGroup, opts, now)

	case o11yv1.BlueGreenSubStepVMStorageRollout:
		// Skip VMStorage rollout if not configured
		if opts.VMStorageSTSName == "" {
			op.FailOperation("VMStorage STS name is not configured, cannot proceed")
			return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
		}
		return r.handleVMStorageRollout(ctx, op, provider, newNg, opts, now)

	case o11yv1.BlueGreenSubStepDeleteOldNodeGroup:
		return r.handleDeleteOldNodeGroup(ctx, op, provider, nodeGroup, newNg, now)

	default:
		op.FailOperation(fmt.Sprintf("unknown blue green sub-step: %s", currentSubStep))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}
}

// handleCreateNodeGroup handles the create-nodegroup sub-step
func (r *ClusterUpgradeOperationReconciler) handleCreateNodeGroup(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	provider providers.CloudProviderInterface,
	nodeGroup *providers.NodeGroup,
	newGroupName string,
	now time.Time,
) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	// Create new nodegroup
	newSpec := providers.NodeGroup{
		Name:         newGroupName,
		Version:      op.Spec.TargetVersion,
		Labels:       nodeGroup.Labels,
		Tags:         nodeGroup.Tags,
		NodeRole:     nodeGroup.NodeRole,
		InstanceType: nodeGroup.InstanceType,
		MinSize:      nodeGroup.DesiredSize,
		MaxSize:      nodeGroup.MaxSize,
		DesiredSize:  nodeGroup.DesiredSize,
		Subnets:      append([]string(nil), nodeGroup.Subnets...),
		Locations:    append([]string(nil), nodeGroup.Locations...),
		Taints:       append([]providers.NodeGroupTaint(nil), nodeGroup.Taints...),
		AmiType:      nodeGroup.AmiType,
		CapacityType: nodeGroup.CapacityType,
		DiskSize:     nodeGroup.DiskSize,

		LaunchTemplateVersion: nodeGroup.LaunchTemplateVersion,
		LaunchTemplateId:      nodeGroup.LaunchTemplateId,
		LaunchTemplateName:    nodeGroup.LaunchTemplateName,
	}
	logger.Info("Creating new nodegroup", "nodeGroupName", nodeGroup.Name, "newNG", newSpec.Name)
	if _, err := provider.CreateNodeGroup(ctx, op.Spec.K8sClusterName, newSpec); err != nil {
		if errors.Is(err, providers.ErrNodeGroupLaunchTemplateCustomAMI) {
			op.FailOperation(fmt.Sprintf("nodegroup %s launch template specifies a custom AMI; blue-green upgrade cannot reuse it with target Kubernetes version %s", nodeGroup.Name, op.Spec.TargetVersion))
			return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
		}
		logger.Info("Failed to create new nodegroup, will retry", "nodegroup", nodeGroup.Name, "error", err)
		return ctrl.Result{RequeueAfter: 3 * interval}, nil
	}
	setBlueGreenSubStep(ctx, op, o11yv1.BlueGreenSubStepWaitReady, now)
	logger.Info("Creating new nodegroup", "nodegroup", nodeGroup.Name)
	if err := r.updateStatusWithRetry(ctx, op); err != nil {
		return ctrl.Result{}, err
	}
	return ctrl.Result{RequeueAfter: interval}, nil
}

// handleWaitReady handles the wait-ready sub-step
// Precondition: newNg is not nil (checked in the caller's switch case)
func (r *ClusterUpgradeOperationReconciler) handleWaitReady(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	newNg *providers.NodeGroup,
	now time.Time,
) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	if newNg.Status == providers.NodeGroupStatusCreateFailed {
		op.FailOperation(fmt.Sprintf("new nodegroup %s is in failure state: %s", newNg.Name, newNg.Status))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	if newNg.Status != providers.NodeGroupStatusActive {
		logger.Info("Waiting for nodegroup to become active", "nodegroup", newNg.Name, "status", newNg.Status)
		return ctrl.Result{RequeueAfter: interval}, nil
	}

	logger.Info("New nodegroup is ready", "nodegroup", newNg.Name, "id", newNg.ID)
	// Move to next step, let controller handle cordon on next reconcile
	setBlueGreenSubStep(ctx, op, o11yv1.BlueGreenSubStepCordon, now)
	if err := r.updateStatusWithRetry(ctx, op); err != nil {
		return ctrl.Result{}, err
	}
	return requeue, nil
}

// handleCordon handles the cordon sub-step
func (r *ClusterUpgradeOperationReconciler) handleCordon(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	provider providers.CloudProviderInterface,
	nodeGroup *providers.NodeGroup,
	opts o11yv1.NodeGroupUpgradeOptions,
	now time.Time,
) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	logger.Info("Cordoning nodes in old nodegroup", "nodegroup", nodeGroup.Name)
	if err := r.cordonNodesByNodeGroup(ctx, provider, op, nodeGroup); err != nil {
		return ctrl.Result{}, fmt.Errorf("failed to cordon nodegroup %s: %w", nodeGroup.Name, err)
	}

	// Move to next step, let controller handle it on next reconcile
	if opts.RequiresVMStorageRollout {
		setBlueGreenSubStep(ctx, op, o11yv1.BlueGreenSubStepVMStorageRollout, now)
	} else {
		setBlueGreenSubStep(ctx, op, o11yv1.BlueGreenSubStepDeleteOldNodeGroup, now)
	}
	if err := r.updateStatusWithRetry(ctx, op); err != nil {
		return ctrl.Result{}, err
	}
	return requeue, nil
}

// handleVMStorageRollout handles the vmstorage-rollout sub-step
func (r *ClusterUpgradeOperationReconciler) handleVMStorageRollout(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	provider providers.CloudProviderInterface,
	newNg *providers.NodeGroup,
	opts o11yv1.NodeGroupUpgradeOptions,
	now time.Time,
) (ctrl.Result, error) {
	logger := log.FromContext(ctx)
	sts := &appsv1.StatefulSet{}
	if err := r.Get(ctx, types.NamespacedName{Namespace: opts.VMStorageSTSNamespace, Name: opts.VMStorageSTSName}, sts); err != nil {
		return ctrl.Result{}, err
	}

	// 1. Check if triggered from the source-of-truth VMCluster CR.
	triggered := false
	if opts.VMClusterHelmValues != nil {
		var err error
		triggered, err = r.hasVMClusterUpgradeAnnotation(ctx, op)
		if err != nil {
			return ctrl.Result{}, err
		}
	}
	if !triggered {
		logger.Info("VMStorage rolling update not triggered, triggering now")
		if err := r.applyAndTriggerVMStorageUpdate(ctx, op, opts); err != nil {
			return ctrl.Result{}, err
		}
		// Requeue to observe status change
		return ctrl.Result{RequeueAfter: interval}, nil
	}

	// 2. Check stability
	isRolledOut := r.isStatefulSetRolledOut(sts)
	if !isRolledOut {
		logger.Info("VMStorage rollout in progress",
			"generation", sts.Generation,
			"observedGeneration", sts.Status.ObservedGeneration,
			"currentRevision", sts.Status.CurrentRevision,
			"updateRevision", sts.Status.UpdateRevision,
			"replicas", *sts.Spec.Replicas,
			"readyReplicas", sts.Status.ReadyReplicas,
			"updatedReplicas", sts.Status.UpdatedReplicas)
		return ctrl.Result{RequeueAfter: interval}, nil
	}

	// 3. Check if pods are on new nodegroups (Goal)
	onNewNodes, err := r.arePodsOnNewNodes(ctx, provider, op, sts, newNg)
	if err != nil {
		return ctrl.Result{}, err
	}
	if !onNewNodes {
		waitStart := time.Time{}
		if op.Status.CurrentStep != nil {
			if op.Status.CurrentStep.SubStepBeginTime != nil {
				waitStart = op.Status.CurrentStep.SubStepBeginTime.Time
			} else if op.Status.CurrentStep.BeginTime != nil {
				waitStart = op.Status.CurrentStep.BeginTime.Time
			}
		}
		if !waitStart.IsZero() && time.Since(waitStart) <= nodeGroupUpgradeTimeout {
			logger.Info("VMStorage pods are not on new nodes yet, waiting for VM operator rollout")
			return ctrl.Result{RequeueAfter: interval}, nil
		}
		// 4. Stable but still on old nodes after timeout - Manual intervention needed
		op.FailOperation("VMStorage stable but pods still on old nodes, manual intervention needed")
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	// Move to next step, let controller handle delete on next reconcile
	setBlueGreenSubStep(ctx, op, o11yv1.BlueGreenSubStepDeleteOldNodeGroup, now)
	if err := r.updateStatusWithRetry(ctx, op); err != nil {
		return ctrl.Result{}, err
	}
	return requeue, nil
}

// handleDeleteOldNodeGroup handles the delete-old-nodegroup sub-step
func (r *ClusterUpgradeOperationReconciler) handleDeleteOldNodeGroup(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	provider providers.CloudProviderInterface,
	nodeGroup *providers.NodeGroup,
	newNg *providers.NodeGroup,
	now time.Time,
) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	// nodeGroup is never nil when we reach DeleteOldNodeGroup (caller ensures it)
	status, err := provider.GetNodeGroupStatus(ctx, op.Spec.K8sClusterName, nodeGroup.ID)
	if err != nil {
		if !errors.Is(err, providers.ErrNodeGroupNotFound) {
			return ctrl.Result{}, fmt.Errorf("failed to get old nodegroup status: %w", err)
		}
		// Nodegroup doesn't exist anymore (e.g. already deleted), we're done
		logger.Info("Old nodegroup no longer exists, upgrade complete", "nodegroup", nodeGroup.Name)
		if err := r.restoreNewNodeGroupMinSize(ctx, op, provider, newNg); err != nil {
			return ctrl.Result{}, err
		}
		op.FinishCurrentStep()
		op.UpdateCurrentStep(now)
		if err := r.updateStatusWithRetry(ctx, op); err != nil {
			return ctrl.Result{}, err
		}
		return requeue, nil
	}

	if status == providers.NodeGroupStatusDeleteFailed {
		op.FailOperation(fmt.Sprintf("old nodegroup %s is in failure state: %s", nodeGroup.Name, status))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	if status == providers.NodeGroupStatusDeleting {
		logger.Info("Old nodegroup is already being deleted", "nodegroup", nodeGroup.Name)
		return ctrl.Result{RequeueAfter: interval}, nil
	}

	if status == providers.NodeGroupStatusActive {
		logger.Info("Deleting old nodegroup", "nodegroup", nodeGroup.Name)
		if err := provider.DeleteNodeGroup(ctx, op.Spec.K8sClusterName, nodeGroup.ID); err != nil {
			return ctrl.Result{}, fmt.Errorf("failed to delete nodegroup %s: %w", nodeGroup.Name, err)
		}
		return ctrl.Result{RequeueAfter: interval}, nil
	}

	// Unknown status, requeue
	return ctrl.Result{RequeueAfter: interval}, nil
}

func (r *ClusterUpgradeOperationReconciler) restoreNewNodeGroupMinSize(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	provider providers.CloudProviderInterface,
	newNg *providers.NodeGroup,
) error {
	if newNg == nil || newNg.MinSize == 1 {
		return nil
	}
	nodeGroupID := newNg.ID
	if nodeGroupID == "" {
		nodeGroupID = newNg.Name
	}
	logger := log.FromContext(ctx)
	logger.Info("Restoring new nodegroup min size", "nodegroup", newNg.Name, "from", newNg.MinSize, "to", 1)
	if err := provider.UpdateNodeGroupMinSize(ctx, op.Spec.K8sClusterName, nodeGroupID, 1); err != nil {
		return fmt.Errorf("failed to restore new nodegroup min size: %w", err)
	}
	return nil
}

// isStatefulSetRolledOut checks if a StatefulSet is fully rolled out and stable
// Conditions:
// 1. ObservedGeneration >= Generation (Controller acknowledged update)
// 2. UpdatedReplicas == Replicas (All replicas on new revision - primary indicator)
// 3. ReadyReplicas == Replicas (All replicas healthy)
//
// Note: We intentionally do NOT require CurrentRevision == UpdateRevision.
// Kubernetes updates CurrentRevision in completeRollingUpdate only after the above
// conditions are met, but there can be timing delays or edge cases where
// CurrentRevision lags behind. UpdatedReplicas is the authoritative signal that
// all pods have been updated to the new revision.
func (r *ClusterUpgradeOperationReconciler) isStatefulSetRolledOut(sts *appsv1.StatefulSet) bool {
	if sts.Status.ObservedGeneration < sts.Generation {
		return false
	}
	if sts.Status.UpdatedReplicas != *sts.Spec.Replicas {
		return false
	}
	if sts.Status.ReadyReplicas != *sts.Spec.Replicas {
		return false
	}
	return true
}

// arePodsOnNewNodes checks if all pods of the StatefulSet are scheduled on the new nodegroups
func (r *ClusterUpgradeOperationReconciler) arePodsOnNewNodes(
	ctx context.Context,
	provider providers.CloudProviderInterface,
	op *o11yv1.ClusterUpgradeOperation,
	sts *appsv1.StatefulSet,
	newNodeGroups *providers.NodeGroup,
) (bool, error) {
	logger := log.FromContext(ctx)

	// Identify nodes belonging to new nodegroups
	newNodeNames := make(map[string]bool)
	// Get label selector for this nodegroup
	labelSelector, err := provider.GetNodeGroupLabelSelector(ctx, op.Spec.K8sClusterName, newNodeGroups.ID)
	if err != nil {
		return false, fmt.Errorf("failed to get label selector for nodegroup %s: %w", newNodeGroups.Name, err)
	}

	// List nodes
	nodeList := &corev1.NodeList{}
	if err := r.List(ctx, nodeList, client.MatchingLabels(labelSelector)); err != nil {
		return false, fmt.Errorf("failed to list nodes for nodegroup %s: %w", newNodeGroups.Name, err)
	}
	for _, node := range nodeList.Items {
		newNodeNames[node.Name] = true
	}

	// List pods of STS
	podList := &corev1.PodList{}
	if err := r.List(ctx, podList, client.InNamespace(sts.Namespace), client.MatchingLabels(sts.Spec.Selector.MatchLabels)); err != nil {
		return false, err
	}

	// Check validity of pod list
	if len(podList.Items) == 0 {
		if *sts.Spec.Replicas > 0 {
			// Pods missing but replicas > 0
			return false, nil
		}
		return true, nil
	}

	for _, pod := range podList.Items {
		if !newNodeNames[pod.Spec.NodeName] {
			logger.Info("Pod still on old node", "pod", pod.Name, "node", pod.Spec.NodeName)
			return false, nil
		}
	}

	return true, nil
}

// hasVMClusterUpgradeAnnotation checks if the source-of-truth VMCluster has the update annotation.
func (r *ClusterUpgradeOperationReconciler) hasVMClusterUpgradeAnnotation(ctx context.Context, op *o11yv1.ClusterUpgradeOperation) (bool, error) {
	vmcluster := newVMClusterObject()
	key := types.NamespacedName{
		Namespace: constants.ObservabilityNamespace,
		Name:      defaultVMClusterName,
	}
	if err := r.Get(ctx, key, vmcluster); err != nil {
		return false, fmt.Errorf("get VMCluster %s/%s: %w", key.Namespace, key.Name, err)
	}
	return hasClusterUpgradeAnnotation(vmcluster.GetAnnotations(), op), nil
}

func hasClusterUpgradeAnnotation(annotations map[string]string, op *o11yv1.ClusterUpgradeOperation) bool {
	if annotations == nil {
		return false
	}
	if owner, ok := annotations[o11yv1.ClusterUpgradeAnnotationLastUpdatedBy]; ok && owner == op.Name {
		return true
	}
	if val, ok := annotations[o11yv1.ClusterUpgradeAnnotationLastUpdatedAt]; ok {
		lastUpdate, err := time.Parse(time.RFC3339, val)
		if err == nil {
			return !lastUpdate.Before(op.Spec.Schedule.Time)
		}
	}
	return false
}

// applyAndTriggerVMStorageUpdate applies VMCluster spec overlay and annotates the VMCluster CR to trigger rollout
func (r *ClusterUpgradeOperationReconciler) applyAndTriggerVMStorageUpdate(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	opts o11yv1.NodeGroupUpgradeOptions,
) error {
	if opts.VMClusterHelmValues == nil {
		return fmt.Errorf("vmClusterHelmValues is required for VMStorage rollout; direct StatefulSet patching is not supported")
	}
	if err := r.applyVMClusterHelmValues(ctx, op, opts); err != nil {
		return fmt.Errorf("failed to apply VMCluster helm values: %v", err)
	}
	return nil
}

// cordonNodesByNodeGroup cordons all nodes in a nodegroup using K8s client directly
// This is cloud-agnostic since it operates on K8s Node resources
func (r *ClusterUpgradeOperationReconciler) cordonNodesByNodeGroup(
	ctx context.Context,
	provider providers.CloudProviderInterface,
	op *o11yv1.ClusterUpgradeOperation,
	ng *providers.NodeGroup,
) error {
	logger := log.FromContext(ctx)

	// Get label selector for this nodegroup from cloud provider
	labelSelector, err := provider.GetNodeGroupLabelSelector(ctx, op.Spec.K8sClusterName, ng.ID)
	if err != nil {
		return fmt.Errorf("failed to get label selector for nodegroup %s: %w", ng.Name, err)
	}

	// List nodes matching the selector
	nodeList := &corev1.NodeList{}
	if err := r.List(ctx, nodeList, client.MatchingLabels(labelSelector)); err != nil {
		return fmt.Errorf("failed to list nodes for nodegroup %s: %w", ng.Name, err)
	}

	logger.Info("Cordoning nodes in nodegroup", "nodegroup", ng.Name, "nodeCount", len(nodeList.Items))

	for _, node := range nodeList.Items {
		if node.Spec.Unschedulable {
			// Already cordoned
			continue
		}

		patch := client.MergeFrom(node.DeepCopy())
		node.Spec.Unschedulable = true
		if err := r.Patch(ctx, &node, patch); err != nil {
			return fmt.Errorf("failed to cordon node %s: %w", node.Name, err)
		}
		logger.Info("Cordoned node", "node", node.Name)
	}

	return nil
}

// processInPlaceUpgrade handles in-place nodegroup upgrade
func (r *ClusterUpgradeOperationReconciler) processInPlaceUpgrade(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	provider providers.CloudProviderInterface,
	nodeGroup *providers.NodeGroup,
	opts o11yv1.NodeGroupUpgradeOptions,
) (ctrl.Result, error) {
	logger := log.FromContext(ctx)

	logger.Info("Processing in-place upgrade for nodegroup", "nodegroup", nodeGroup.Name,
		"targetVersion", op.Spec.TargetVersion,
		"currentStatus", nodeGroup.Status,
		"currentVersion", nodeGroup.Version,
		"currentRuntimeVersion", nodeGroup.RuntimeVersion,
		"cloudOperationID", op.Status.CurrentStep.CloudOperationID)

	clusterName := op.Spec.K8sClusterName
	targetVersion := op.Spec.TargetVersion
	currentStatus := nodeGroup.Status
	currentVersion := nodeGroup.Version
	currentRuntimeVersion := nodeGroup.RuntimeVersion
	// It's for ACK. In ack we couldn't get the kubelet version from the nodegroup.
	if currentVersion == "" {
		currentVersion = GetNodeGroupVersionFromNodes(ctx, r.Client, provider, clusterName, nodeGroup)
	}

	// 1) Failure states: fail fast to avoid infinite polling
	if currentStatus == providers.NodeGroupStatusCreateFailed || currentStatus == providers.NodeGroupStatusDeleteFailed {
		op.FailOperation(fmt.Sprintf("nodegroup %s is in failure state: %s", nodeGroup.Name, currentStatus))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	// 2) If we have a cloud operation ID (e.g. GKE), poll it instead of relying on node group status
	if op.Status.CurrentStep.CloudOperationID != "" {
		done, err := provider.GetNodeGroupUpgradeOperationStatus(ctx, clusterName, nodeGroup.ID, op.Status.CurrentStep.CloudOperationID)
		if err != nil {
			op.FailOperation(fmt.Sprintf("nodegroup %s upgrade failed: %v", nodeGroup.Name, err))
			return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
		}
		if !done {
			logger.Info("Nodegroup upgrade operation in progress", "nodegroup", nodeGroup.Name, "operationID", op.Status.CurrentStep.CloudOperationID)
			return ctrl.Result{RequeueAfter: interval}, nil
		}
		op.FinishCurrentStep()
		op.UpdateCurrentStep(time.Now())
		if err := r.updateStatusWithRetry(ctx, op); err != nil {
			return ctrl.Result{}, err
		}
		return requeue, nil
	}

	// 3) If already updating (from GetNodeGroupStatus), just poll again
	if currentStatus == providers.NodeGroupStatusUpdating {
		logger.Info("Nodegroup upgrade in progress", "nodegroup", nodeGroup.Name, "currentVersion", currentVersion)
		return ctrl.Result{RequeueAfter: interval}, nil
	}

	// 4) Already at target and stable -> finish step
	if currentStatus == providers.NodeGroupStatusActive &&
		(currentVersion == "" || k8sVersionsEqual(targetVersion, currentVersion)) &&
		(opts.TargetRuntimeVersion == "" || opts.TargetRuntimeVersion == currentRuntimeVersion) {
		op.FinishCurrentStep()
		op.UpdateCurrentStep(time.Now())
		if err := r.updateStatusWithRetry(ctx, op); err != nil {
			return ctrl.Result{}, err
		}
		return requeue, nil
	}

	// 5) Timeout guard
	if op.Status.CurrentStep.BeginTime != nil &&
		time.Since(op.Status.CurrentStep.BeginTime.Time) > nodeGroupUpgradeTimeout {
		op.FailOperation(fmt.Sprintf("nodegroup upgrade did not complete within %v (status: %s, current version: %q, target: %q)",
			nodeGroupUpgradeTimeout, currentStatus, currentVersion, targetVersion))
		return ctrl.Result{}, r.updateStatusWithRetry(ctx, op)
	}

	// 6) If nodegroup is not ACTIVE, just wait until it's stable before trying to trigger an in-place upgrade.
	if currentStatus != providers.NodeGroupStatusActive {
		logger.Info("Waiting for nodegroup to become ACTIVE", "nodegroup", nodeGroup.Name, "status", currentStatus)
		return ctrl.Result{RequeueAfter: interval}, nil
	}

	// 7) Trigger upgrade; retry on failure (controller will requeue until success).
	operationID, err := provider.UpgradeNodeGroup(ctx, clusterName, nodeGroup.ID, targetVersion, &opts)
	if err != nil {
		logger.Info("UpgradeNodeGroup failed, will retry", "nodegroup", nodeGroup.Name, "error", err)
		return ctrl.Result{RequeueAfter: interval}, nil
	}
	if operationID != "" && op.Status.CurrentStep != nil {
		op.Status.CurrentStep.CloudOperationID = operationID
	}
	logger.Info("Triggered in-place upgrade for nodegroup", "nodegroup", nodeGroup.Name, "from", currentVersion, "to", targetVersion, "operationID", operationID)

	return ctrl.Result{RequeueAfter: interval}, r.updateStatusWithRetry(ctx, op)

}

// applyVMClusterHelmValues applies a VMCluster.spec overlay to trigger VM operator rollout
func (r *ClusterUpgradeOperationReconciler) applyVMClusterHelmValues(
	ctx context.Context,
	op *o11yv1.ClusterUpgradeOperation,
	opts o11yv1.NodeGroupUpgradeOptions,
) error {
	logger := log.FromContext(ctx)
	if opts.VMClusterHelmValues == nil || len(opts.VMClusterHelmValues.Raw) == 0 {
		return nil
	}

	// NOTE: Despite the field name (HelmValues), in our dataplane the source-of-truth is VM operator's VMCluster CR.
	// We treat this JSON blob as a partial "spec" overlay and merge it into the existing VMCluster.spec.
	logger.Info("Applying VMCluster spec overlay for VM operator", "values", string(opts.VMClusterHelmValues.Raw))

	var overlayAny any
	if err := json.Unmarshal(opts.VMClusterHelmValues.Raw, &overlayAny); err != nil {
		return fmt.Errorf("invalid vmClusterHelmValues (must be JSON object): %w", err)
	}
	overlayObj, ok := overlayAny.(map[string]any)
	if !ok {
		return fmt.Errorf("invalid vmClusterHelmValues: expected JSON object, got %T", overlayAny)
	}
	// Support passing either a full resource (with {"spec": {...}}) or a bare spec overlay ({"vmstorage": {...}}).
	if specAny, ok := overlayObj["spec"]; ok {
		if specMap, ok := specAny.(map[string]any); ok {
			overlayObj = specMap
		}
	}

	vmcluster := newVMClusterObject()

	key := types.NamespacedName{
		Namespace: constants.ObservabilityNamespace,
		Name:      defaultVMClusterName,
	}
	if err := r.Get(ctx, key, vmcluster); err != nil {
		return fmt.Errorf("get VMCluster %s/%s: %w", key.Namespace, key.Name, err)
	}

	old := vmcluster.DeepCopy()

	currentSpec, found, err := unstructured.NestedMap(vmcluster.Object, "spec")
	if err != nil {
		return fmt.Errorf("read VMCluster.spec: %w", err)
	}
	if !found || currentSpec == nil {
		currentSpec = map[string]any{}
	}

	deepMergeMap(currentSpec, overlayObj)
	if err := unstructured.SetNestedMap(vmcluster.Object, currentSpec, "spec"); err != nil {
		return fmt.Errorf("set VMCluster.spec: %w", err)
	}

	ann := vmcluster.GetAnnotations()
	if ann == nil {
		ann = map[string]string{}
	}
	ann[o11yv1.ClusterUpgradeAnnotationLastUpdatedAt] = time.Now().Format(time.RFC3339)
	ann[o11yv1.ClusterUpgradeAnnotationLastUpdatedBy] = op.Name
	vmcluster.SetAnnotations(ann)

	if err := r.Patch(ctx, vmcluster, client.MergeFrom(old)); err != nil {
		return fmt.Errorf("patch VMCluster %s/%s: %w", key.Namespace, key.Name, err)
	}
	return nil
}

// deepMergeMap recursively merges src into dst.
// - map values are merged recursively
// - non-map values (including arrays) overwrite dst
// - null values delete the key from dst
func deepMergeMap(dst map[string]any, src map[string]any) {
	for k, v := range src {
		if v == nil {
			delete(dst, k)
			continue
		}
		srcMap, srcIsMap := v.(map[string]any)
		if !srcIsMap {
			dst[k] = v
			continue
		}
		if cur, ok := dst[k]; ok {
			if curMap, ok := cur.(map[string]any); ok {
				deepMergeMap(curMap, srcMap)
				dst[k] = curMap
				continue
			}
		}
		// Replace when dst is not a map (or missing).
		dst[k] = srcMap
	}
}

func newVMClusterObject() *unstructured.Unstructured {
	vmcluster := &unstructured.Unstructured{}
	vmcluster.SetGroupVersionKind(schema.GroupVersionKind{
		Group:   "operator.victoriametrics.com",
		Version: "v1beta1",
		Kind:    "VMCluster",
	})
	return vmcluster
}