package server

import (
	"context"
	"fmt"
	"slices"
	"sort"
	"strings"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	"github.com/tidbcloud/cluster-service-ng/pkg/nodegroup"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"github.com/tidbcloud/cluster-service-ng/pkg/tioperator"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
)

func (s *serverImpl) CreateNodeInstanceGroup(ctx context.Context, group *entity.NodeInstanceGroup) (*entity.NodeInstanceGroup, error) {
	if group == nil {
		return nil, status.Errorf(codes.InvalidArgument, "node instance group is nil")
	}
	if group.GetNodeSpec() == nil {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec is required")
	}

	nodeSpec := entity.GetComponentNodeSpec(group.GetNodeSpec()) // canonical uppercase, e.g. "8C16G"
	if err := nodegroup.ValidateSpec(nodeSpec); err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid node_spec: %v", err)
	}
	if err := nodegroup.ValidateCloudInstanceType(group.GetPrimaryInstanceType()); err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid primary_instance_type: %v", err)
	}

	fallbacks := group.GetFallbackInstances()
	if err := validateFallbackInstances(fallbacks, group.GetNodeSpec()); err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid fallback_instances: %v", err)
	}
	for _, fb := range fallbacks {
		if fb.GetInstanceType() == group.GetPrimaryInstanceType() {
			return nil, status.Errorf(codes.InvalidArgument, "fallback instance_type cannot be the same as primary: %s", group.GetPrimaryInstanceType())
		}
	}
	if group.GetNodeSpec().GetWithLocalDisk() && len(fallbacks) > 0 {
		return nil, status.Errorf(codes.InvalidArgument, "fallback_instances is not allowed when with_local_disk is true")
	}

	cpu, memGi, err := nodegroup.ParseSpec(nodeSpec)
	if err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid node_spec: %v", err)
	}
	if group.GetNodeSpec().GetCpu() != 0 && group.GetNodeSpec().GetCpu() != cpu {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec cpu mismatch: %d != %d", group.GetNodeSpec().GetCpu(), cpu)
	}
	if group.GetNodeSpec().GetMemoryGi() != 0 && group.GetNodeSpec().GetMemoryGi() != memGi {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec memory mismatch: %d != %d", group.GetNodeSpec().GetMemoryGi(), memGi)
	}
	if group.GetNodeSpec().GetCpu() == 0 {
		group.NodeSpec.Cpu = cpu
	}
	if group.GetNodeSpec().GetMemoryGi() == 0 {
		group.NodeSpec.MemoryGi = memGi
	}

	group.Name = nodeSpec
	group.Status = apiv1.NodeInstanceGroup_PENDING
	group.FallbackInstances = fallbacks
	group.DisabledInstanceTypes = nil
	if group.GetNodeSpec().GetAvailableCpu() == 0 {
		group.NodeSpec.AvailableCpu = defaultAvailableCPU(group.GetNodeSpec().GetCpu())
	}
	if group.GetNodeSpec().GetAvailableMemoryMi() == 0 {
		group.NodeSpec.AvailableMemoryMi = uint32(tioperator.ComputeAvailableMemory(uint64(group.GetNodeSpec().GetMemoryGi()) * 1024))
	}
	// Priority defaults to 1000 for primary instances if not specified.
	// This ensures the "primary >= 1000" invariant and consistent API responses.
	if group.GetPriority() == 0 {
		group.Priority = nodegroup.PriorityPrimary
	}

	var created *entity.NodeInstanceGroup
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		existing, err := s.nodeGroupRepo.LoadBySpecAndPrimary(ctx, nodeSpec, group.GetPrimaryInstanceType())
		if err != nil {
			return status.Errorf(codes.Internal, "failed to load node instance group: %v", err)
		}
		if existing != nil {
			return status.Errorf(codes.AlreadyExists, "node instance group already exists for node_spec %s with primary %s", nodeSpec, group.GetPrimaryInstanceType())
		}

		if err := s.nodeGroupRepo.Save(ctx, group); err != nil {
			return status.Errorf(codes.Internal, "failed to save node instance group: %v", err)
		}
		created = group

		return nil
	}); err != nil {
		return nil, err
	}

	return created, nil
}

func (s *serverImpl) ListNodeInstanceGroups(ctx context.Context, nodeSpec string) ([]*entity.NodeInstanceGroup, error) {
	if nodeSpec == "" {
		return s.nodeGroupRepo.ListAll(ctx)
	}

	return s.nodeGroupRepo.ListBySpec(ctx, nodeSpec)
}

func (s *serverImpl) DeprecateNodeInstanceGroup(ctx context.Context, nodeSpec, primaryInstance string) (*entity.NodeInstanceGroup, error) {
	if nodeSpec == "" || primaryInstance == "" {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec and primary_instance are required")
	}
	return s.SetNodeInstanceGroupStatus(ctx, nodeSpec, primaryInstance, apiv1.NodeInstanceGroup_DEPRECATED)
}

func (s *serverImpl) SetNodeInstanceGroupStatus(ctx context.Context, nodeSpec, primaryInstance string, nodeGroupStatus apiv1.NodeInstanceGroup_Status) (*entity.NodeInstanceGroup, error) {
	if nodeSpec == "" || primaryInstance == "" {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec and primary_instance are required")
	}
	switch nodeGroupStatus {
	case apiv1.NodeInstanceGroup_ACTIVE, apiv1.NodeInstanceGroup_DEPRECATED, apiv1.NodeInstanceGroup_RETIRING:
	default:
		return nil, status.Errorf(codes.InvalidArgument, "invalid status: %v, only ACTIVE, DEPRECATED, and RETIRING are allowed", nodeGroupStatus)
	}

	// Fetch non-deleting dataplanes before the transaction to avoid holding DB locks during external gRPC calls.
	// Deleting dataplanes are filtered out at the list level so the transaction loop needs no per-dataplane gRPC check.
	dataplaneNames, err := s.listNonDeletingDataplaneNames(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list dataplanes: %v", err)
	}

	var updated *entity.NodeInstanceGroup
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		group, err := s.nodeGroupRepo.LoadBySpecAndPrimary(ctx, nodeSpec, primaryInstance)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to load node instance group: %v", err)
		}
		if group == nil {
			return status.Errorf(codes.NotFound, "node instance group not found: %s/%s", nodeSpec, primaryInstance)
		}

		if group.GetStatus() == apiv1.NodeInstanceGroup_RETIRING || group.GetStatus() == apiv1.NodeInstanceGroup_RETIRED {
			return status.Errorf(codes.FailedPrecondition, "cannot transition from %v to %v", group.GetStatus(), nodeGroupStatus)
		}

		// RETIRING requires the group to be DEPRECATED and all dataplane overrides for this
		// (nodeSpec, primaryInstance) to also be DEPRECATED (or absent, meaning they follow the region-level status).
		if nodeGroupStatus == apiv1.NodeInstanceGroup_RETIRING {
			if group.GetStatus() != apiv1.NodeInstanceGroup_DEPRECATED {
				return status.Errorf(codes.FailedPrecondition, "can only transition to RETIRING from DEPRECATED, current status is %v", group.GetStatus())
			}
			overrides, err := s.dataplaneNodeGroupRepo.ListBySpec(ctx, nodeSpec)
			if err != nil {
				return status.Errorf(codes.Internal, "failed to list dataplane overrides: %v", err)
			}
			for _, override := range overrides {
				if override.GetPrimaryInstanceType() != primaryInstance {
					continue
				}
				if override.GetStatus() != apiv1.DataplaneNodeInstanceGroup_DEPRECATED {
					return status.Errorf(codes.FailedPrecondition,
						"dataplane override on %s is not DEPRECATED (status=%v), all overrides must be DEPRECATED before RETIRING",
						override.GetDataplaneName(), override.GetStatus())
				}
			}
		}

		oldGroupStatus := group.GetStatus()
		group.Status = nodeGroupStatus
		if err := s.nodeGroupRepo.Save(ctx, group); err != nil {
			return status.Errorf(codes.Internal, "failed to update node instance group: %v", err)
		}
		updated = group

		for _, dataplaneName := range dataplaneNames {
			override, err := s.dataplaneNodeGroupRepo.LoadByDataplaneAndSpec(ctx, dataplaneName, nodeSpec, group.GetPrimaryInstanceType())
			if err != nil {
				return status.Errorf(codes.Internal, "failed to load dataplane override: %v", err)
			}

			switch nodeGroupStatus {
			case apiv1.NodeInstanceGroup_ACTIVE:
				if override != nil && override.GetStatus() == apiv1.DataplaneNodeInstanceGroup_DEPRECATED {
					// Respect per-dataplane exclusion: skip dataplanes that explicitly deprecated this group.
					continue
				}
				// Clean up redundant ACTIVE overrides: region-level ACTIVE already covers all dataplanes.
				if override != nil && override.GetStatus() == apiv1.DataplaneNodeInstanceGroup_ACTIVE {
					if err := s.dataplaneNodeGroupRepo.Delete(ctx, dataplaneName, nodeSpec, group.GetPrimaryInstanceType()); err != nil {
						return status.Errorf(codes.Internal, "failed to delete redundant dataplane override: %v", err)
					}
					continue
				}
				planName := dag.PlanCreateNodeInstanceGroup
				execIDPrefix := "set-nig-active"
				if oldGroupStatus == apiv1.NodeInstanceGroup_DEPRECATED {
					planName = dag.PlanReactivateNodeInstanceGroup
					execIDPrefix = "set-nig-reactivate"
				}
				placementGroupEnabled, err := s.placementGroupEnabledForDataplane(ctx, dataplaneName)
				if err != nil {
					return status.Errorf(codes.Internal, "failed to resolve placement group eligibility: %v", err)
				}
				args, err := buildNodeInstanceGroupDagArgs(ctx, group, dataplaneName, placementGroupEnabled)
				if err != nil {
					return status.Errorf(codes.Internal, "failed to build dag args: %v", err)
				}
				execID := utils.GeneratedDagExecId(ctx, fmt.Sprintf("%s-%s-%s-%s", execIDPrefix, truncateForExecId(dataplaneName), strings.ToLower(nodeSpec), normalizeInstanceType(group.GetPrimaryInstanceType())))
				if err := s.dagExecutor.Submit(ctx, planName, args, execID); err != nil {
					return status.Errorf(codes.Internal, "failed to submit node instance group dag: %v", err)
				}

			case apiv1.NodeInstanceGroup_DEPRECATED:
				// Only deprecate on dataplanes where the group was effectively active before this transition.
				// This prevents creating spurious state records on dataplanes that never had the node pool
				// (e.g., when transitioning from PENDING → DEPRECATED).
				if !isNodeGroupEffectivelyActive(oldGroupStatus, override) {
					continue
				}
				// Skip dataplanes with ACTIVE overlay — the override takes precedence over region-level status.
				if override != nil && override.GetStatus() == apiv1.DataplaneNodeInstanceGroup_ACTIVE {
					continue
				}
				// DEPRECATED plans run the DeprecateNodeInstanceGroupAction which only
				// uses the node instance group data (not the placement_group_enabled
				// field), so the value is unused. Pass the default false and skip the
				// shared-pool repo round-trip entirely.
				args, err := buildNodeInstanceGroupDagArgs(ctx, group, dataplaneName, false)
				if err != nil {
					return status.Errorf(codes.Internal, "failed to build dag args: %v", err)
				}
				execID := utils.GeneratedDagExecId(ctx, fmt.Sprintf("set-nig-deprecated-%s-%s-%s", truncateForExecId(dataplaneName), strings.ToLower(nodeSpec), normalizeInstanceType(group.GetPrimaryInstanceType())))
				if err := s.dagExecutor.Submit(ctx, dag.PlanDeprecateNodeInstanceGroup, args, execID); err != nil {
					return status.Errorf(codes.Internal, "failed to submit deprecate node instance group dag: %v", err)
				}
				// Best-effort cleanup of shortage detector state for this dataplane.
				// Failure is non-fatal: the shortage detector's isNodeInstanceGroupEnabled check
				// will skip DEPRECATED groups, and the RETIRING bulk-clean handles stragglers.
				if err := s.nodeGroupStateRepo.DeleteByDataplaneAndSpec(ctx, dataplaneName, nodeSpec, group.GetPrimaryInstanceType()); err != nil {
					tilog.Warn(ctx, "failed to delete dataplane state for deprecated group (best-effort)",
						zap.String("dataplaneName", dataplaneName),
						zap.String("nodeSpec", nodeSpec),
						zap.String("primaryInstance", group.GetPrimaryInstanceType()),
						zap.Error(err))
				}

			case apiv1.NodeInstanceGroup_RETIRING:
				// Ensure every dataplane has an explicit RETIRING override so the retiring
				// processor can track per-dataplane progress (it checks override status to
				// skip already-RETIRED dataplanes).
				if override == nil {
					override = &entity.DataplaneNodeInstanceGroup{
						DataplaneName:       dataplaneName,
						NodeSpec:            group.GetNodeSpec(),
						PrimaryInstanceType: group.GetPrimaryInstanceType(),
					}
				}
				override.Status = apiv1.DataplaneNodeInstanceGroup_RETIRING
				if err := s.dataplaneNodeGroupRepo.Save(ctx, override); err != nil {
					return status.Errorf(codes.Internal, "failed to update dataplane override to RETIRING: %v", err)
				}
			}
		}

		// Best-effort bulk-clean all shortage detector states when entering RETIRING.
		// At this point every dataplane override is RETIRING (all were DEPRECATED before,
		// and the DEPRECATED case above already cleaned individual dataplanes). This catch-all
		// handles any stragglers (e.g. dataplanes that had no override or were added between transitions).
		if nodeGroupStatus == apiv1.NodeInstanceGroup_RETIRING {
			if err := s.nodeGroupStateRepo.DeleteBySpecAndPrimary(ctx, nodeSpec, primaryInstance); err != nil {
				tilog.Warn(ctx, "failed to bulk-delete dataplane states for retiring group (best-effort)",
					zap.String("nodeSpec", nodeSpec),
					zap.String("primaryInstance", primaryInstance),
					zap.Error(err))
			}
		}

		return nil
	}); err != nil {
		return nil, err
	}

	return updated, nil
}

func (s *serverImpl) ListDataplaneNodeInstanceGroups(ctx context.Context, dataplaneName string) ([]*entity.DataplaneNodeInstanceGroup, error) {
	if dataplaneName == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dataplane_name is required")
	}

	// Load all region-level groups and all overlays for this dataplane.
	allGroups, err := s.nodeGroupRepo.ListAll(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list node instance groups: %v", err)
	}
	overrides, err := s.dataplaneNodeGroupRepo.ListByDataplane(ctx, dataplaneName)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list dataplane overrides: %v", err)
	}

	// Build override lookup: key = "nodeSpec/primaryInstance"
	overrideMap := make(map[string]*entity.DataplaneNodeInstanceGroup, len(overrides))
	for _, o := range overrides {
		spec := entity.GetComponentNodeSpec(o.GetNodeSpec())
		key := spec + "/" + o.GetPrimaryInstanceType()
		overrideMap[key] = o
	}

	// Merge: for each region-level group, compute the effective dataplane status.
	results := make([]*entity.DataplaneNodeInstanceGroup, 0, len(allGroups))
	for _, group := range allGroups {
		spec := entity.GetComponentNodeSpec(group.GetNodeSpec())
		key := spec + "/" + group.GetPrimaryInstanceType()

		if override, ok := overrideMap[key]; ok {
			// Overlay exists — use override status directly.
			results = append(results, override)
			delete(overrideMap, key)
		} else {
			// No overlay — derive effective status from region-level status.
			var effectiveStatus apiv1.DataplaneNodeInstanceGroup_Status
			switch group.GetStatus() {
			case apiv1.NodeInstanceGroup_ACTIVE:
				effectiveStatus = apiv1.DataplaneNodeInstanceGroup_ACTIVE
			case apiv1.NodeInstanceGroup_DEPRECATED:
				effectiveStatus = apiv1.DataplaneNodeInstanceGroup_DEPRECATED
			case apiv1.NodeInstanceGroup_RETIRING:
				effectiveStatus = apiv1.DataplaneNodeInstanceGroup_RETIRING
			case apiv1.NodeInstanceGroup_RETIRED:
				effectiveStatus = apiv1.DataplaneNodeInstanceGroup_RETIRED
			default:
				// PENDING groups without overlay are not effective on any dataplane.
				continue
			}
			results = append(results, &entity.DataplaneNodeInstanceGroup{
				DataplaneName:       dataplaneName,
				NodeSpec:            group.GetNodeSpec(),
				PrimaryInstanceType: group.GetPrimaryInstanceType(),
				Status:              effectiveStatus,
			})
		}
	}

	// Include any remaining overrides that don't match a current region-level group.
	for _, o := range overrideMap {
		results = append(results, o)
	}

	return results, nil
}

func (s *serverImpl) SetDataplaneNodeInstanceGroupStatus(ctx context.Context, dataplaneName, nodeSpec, primaryInstance string, overrideStatus apiv1.DataplaneNodeInstanceGroup_Status) (*entity.DataplaneNodeInstanceGroup, error) {
	if dataplaneName == "" || nodeSpec == "" || primaryInstance == "" {
		return nil, status.Errorf(codes.InvalidArgument, "dataplane_name, node_spec, and primary_instance are required")
	}
	switch overrideStatus {
	case apiv1.DataplaneNodeInstanceGroup_ACTIVE, apiv1.DataplaneNodeInstanceGroup_DEPRECATED:
	default:
		return nil, status.Errorf(codes.InvalidArgument, "invalid status: %v, only ACTIVE and DEPRECATED are allowed", overrideStatus)
	}

	// Check dataplane status - cannot set node instance group status if dataplane is deleting
	isDeleting, err := s.isDataplaneDeleting(ctx, dataplaneName)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to check dataplane status: %v", err)
	}
	if isDeleting {
		return nil, status.Errorf(codes.FailedPrecondition, "dataplane %s is being deleted", dataplaneName)
	}

	group, err := s.nodeGroupRepo.LoadBySpecAndPrimary(ctx, nodeSpec, primaryInstance)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to load node instance group: %v", err)
	}
	if group == nil {
		return nil, status.Errorf(codes.NotFound, "node instance group not found: %s", nodeSpec)
	}

	// Forbid any manual override changes when the region-level group is RETIRING or RETIRED.
	// ACTIVE would reintroduce workloads after retirement; DEPRECATED would persist because
	// the retiring processor only scans region groups in RETIRING state and would never
	// reconcile this dataplane back to RETIRED automatically.
	if group.GetStatus() == apiv1.NodeInstanceGroup_RETIRING || group.GetStatus() == apiv1.NodeInstanceGroup_RETIRED {
		return nil, status.Errorf(codes.FailedPrecondition,
			"cannot change dataplane override: region-level group is %v, manual overrides are forbidden during retirement",
			group.GetStatus())
	}

	override, err := s.dataplaneNodeGroupRepo.LoadByDataplaneAndSpec(ctx, dataplaneName, nodeSpec, primaryInstance)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to load dataplane override: %v", err)
	}
	oldOverrideStatus := apiv1.DataplaneNodeInstanceGroup_STATUS_UNSPECIFIED
	if override == nil {
		cpu, memGi, err := nodegroup.ParseSpec(nodeSpec)
		if err != nil {
			return nil, status.Errorf(codes.InvalidArgument, "invalid node_spec: %v", err)
		}
		override = &entity.DataplaneNodeInstanceGroup{
			DataplaneName:       dataplaneName,
			NodeSpec:            &apiv1.NodeSpec{Cpu: cpu, MemoryGi: memGi},
			PrimaryInstanceType: primaryInstance,
		}
	} else {
		oldOverrideStatus = override.GetStatus()
	}
	override.Status = overrideStatus
	if err := s.dataplaneNodeGroupRepo.Save(ctx, override); err != nil {
		return nil, status.Errorf(codes.Internal, "failed to save dataplane override: %v", err)
	}

	switch overrideStatus {
	case apiv1.DataplaneNodeInstanceGroup_ACTIVE:
		planName := dag.PlanCreateNodeInstanceGroup
		execIDPrefix := "set-dp-nig-active"
		if oldOverrideStatus == apiv1.DataplaneNodeInstanceGroup_DEPRECATED || group.GetStatus() == apiv1.NodeInstanceGroup_DEPRECATED {
			planName = dag.PlanReactivateNodeInstanceGroup
			execIDPrefix = "set-dp-nig-reactivate"
		}
		placementGroupEnabled, err := s.placementGroupEnabledForDataplane(ctx, dataplaneName)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to resolve placement group eligibility: %v", err)
		}
		args, err := buildNodeInstanceGroupDagArgs(ctx, group, dataplaneName, placementGroupEnabled)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to build dag args: %v", err)
		}
		execID := utils.GeneratedDagExecId(ctx, fmt.Sprintf("%s-%s-%s-%s", execIDPrefix, truncateForExecId(dataplaneName), strings.ToLower(nodeSpec), normalizeInstanceType(group.GetPrimaryInstanceType())))
		if err := s.dagExecutor.Submit(ctx, planName, args, execID); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to submit node instance group dag: %v", err)
		}
	case apiv1.DataplaneNodeInstanceGroup_DEPRECATED:
		// DeprecateNodeInstanceGroupAction does not consume
		// placement_group_enabled, so the value is unused. Pass the
		// default false and skip the shared-pool repo round-trip.
		args, err := buildNodeInstanceGroupDagArgs(ctx, group, dataplaneName, false)
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to build dag args: %v", err)
		}
		execID := utils.GeneratedDagExecId(ctx, fmt.Sprintf("set-dp-nig-deprecated-%s-%s-%s", truncateForExecId(dataplaneName), strings.ToLower(nodeSpec), normalizeInstanceType(group.GetPrimaryInstanceType())))
		if err := s.dagExecutor.Submit(ctx, dag.PlanDeprecateNodeInstanceGroup, args, execID); err != nil {
			return nil, status.Errorf(codes.Internal, "failed to submit deprecate node instance group dag: %v", err)
		}
		// Best-effort cleanup of shortage detector state for this dataplane.
		// The override and DAG have already been committed, so returning an error here
		// would cause callers to retry and submit duplicate DAGs.
		if err := s.nodeGroupStateRepo.DeleteByDataplaneAndSpec(ctx, dataplaneName, nodeSpec, primaryInstance); err != nil {
			tilog.Warn(ctx, "failed to delete dataplane state for deprecated group (best-effort)",
				zap.String("dataplaneName", dataplaneName),
				zap.String("nodeSpec", nodeSpec),
				zap.String("primaryInstance", primaryInstance),
				zap.Error(err))
		}
	}

	return override, nil
}

func (s *serverImpl) SetNodeInstanceGroupInstanceTypeStatus(ctx context.Context, nodeSpec, primaryInstance, instanceType string, instanceStatus apiv1.InstanceTypeStatus) (*entity.NodeInstanceGroup, error) {
	if nodeSpec == "" || primaryInstance == "" || instanceType == "" {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec, primary_instance, and instance_type are required")
	}
	switch instanceStatus {
	case apiv1.InstanceTypeStatus_NORMAL, apiv1.InstanceTypeStatus_DISABLED:
	default:
		return nil, status.Errorf(codes.InvalidArgument, "invalid instance type status: %v", instanceStatus)
	}
	if err := nodegroup.ValidateCloudInstanceType(instanceType); err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid instance_type: %v", err)
	}

	// Fetch non-deleting dataplanes before the transaction to avoid holding DB locks during external gRPC calls.
	// Deleting dataplanes are filtered out at the list level.
	dataplaneNames, err := s.listNonDeletingDataplaneNames(ctx)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list dataplanes: %v", err)
	}

	var updated *entity.NodeInstanceGroup
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		group, err := s.nodeGroupRepo.LoadBySpecAndPrimary(ctx, nodeSpec, primaryInstance)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to load node instance group: %v", err)
		}
		if group == nil {
			return status.Errorf(codes.NotFound, "node instance group not found: %s/%s", nodeSpec, primaryInstance)
		}

		if instanceType != group.GetPrimaryInstanceType() {
			found := false
			for _, fallback := range group.GetFallbackInstances() {
				if fallback != nil && fallback.GetInstanceType() == instanceType {
					found = true
					break
				}
			}
			if !found {
				return status.Errorf(codes.InvalidArgument, "instance_type %s not found in node_spec %s", instanceType, nodeSpec)
			}
		}

		disabled := append([]string(nil), group.GetDisabledInstanceTypes()...)
		if instanceStatus == apiv1.InstanceTypeStatus_DISABLED {
			disabled = addUniqueString(disabled, instanceType)
		} else {
			disabled = removeString(disabled, instanceType)
		}
		group.DisabledInstanceTypes = disabled
		if err := s.nodeGroupRepo.Save(ctx, group); err != nil {
			return status.Errorf(codes.Internal, "failed to update node instance group: %v", err)
		}
		updated = group

		instanceStatusValue := instanceTypeStatusString(instanceStatus)
		for _, dataplaneName := range dataplaneNames {
			override, err := s.dataplaneNodeGroupRepo.LoadByDataplaneAndSpec(ctx, dataplaneName, nodeSpec, primaryInstance)
			if err != nil {
				return status.Errorf(codes.Internal, "failed to load dataplane override: %v", err)
			}
			// Skip dataplanes where the group is not effectively active.
			if !isNodeGroupEffectivelyActive(group.GetStatus(), override) {
				continue
			}
			// SetNodeInstanceGroupInstanceTypeStatusAction only taints
			// the affected instance types; it does not create NodePools
			// and never reads placement_group_enabled. Pass the literal
			// false and skip the SharedPool repo round-trip entirely.
			args, err := buildNodeInstanceGroupDagArgs(ctx, group, dataplaneName, false)
			if err != nil {
				return status.Errorf(codes.Internal, "failed to build dag args: %v", err)
			}
			args[workflow.InstanceTypeKey] = actions.ActionArgString(instanceType)
			args[workflow.InstanceTypeStatusKey] = actions.ActionArgString(instanceStatusValue)
			execID := utils.GeneratedDagExecId(ctx, fmt.Sprintf("set-nig-type-status-%s-%s-%s-%s", truncateForExecId(dataplaneName), strings.ToLower(nodeSpec), normalizeInstanceType(primaryInstance), normalizeInstanceType(instanceType)))
			if err := s.dagExecutor.Submit(ctx, dag.PlanSetNodeInstanceGroupInstanceTypeStatus, args, execID); err != nil {
				return status.Errorf(codes.Internal, "failed to submit set instance type status dag: %v", err)
			}
		}
		return nil
	}); err != nil {
		return nil, err
	}

	return updated, nil
}

func (s *serverImpl) DeleteNodeInstanceGroup(ctx context.Context, nodeSpec, primaryInstance string) (*entity.NodeInstanceGroup, error) {
	if nodeSpec == "" || primaryInstance == "" {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec and primary_instance are required")
	}

	var deleted *entity.NodeInstanceGroup
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		target, err := s.nodeGroupRepo.LoadBySpecAndPrimary(ctx, nodeSpec, primaryInstance)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to load node instance group: %v", err)
		}
		if target == nil {
			return status.Errorf(codes.NotFound, "node instance group not found: %s/%s", nodeSpec, primaryInstance)
		}

		// Only RETIRED groups can be deleted — RETIRED guarantees all nodepools have already been
		// cleaned up by the retiring processor cronjob.
		if target.GetStatus() != apiv1.NodeInstanceGroup_RETIRED {
			return status.Errorf(codes.FailedPrecondition, "node instance group is not retired: %s/%s (status=%v), only RETIRED groups can be deleted", nodeSpec, primaryInstance, target.GetStatus())
		}

		overrides, err := s.dataplaneNodeGroupRepo.ListBySpec(ctx, nodeSpec)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to list dataplane overrides: %v", err)
		}

		if err := s.nodeGroupRepo.Delete(ctx, nodeSpec, primaryInstance); err != nil {
			return status.Errorf(codes.Internal, "failed to delete node instance group: %v", err)
		}
		for _, override := range overrides {
			if override.GetPrimaryInstanceType() != primaryInstance {
				continue
			}
			if err := s.dataplaneNodeGroupRepo.Delete(ctx, override.GetDataplaneName(), nodeSpec, primaryInstance); err != nil {
				return status.Errorf(codes.Internal, "failed to delete dataplane override: %v", err)
			}
		}

		// Safety-net: clean up any remaining shortage detector state rows.
		// States should already be removed during DEPRECATED/RETIRING transitions, but
		// this handles edge cases (e.g., states created between transitions).
		if err := s.nodeGroupStateRepo.DeleteBySpecAndPrimary(ctx, nodeSpec, primaryInstance); err != nil {
			return status.Errorf(codes.Internal, "failed to delete dataplane state rows: %v", err)
		}

		deleted = target
		return nil
	}); err != nil {
		return nil, err
	}

	return deleted, nil
}

// UpdateNodeInstanceGroup updates the fallback instances for a node instance group.
// Only fallback_instances can be modified; the primary instance type and node spec are immutable.
// This is a region-level metadata update that does not require per-dataplane propagation.
func (s *serverImpl) UpdateNodeInstanceGroup(ctx context.Context, nodeSpec, primaryInstance string, fallbackInstances []*apiv1.NodeInstanceGroup_FallbackInstance) (*entity.NodeInstanceGroup, error) {
	if nodeSpec == "" || primaryInstance == "" {
		return nil, status.Errorf(codes.InvalidArgument, "node_spec and primary_instance are required")
	}

	if err := validateFallbackInstances(fallbackInstances, nil); err != nil {
		return nil, status.Errorf(codes.InvalidArgument, "invalid fallback_instances: %v", err)
	}
	for _, fb := range fallbackInstances {
		if fb.GetInstanceType() == primaryInstance {
			return nil, status.Errorf(codes.InvalidArgument, "fallback instance_type cannot be the same as primary: %s", primaryInstance)
		}
	}

	var updated *entity.NodeInstanceGroup
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		group, err := s.nodeGroupRepo.LoadBySpecAndPrimary(ctx, nodeSpec, primaryInstance)
		if err != nil {
			return status.Errorf(codes.Internal, "failed to load node instance group: %v", err)
		}
		if group == nil {
			return status.Errorf(codes.NotFound, "node instance group not found: %s/%s", nodeSpec, primaryInstance)
		}

		// Re-validate with primary spec now that we have the group.
		if err := validateFallbackInstances(fallbackInstances, group.GetNodeSpec()); err != nil {
			return status.Errorf(codes.InvalidArgument, "invalid fallback_instances: %v", err)
		}

		if group.GetNodeSpec().GetWithLocalDisk() && len(fallbackInstances) > 0 {
			return status.Errorf(codes.InvalidArgument, "fallback_instances is not allowed when with_local_disk is true")
		}

		// Prevent removal of existing fallback instances to avoid orphaned nodepools.
		if err := validateFallbackNotRemoved(group.GetFallbackInstances(), fallbackInstances); err != nil {
			return err
		}

		group.FallbackInstances = fallbackInstances
		if err := s.nodeGroupRepo.Save(ctx, group); err != nil {
			return status.Errorf(codes.Internal, "failed to save node instance group: %v", err)
		}
		updated = group
		return nil
	}); err != nil {
		return nil, err
	}

	tilog.Info(ctx, "updated node instance group fallback instances",
		zap.String("nodeSpec", nodeSpec),
		zap.String("primaryInstance", primaryInstance),
		zap.Int("fallbackCount", len(fallbackInstances)))

	return updated, nil
}

// normalizeInstanceType converts an instance type like "r8g.xlarge" to "r8g-xlarge" for use in task names.
func normalizeInstanceType(instanceType string) string {
	return strings.ReplaceAll(strings.ToLower(strings.TrimSpace(instanceType)), ".", "-")
}

// truncateForExecId truncates a dataplane name for use in exec ID generation.
// It removes the "dp-" prefix (if present) to preserve more meaningful characters,
// then truncates to fit within the 127 character limit for K8s resource names.
const execIdDataplaneMaxLen = 20

func truncateForExecId(s string) string {
	// Remove "dp-" prefix to preserve more meaningful characters
	s = strings.TrimPrefix(s, "dp-")
	if len(s) <= execIdDataplaneMaxLen {
		return s
	}
	return s[:execIdDataplaneMaxLen]
}

func validateFallbackInstances(values []*entity.NodeInstanceGroupFallbackInstance, primaryNodeSpec *apiv1.NodeSpec) error {
	seen := make(map[string]struct{}, len(values))
	for _, value := range values {
		if value == nil {
			continue
		}
		if err := nodegroup.ValidateCloudInstanceType(value.GetInstanceType()); err != nil {
			return err
		}
		if value.GetNodeSpec() == nil {
			return fmt.Errorf("fallback node_spec is nil")
		}
		fallbackSpec := entity.GetComponentNodeSpec(value.GetNodeSpec())
		if err := nodegroup.ValidateSpec(fallbackSpec); err != nil {
			return err
		}
		// Fallback instance must provide at least as much capacity as primary.
		// Fallback nodepools use primary spec for resource advertisement, so a
		// smaller fallback would overcommit the underlying machine.
		if primaryNodeSpec != nil {
			if value.GetNodeSpec().GetCpu() < primaryNodeSpec.GetCpu() || value.GetNodeSpec().GetMemoryGi() < primaryNodeSpec.GetMemoryGi() {
				return fmt.Errorf("fallback instance %s spec (%dC%dG) must be >= primary spec (%dC%dG)",
					value.GetInstanceType(),
					value.GetNodeSpec().GetCpu(), value.GetNodeSpec().GetMemoryGi(),
					primaryNodeSpec.GetCpu(), primaryNodeSpec.GetMemoryGi())
			}
		}
		if value.GetPriority() == 0 {
			return fmt.Errorf("fallback priority is required for instance %s", value.GetInstanceType())
		}
		if value.GetPriority() >= nodegroup.PriorityPrimary {
			return fmt.Errorf("fallback priority must be less than primary (%d): got %d", nodegroup.PriorityPrimary, value.GetPriority())
		}
		key := fmt.Sprintf("%s/%s", value.GetInstanceType(), fallbackSpec)
		if _, ok := seen[key]; ok {
			return fmt.Errorf("duplicated fallback instance: %s", value.GetInstanceType())
		}
		seen[key] = struct{}{}
	}
	return nil
}

// validateFallbackNotRemoved ensures that no existing fallback instance is removed during an update.
// Removing a fallback would leave orphaned nodepools in dataplanes.
func validateFallbackNotRemoved(existing, updated []*entity.NodeInstanceGroupFallbackInstance) error {
	if len(existing) == 0 {
		return nil
	}
	updatedKeys := make(map[string]struct{}, len(updated))
	for _, item := range updated {
		if item == nil || item.GetNodeSpec() == nil {
			continue
		}
		key := fmt.Sprintf("%s/%s", item.GetInstanceType(), entity.GetComponentNodeSpec(item.GetNodeSpec()))
		updatedKeys[key] = struct{}{}
	}
	for _, item := range existing {
		if item == nil || item.GetNodeSpec() == nil {
			continue
		}
		key := fmt.Sprintf("%s/%s", item.GetInstanceType(), entity.GetComponentNodeSpec(item.GetNodeSpec()))
		if _, ok := updatedKeys[key]; !ok {
			return status.Errorf(codes.InvalidArgument, "fallback_instances cannot remove existing instance: %s", item.GetInstanceType())
		}
	}
	return nil
}

func addUniqueString(values []string, value string) []string {
	if slices.Contains(values, value) {
		return values
	}
	return append(values, value)
}

func removeString(values []string, value string) []string {
	if len(values) == 0 {
		return values
	}
	result := make([]string, 0, len(values))
	for _, existing := range values {
		if existing == value {
			continue
		}
		result = append(result, existing)
	}
	return result
}

type nodePoolRef struct {
	dataPlaneID string
	nodePool    *dplanemgrv1.NodePool
}

func (s *serverImpl) listNodePoolsFromDataplanes(ctx context.Context, dataplaneNames []string) ([]nodePoolRef, error) {
	var results []nodePoolRef
	for _, dataplaneName := range dataplaneNames {
		resp, err := s.dataplaneClient.ListNodePools(ctx, &dplanemgrv1.ListNodePoolsRequest{
			DataPlaneId: dataplaneName,
		})
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to list nodepools for dataplane %s: %v", dataplaneName, err)
		}
		for _, nodePool := range resp.GetNodePools() {
			results = append(results, nodePoolRef{
				dataPlaneID: dataplaneName,
				nodePool:    nodePool,
			})
		}
	}
	return results, nil
}

func (s *serverImpl) listNonDeletingDataplaneNames(ctx context.Context) ([]string, error) {
	resp, err := s.dataplaneClient.ListDataPlanes(ctx, &dplanemgrv1.ListDataPlanesRequest{})
	if err != nil {
		return nil, err
	}
	names := make([]string, 0, len(resp.GetDataPlanes()))
	for _, dataplane := range resp.GetDataPlanes() {
		if dataplane == nil || dataplane.GetId() == "" {
			continue
		}
		if dataplane.GetStatus() == dplanemgrv1.DataPlane_Deleting {
			tilog.Info(ctx, "skipping deleting dataplane",
				zap.String("dataplaneName", dataplane.GetId()))
			continue
		}
		names = append(names, dataplane.GetId())
	}
	sort.Strings(names)
	return names, nil
}

// isDataplaneDeleting checks if the dataplane is in Deleting status.
func (s *serverImpl) isDataplaneDeleting(ctx context.Context, dataplaneName string) (bool, error) {
	dp, err := s.dataplaneClient.GetDataPlane(ctx, &dplanemgrv1.GetDataPlaneRequest{DataPlaneId: dataplaneName})
	if err != nil {
		return false, err
	}
	return dp.GetStatus() == dplanemgrv1.DataPlane_Deleting, nil
}

func (s *serverImpl) ensureNoActiveNodePoolsForGroup(ctx context.Context, group *entity.NodeInstanceGroup, dataplaneNames []string) error {
	if group == nil {
		return status.Errorf(codes.InvalidArgument, "node instance group is nil")
	}
	if group.GetNodeSpec() == nil {
		return status.Errorf(codes.InvalidArgument, "node_spec is nil")
	}
	groupSpec := entity.GetComponentNodeSpec(group.GetNodeSpec())
	// Parse groupSpec to get CPU and Memory for matching dataplane default pools
	cpu, mem, err := nodegroup.ParseSpec(groupSpec)
	if err != nil {
		return err
	}
	nodePools, err := s.listNodePoolsFromDataplanes(ctx, dataplaneNames)
	if err != nil {
		return err
	}
	instanceTypes := nodeInstanceGroupInstanceTypes(group)
	for _, ref := range nodePools {
		labels := ref.nodePool.GetLabels()
		isMatch := false

		if nodegroup.IsManagedByInstanceTypeManager(labels) {
			// Case 1: InstanceTypeManager managed pools - match by node-spec label
			if labels != nil && labels[nodegroup.LabelNodeSpec] == groupSpec {
				isMatch = true
			}
		} else {
			// Case 2: Dataplane default pools - match by CPU/Memory
			if ref.nodePool.GetCpuCore() == cpu && ref.nodePool.GetMemoryGi() == mem {
				isMatch = true
			}
		}

		if !isMatch {
			continue
		}
		if !shouldCheckNodePoolForGroup(group, instanceTypes, ref.nodePool) {
			continue
		}
		if ref.nodePool.GetCurrentReplicas() > 0 {
			return status.Errorf(codes.FailedPrecondition, "nodepool still has replicas for node_spec %s primary_instance %s", entity.GetComponentNodeSpec(group.GetNodeSpec()), group.GetPrimaryInstanceType())
		}
		if !isDeprecatedNodePool(ref.nodePool) {
			return status.Errorf(codes.FailedPrecondition, "nodepool already exists for node_spec %s primary_instance %s", entity.GetComponentNodeSpec(group.GetNodeSpec()), group.GetPrimaryInstanceType())
		}
	}
	return nil
}

func shouldCheckNodePoolForGroup(group *entity.NodeInstanceGroup, instanceTypes map[string]nodegroup.NodePoolRole, nodePool *dplanemgrv1.NodePool) bool {
	if nodePool == nil {
		return false
	}
	instanceType := nodePool.GetInstanceType()
	if instanceType == group.GetPrimaryInstanceType() {
		return nodegroup.NodePoolMatchesRole(nodePool, nodegroup.NodePoolRolePrimary)
	}
	expectedRole, ok := instanceTypes[instanceType]
	if !ok {
		return false
	}
	return nodegroup.NodePoolMatchesRole(nodePool, expectedRole)
}

func nodeInstanceGroupInstanceTypes(group *entity.NodeInstanceGroup) map[string]nodegroup.NodePoolRole {
	types := map[string]nodegroup.NodePoolRole{
		group.GetPrimaryInstanceType(): nodegroup.NodePoolRolePrimary,
	}
	for _, fallback := range group.GetFallbackInstances() {
		if fallback == nil || fallback.GetInstanceType() == "" {
			continue
		}
		if _, exists := types[fallback.GetInstanceType()]; exists {
			continue
		}
		types[fallback.GetInstanceType()] = nodegroup.NodePoolRoleFallback
	}
	return types
}

func isDeprecatedNodePool(nodePool *dplanemgrv1.NodePool) bool {
	if nodePool == nil {
		return false
	}
	deprecated := nodegroup.DeprecatedTaint()
	for _, taint := range nodePool.GetTaints() {
		if taint == nil {
			continue
		}
		if taint.GetKey() == deprecated.Key && taint.GetValue() == deprecated.Value {
			return true
		}
	}
	return false
}

// placementGroupEnabledForDataplanes returns the per-dataplane placement
// group eligibility used by the NIG creation DAG. The value is the durable
// per-pool migration gate stored on the shared pool as
// Provider.PlacementGroupEnabled (set at CreateSharedPool time from the
// workload's node-zone count, "AWS + single-AZ?"). It is the same source of
// truth shared with the NodePool infrastructure side (DPM gRPC and EKS
// controller), the component templates, and the shortage-detector fallback.
// Re-deriving here would silently flip legacy single-AZ/PG-disabled pools.
//
// The returned map always has one entry per requested dataplane. Dataplanes
// with no shared pool silently default to false (proto spec default). The
// loop only writes entries for dataplanes in the requested set, so a repo
// that ignores the filter cannot widen the result map.
func (s *serverImpl) placementGroupEnabledForDataplanes(ctx context.Context, dataplaneNames []string) (map[string]bool, error) {
	out := make(map[string]bool, len(dataplaneNames))
	if s.sharedPoolRepo == nil {
		return nil, status.Error(codes.Internal, "shared pool repository is not configured")
	}
	if len(dataplaneNames) == 0 {
		return out, nil
	}
	for _, name := range dataplaneNames {
		out[name] = false
	}
	sharedPools, err := s.sharedPoolRepo.List(ctx, &repo.SharedPoolFilter{DataplaneNames: dataplaneNames})
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to list shared pools for placement group eligibility: %v", err)
	}
	for _, sharedPool := range sharedPools {
		if sharedPool == nil || sharedPool.GetProvider() == nil {
			continue
		}
		name := sharedPool.GetDataplaneName()
		if _, requested := out[name]; !requested {
			continue
		}
		out[name] = sharedPool.GetProvider().GetPlacementGroupEnabled()
	}
	return out, nil
}

// placementGroupEnabledForDataplane resolves the placement-group
// eligibility for a single dataplane. It is the per-call shape used by the
// NIG action's DAG submission, where the active dataplane is known at the
// point the value is needed and only that one lookup is required.
func (s *serverImpl) placementGroupEnabledForDataplane(ctx context.Context, dataplaneName string) (bool, error) {
	if s.sharedPoolRepo == nil {
		return false, status.Error(codes.Internal, "shared pool repository is not configured")
	}
	sharedPools, err := s.sharedPoolRepo.List(ctx, &repo.SharedPoolFilter{DataplaneNames: []string{dataplaneName}})
	if err != nil {
		return false, status.Errorf(codes.Internal, "failed to list shared pools for placement group eligibility: %v", err)
	}
	for _, sp := range sharedPools {
		if sp == nil || sp.GetProvider() == nil {
			continue
		}
		if sp.GetDataplaneName() == dataplaneName {
			return sp.GetProvider().GetPlacementGroupEnabled(), nil
		}
	}
	return false, nil
}

func defaultAvailableCPU(cpu uint32) uint32 {
	if cpu <= 1 {
		return 0
	}
	return cpu - 1
}

func buildNodeInstanceGroupDagArgs(ctx context.Context, group *entity.NodeInstanceGroup, dataplaneName string, placementGroupEnabled bool) (actions.ActionArgs, error) {
	groupData, err := protojson.Marshal(group)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal node instance group: %w", err)
	}
	args := actions.ActionArgs{
		workflow.DataplaneIdKey:           actions.ActionArgString(dataplaneName),
		workflow.NodeInstanceGroupKey:     actions.ActionArgString(string(groupData)),
		workflow.PlacementGroupEnabledKey: actions.ActionArgBool(placementGroupEnabled),
		utils.TraceIdKey:                  utils.GetTraceId(ctx),
	}
	return args, nil
}

func instanceTypeStatusString(value apiv1.InstanceTypeStatus) string {
	switch value {
	case apiv1.InstanceTypeStatus_DISABLED:
		return nodegroup.NodeInstanceGroupInstanceTypeStatusDisabled
	default:
		return nodegroup.NodeInstanceGroupInstanceTypeStatusNormal
	}
}

// isNodeGroupEffectivelyActive checks if a NodeInstanceGroup is effectively active
// on a specific dataplane, considering the overlay resolution semantics.
// Returns true only when the resolved effective status is ACTIVE.
func isNodeGroupEffectivelyActive(groupStatus apiv1.NodeInstanceGroup_Status, override *entity.DataplaneNodeInstanceGroup) bool {
	if override != nil {
		return override.GetStatus() == apiv1.DataplaneNodeInstanceGroup_ACTIVE
	}
	return groupStatus == apiv1.NodeInstanceGroup_ACTIVE
}