package nodegroup

import "sort"

// Taint keys and values for node instance group state management.
// All taints use a consistent namespace: tidbcloud.com/nodegroup-state
const (
	// TaintKeyNodeGroupState is the taint key for node instance group state
	TaintKeyNodeGroupState = "tidbcloud.com/nodegroup-state"

	// TaintValueDeprecated is used for manually deprecated node instance groups
	TaintValueDeprecated = "deprecated"

	// TaintValueShortageFallback is used when fallback is being drained
	TaintValueShortageFallback = "shortage-fallback"

	// TaintValueDisabled is used when an instance type is marked as disabled
	TaintValueDisabled = "disabled"

	// TaintEffectNoSchedule prevents new pods from being scheduled
	TaintEffectNoSchedule = "NoSchedule"

	// TaintEffectPreferNoSchedule is a soft preference to avoid scheduling
	TaintEffectPreferNoSchedule = "PreferNoSchedule"
)

// Taint represents a Kubernetes taint.
type Taint struct {
	Key    string
	Value  string
	Effect string
}

// DeprecatedTaint creates a NoSchedule taint for deprecated node instance groups.
func DeprecatedTaint() Taint {
	return Taint{
		Key:    TaintKeyNodeGroupState,
		Value:  TaintValueDeprecated,
		Effect: TaintEffectNoSchedule,
	}
}

// ShortageFallbackTaint creates a NoSchedule taint for fallback node instance groups being drained.
func ShortageFallbackTaint() Taint {
	return Taint{
		Key:    TaintKeyNodeGroupState,
		Value:  TaintValueShortageFallback,
		Effect: TaintEffectNoSchedule,
	}
}

// DisabledTaint creates a NoSchedule taint for disabled instance types.
func DisabledTaint() Taint {
	return Taint{
		Key:    TaintKeyNodeGroupState,
		Value:  TaintValueDisabled,
		Effect: TaintEffectNoSchedule,
	}
}

// HasTaint checks if the taints slice contains a specific taint.
func HasTaint(taints []Taint, key, value string) bool {
	for _, t := range taints {
		if t.Key == key && t.Value == value {
			return true
		}
	}
	return false
}

// HasDeprecatedTaint checks if the taints indicate a deprecated node instance group.
func HasDeprecatedTaint(taints []Taint) bool {
	return HasTaint(taints, TaintKeyNodeGroupState, TaintValueDeprecated)
}

// HasShortageFallbackTaint checks if the taints indicate a draining fallback.
func HasShortageFallbackTaint(taints []Taint) bool {
	return HasTaint(taints, TaintKeyNodeGroupState, TaintValueShortageFallback)
}

// HasDisabledTaint checks if the taints indicate a disabled instance type.
func HasDisabledTaint(taints []Taint) bool {
	return HasTaint(taints, TaintKeyNodeGroupState, TaintValueDisabled)
}

// MergeTaints merges existing taints with new taints.
// Uses Key as the dedup key: each Key can only appear once in the result.
// New taints with the same Key will override existing ones (last-write-wins).
// This matches our domain convention where each taint key (e.g., tidbcloud.com/nodegroup-state)
// has exactly one active value at a time.
// The result is sorted by Key for deterministic output.
// Returns a new slice without modifying the inputs.
func MergeTaints(existing, newTaints []Taint) []Taint {
	// Build a map of existing taints by key
	taintMap := make(map[string]Taint, len(existing)+len(newTaints))

	for _, t := range existing {
		taintMap[t.Key] = t
	}

	for _, t := range newTaints {
		taintMap[t.Key] = t
	}

	// Convert back to slice, sorted by key for deterministic output
	result := make([]Taint, 0, len(taintMap))
	for _, t := range taintMap {
		result = append(result, t)
	}
	sort.Slice(result, func(i, j int) bool {
		return result[i].Key < result[j].Key
	})

	return result
}

// RemoveTaint removes a taint with the specified key from the slice.
// Returns a new slice without modifying the input.
func RemoveTaint(taints []Taint, key string) []Taint {
	result := make([]Taint, 0, len(taints))
	for _, t := range taints {
		if t.Key != key {
			result = append(result, t)
		}
	}
	return result
}

// RemoveTaintByKeyValue removes a taint matching both key and value from the slice.
// Unlike RemoveTaint which removes all taints with a given key, this only removes the
// exact key+value pair. This is important when multiple taint values share the same key
// (e.g., tidbcloud.com/nodegroup-state can be "disabled", "deprecated", or "shortage-fallback").
// Returns a new slice without modifying the input.
func RemoveTaintByKeyValue(taints []Taint, key, value string) []Taint {
	result := make([]Taint, 0, len(taints))
	for _, t := range taints {
		if t.Key == key && t.Value == value {
			continue
		}
		result = append(result, t)
	}
	return result
}

// AddTaint adds or replaces a taint by Key in the slice.
// If a taint with the same Key already exists, it is replaced (consistent with MergeTaints semantics
// where each Key has exactly one active value at a time).
// The result is sorted by Key for deterministic output (consistent with MergeTaints).
// Returns a new slice without modifying the input.
func AddTaint(taints []Taint, taint Taint) []Taint {
	result := RemoveTaint(taints, taint.Key)
	result = append(result, taint)
	sort.Slice(result, func(i, j int) bool {
		return result[i].Key < result[j].Key
	})
	return result
}