resource.go

/*
 * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 controller_common

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
	"strconv"

	"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
	"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
	"github.com/google/go-cmp/cmp"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/record"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/apiutil"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

const (
	// NvidiaAnnotationHashKey indicates annotation name for last applied hash by the operator
	NvidiaAnnotationHashKey = "nvidia.com/last-applied-hash"
	// NvidiaAnnotationGenerationKey indicates annotation name for last applied generation by the operator
	// This is used to detect manual changes to resources
	NvidiaAnnotationGenerationKey = "nvidia.com/last-applied-generation"
)

type Reconciler interface {
	client.Client
	GetRecorder() record.EventRecorder
}

// ResourceGenerator is a function that generates a resource.
// it must return the resource, a boolean indicating if the resource should be deleted, and an error
// if the resource should be deleted, the returned resource must contain the necessary information to delete it (name and namespace)
type ResourceGenerator[T client.Object] func(ctx context.Context) (T, bool, error)

//nolint:nakedret
func SyncResource[T client.Object](ctx context.Context, r Reconciler, parentResource client.Object, generateResource ResourceGenerator[T]) (modified bool, res T, err error) {
	logs := log.FromContext(ctx)

	resource, toDelete, err := generateResource(ctx)
	if err != nil {
		return
	}
	resourceNamespace := resource.GetNamespace()
	resourceName := resource.GetName()
	resourceType := reflect.TypeOf(resource).Elem().Name()
	logs = logs.WithValues("namespace", resourceNamespace, "resourceName", resourceName, "resourceType", resourceType)

	// Retrieve the GroupVersionKind (GVK) of the desired object
	gvk, err := apiutil.GVKForObject(resource, r.Scheme())
	if err != nil {
		logs.Error(err, "Failed to get GVK for object")
		return
	}

	// Create a new instance of the object
	obj, err := r.Scheme().New(gvk)
	if err != nil {
		logs.Error(err, "Failed to create a new object for GVK")
		return
	}

	// Type assertion to ensure the object implements client.Object
	oldResource, ok := obj.(T)
	if !ok {
		return
	}

	err = r.Get(ctx, types.NamespacedName{Name: resourceName, Namespace: resourceNamespace}, oldResource)
	oldResourceIsNotFound := errors.IsNotFound(err)
	if err != nil && !oldResourceIsNotFound {
		r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, fmt.Sprintf("Get%s", resourceType), "Failed to get %s %s: %s", resourceType, resourceNamespace, err)
		logs.Error(err, "Failed to get resource.")
		return
	}
	err = nil

	if oldResourceIsNotFound {
		if toDelete {
			logs.Info("Resource not found. Nothing to do.")
			return
		}
		logs.Info("Resource not found. Creating a new one.")

		err = ctrl.SetControllerReference(parentResource, resource, r.Scheme())
		if err != nil {
			logs.Error(err, "Failed to set controller reference.")
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, "SetControllerReference", "Failed to set controller reference for %s %s: %s", resourceType, resourceNamespace, err)
			return
		}

		var hash string
		hash, err = GetSpecHash(resource)
		if err != nil {
			logs.Error(err, "Failed to get spec hash.")
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, "GetSpecHash", "Failed to get spec hash for %s %s: %s", resourceType, resourceNamespace, err)
			return
		}

		// On create, set generation to 1 (new resources start at generation 1)
		updateAnnotations(resource, hash, 1)

		r.GetRecorder().Eventf(parentResource, corev1.EventTypeNormal, fmt.Sprintf("Create%s", resourceType), "Creating a new %s %s", resourceType, resourceNamespace)
		err = r.Create(ctx, resource)
		if err != nil {
			logs.Error(err, "Failed to create Resource.")
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, fmt.Sprintf("Create%s", resourceType), "Failed to create %s %s: %s", resourceType, resourceNamespace, err)
			return
		}
		logs.Info(fmt.Sprintf("%s created.", resourceType))
		r.GetRecorder().Eventf(parentResource, corev1.EventTypeNormal, fmt.Sprintf("Create%s", resourceType), "Created %s %s", resourceType, resourceNamespace)
		modified = true
		res = resource
	} else {
		logs.Info(fmt.Sprintf("%s found.", resourceType))
		if toDelete {
			logs.Info(fmt.Sprintf("%s found. Deleting the existing one.", resourceType))
			err = r.Delete(ctx, oldResource)
			if err != nil {
				logs.Error(err, fmt.Sprintf("Failed to delete %s.", resourceType))
				r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, fmt.Sprintf("Delete%s", resourceType), "Failed to delete %s %s: %s", resourceType, resourceNamespace, err)
				return
			}
			logs.Info(fmt.Sprintf("%s deleted.", resourceType))
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeNormal, fmt.Sprintf("Delete%s", resourceType), "Deleted %s %s", resourceType, resourceNamespace)
			modified = true
			return
		}

		// Check if the Spec has changed and update if necessary
		var changeResult SpecChangeResult
		changeResult, err = GetSpecChangeResult(oldResource, resource)
		if err != nil {
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, fmt.Sprintf("CalculatePatch%s", resourceType), "Failed to calculate patch for %s %s: %s", resourceType, resourceNamespace, err)
			return false, resource, fmt.Errorf("failed to check if spec has changed: %w", err)
		}

		if !changeResult.NeedsUpdate {
			logs.Info(fmt.Sprintf("%s spec is the same. Skipping update.", resourceType))
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeNormal, fmt.Sprintf("Update%s", resourceType), "Skipping update %s %s", resourceType, resourceNamespace)
			res = oldResource
			return
		}

		// Log if manual changes were detected
		if changeResult.ManualChangeDetected {
			logs.Info(fmt.Sprintf("Manual changes detected on %s, will be overwritten", resourceType),
				"currentGeneration", oldResource.GetGeneration(),
				"lastAppliedGeneration", getAnnotation(oldResource, NvidiaAnnotationGenerationKey))
		}

		// Generate and log diff before updating
		diff, diffErr := generateSpecDiff(oldResource, resource)
		if diffErr != nil {
			logs.V(1).Info(fmt.Sprintf("Failed to generate diff for %s: %v", resourceType, diffErr))
		} else if diff != "" {
			logs.Info(fmt.Sprintf("%s spec changes detected", resourceType), "diff", diff)
		}

		// Update the spec of the current object with the desired spec
		err = CopySpec(resource, oldResource)
		if err != nil {
			logs.Error(err, fmt.Sprintf("Failed to copy spec for %s.", resourceType))
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, fmt.Sprintf("CopySpec%s", resourceType), "Failed to copy spec for %s %s: %s", resourceType, resourceNamespace, err)
			return
		}

		updateAnnotations(oldResource, *changeResult.NewHash, changeResult.NewGeneration)

		err = r.Update(ctx, oldResource)
		if err != nil {
			logs.Error(err, fmt.Sprintf("Failed to update %s.", resourceType))
			r.GetRecorder().Eventf(parentResource, corev1.EventTypeWarning, fmt.Sprintf("Update%s", resourceType), "Failed to update %s %s: %s", resourceType, resourceNamespace, err)
			return
		}
		logs.Info(fmt.Sprintf("%s updated.", resourceType))
		r.GetRecorder().Eventf(parentResource, corev1.EventTypeNormal, fmt.Sprintf("Update%s", resourceType), "Updated %s %s", resourceType, resourceNamespace)
		modified = true
		res = oldResource
	}
	return
}

// CopySpec copies only the Spec field from source to destination using Unstructured
func CopySpec(source, destination client.Object) error {
	// Convert source to unstructured
	sourceMap, err := runtime.DefaultUnstructuredConverter.ToUnstructured(source)
	if err != nil {
		return err
	}
	sourceUnstructured := &unstructured.Unstructured{Object: sourceMap}

	// Convert destination to unstructured
	destMap, err := runtime.DefaultUnstructuredConverter.ToUnstructured(destination)
	if err != nil {
		return err
	}
	destUnstructured := &unstructured.Unstructured{Object: destMap}

	// Extract only the spec from source
	sourceSpec, found, err := unstructured.NestedFieldCopy(sourceUnstructured.Object, "spec")
	if err != nil {
		return err
	}
	if !found {
		return fmt.Errorf("spec not found in source object")
	}

	// Set the spec in the destination
	err = unstructured.SetNestedField(destUnstructured.Object, sourceSpec, "spec")
	if err != nil {
		return err
	}

	// Convert back to the original object
	return runtime.DefaultUnstructuredConverter.FromUnstructured(destUnstructured.Object, destination)
}

func getSpec(obj client.Object) (any, error) {
	// Convert source to unstructured
	sourceMap, err := runtime.DefaultUnstructuredConverter.ToUnstructured(obj)
	if err != nil {
		return nil, err
	}
	sourceUnstructured := &unstructured.Unstructured{Object: sourceMap}
	// Extract only the spec from source
	spec, found, err := unstructured.NestedFieldCopy(sourceUnstructured.Object, "spec")
	if err != nil {
		return nil, err
	}
	if !found {
		return nil, nil
	}
	return spec, nil
}

// SpecChangeResult contains the result of spec change detection
type SpecChangeResult struct {
	// NewHash is the hash to set in the annotation (nil if no update needed)
	NewHash *string
	// NewGeneration is the generation to set in the annotation
	NewGeneration int64
	// NeedsUpdate indicates whether the resource needs to be updated
	NeedsUpdate bool
	// ManualChangeDetected indicates whether a manual change was detected
	ManualChangeDetected bool
}

// GetSpecChangeResult determines if a resource needs to be updated by comparing the desired spec hash
// with the last applied hash annotation. It also tracks generation to detect manual changes.
//
// Returns:
//   - SpecChangeResult with update information
//   - error if hash computation fails
func GetSpecChangeResult(current client.Object, desired client.Object) (SpecChangeResult, error) {
	desiredHash, err := GetSpecHash(desired)
	if err != nil {
		return SpecChangeResult{}, err
	}

	lastAppliedHash := getAnnotation(current, NvidiaAnnotationHashKey)
	lastAppliedGenStr := getAnnotation(current, NvidiaAnnotationGenerationKey)
	currentGen := current.GetGeneration()

	// Case 1: Hash annotation missing (external create or pre-upgrade resource)
	// Note: This is not first-time CREATE (handled separately in SyncResource with generation=1).
	// This handles existing resources without our annotations - we're about to update them,
	// so NewGeneration = currentGen + 1 is correct.
	if lastAppliedHash == "" {
		return SpecChangeResult{
			NewHash:       &desiredHash,
			NewGeneration: currentGen + 1,
			NeedsUpdate:   true,
		}, nil
	}

	// Case 2: Hash different (spec changed)
	if desiredHash != lastAppliedHash {
		return SpecChangeResult{
			NewHash:       &desiredHash,
			NewGeneration: currentGen + 1,
			NeedsUpdate:   true,
		}, nil
	}

	// Case 3: Hash same, but generation annotation missing (upgrade scenario)
	// Do a full update to ensure spec is exactly what we want - there could have been
	// manual edits before we added generation tracking. The cost is one extra Update
	// per resource during upgrade, but on next reconcile generations will match.
	if lastAppliedGenStr == "" {
		return SpecChangeResult{
			NewHash:       &desiredHash,
			NewGeneration: currentGen + 1,
			NeedsUpdate:   true,
		}, nil
	}

	// Case 4: Both annotations exist, check for manual changes
	lastAppliedGen, err := strconv.ParseInt(lastAppliedGenStr, 10, 64)
	if err != nil {
		// Corrupted annotation, force update to fix
		return SpecChangeResult{
			NewHash:       &desiredHash,
			NewGeneration: currentGen + 1,
			NeedsUpdate:   true,
		}, nil
	}

	// Detect manual changes: if current generation > last applied generation,
	// someone else modified the resource after our last update
	if currentGen > 0 && currentGen > lastAppliedGen {
		return SpecChangeResult{
			NewHash:              &desiredHash,
			NewGeneration:        currentGen + 1,
			NeedsUpdate:          true,
			ManualChangeDetected: true,
		}, nil
	}

	// No update needed
	return SpecChangeResult{
		NeedsUpdate: false,
	}, nil
}

// getAnnotation safely retrieves an annotation value from an object
func getAnnotation(obj client.Object, key string) string {
	annotations := obj.GetAnnotations()
	if annotations == nil {
		return ""
	}
	return annotations[key]
}

// generateSpecDiff creates a unified diff showing changes between old and new resource specs
func generateSpecDiff(oldResource, newResource client.Object) (string, error) {
	oldSpec, err := getSpec(oldResource)
	if err != nil {
		return "", fmt.Errorf("failed to get old spec: %w", err)
	}

	newSpec, err := getSpec(newResource)
	if err != nil {
		return "", fmt.Errorf("failed to get new spec: %w", err)
	}

	// Generate diff using cmp
	diff := cmp.Diff(oldSpec, newSpec)
	if diff == "" {
		return "", nil
	}

	return diff, nil
}

func GetSpecHash(obj client.Object) (string, error) {
	spec, err := getSpec(obj)
	if err != nil {
		return "", err
	}
	return GetResourceHash(spec)
}

// updateAnnotations sets both hash and generation annotations on an object
func updateAnnotations(obj client.Object, hash string, generation int64) {
	annotations := obj.GetAnnotations()
	if annotations == nil {
		annotations = map[string]string{}
	}
	annotations[NvidiaAnnotationHashKey] = hash
	annotations[NvidiaAnnotationGenerationKey] = strconv.FormatInt(generation, 10)
	obj.SetAnnotations(annotations)
}

// GetResourceHash returns a consistent hash for the given object spec
func GetResourceHash(obj any) (string, error) {
	// Convert obj to a map[string]interface{}
	objMap, err := json.Marshal(obj)
	if err != nil {
		return "", err
	}

	var objData map[string]interface{}
	if err := json.Unmarshal(objMap, &objData); err != nil {
		return "", err
	}

	// Sort keys to ensure consistent serialization
	sortedObjData := SortKeys(objData)

	// Serialize to JSON
	serialized, err := json.Marshal(sortedObjData)
	if err != nil {
		return "", err
	}

	// Compute the hash
	hasher := sha256.New()
	hasher.Write(serialized)
	return fmt.Sprintf("%x", hasher.Sum(nil)), nil
}

// SortKeys recursively sorts the keys of a map to ensure consistent serialization
func SortKeys(obj interface{}) interface{} {
	switch obj := obj.(type) {
	case map[string]interface{}:
		sortedMap := make(map[string]interface{})
		keys := make([]string, 0, len(obj))
		for k := range obj {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			sortedMap[k] = SortKeys(obj[k])
		}
		return sortedMap
	case []interface{}:
		// Check if the slice contains maps and sort them by the "name" field or the first available field
		if len(obj) > 0 {

			if _, ok := obj[0].(map[string]interface{}); ok {
				sort.SliceStable(obj, func(i, j int) bool {
					iMap, iOk := obj[i].(map[string]interface{})
					jMap, jOk := obj[j].(map[string]interface{})
					if iOk && jOk {
						// Try to sort by "name" if present
						iName, iNameOk := iMap["name"].(string)
						jName, jNameOk := jMap["name"].(string)
						if iNameOk && jNameOk {
							return iName < jName
						}

						// If "name" is not available, sort by the first key in each map
						if len(iMap) > 0 && len(jMap) > 0 {
							iFirstKey := firstKey(iMap)
							jFirstKey := firstKey(jMap)
							return iFirstKey < jFirstKey
						}
					}
					// If no valid comparison is possible, maintain the original order
					return false
				})
			}
		}
		for i, v := range obj {
			obj[i] = SortKeys(v)
		}
	}
	return obj
}

// Helper function to get the first key of a map (alphabetically sorted)
func firstKey(m map[string]interface{}) string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys[0]
}

func GetResourcesConfig(resources *v1alpha1.Resources) (*corev1.ResourceRequirements, error) {

	if resources == nil {
		return nil, nil
	}

	currentResources := &corev1.ResourceRequirements{}

	if resources.Limits != nil {
		if resources.Limits.CPU != "" {
			q, err := resource.ParseQuantity(resources.Limits.CPU)
			if err != nil {
				return nil, fmt.Errorf("parse limits cpu quantity: %w", err)
			}
			if currentResources.Limits == nil {
				currentResources.Limits = make(corev1.ResourceList)
			}
			currentResources.Limits[corev1.ResourceCPU] = q
		}
		if resources.Limits.Memory != "" {
			q, err := resource.ParseQuantity(resources.Limits.Memory)
			if err != nil {
				return nil, fmt.Errorf("parse limits memory quantity: %w", err)
			}
			if currentResources.Limits == nil {
				currentResources.Limits = make(corev1.ResourceList)
			}
			currentResources.Limits[corev1.ResourceMemory] = q
		}
		if resources.Limits.GPU != "" {
			q, err := resource.ParseQuantity(resources.Limits.GPU)
			if err != nil {
				return nil, fmt.Errorf("parse limits gpu quantity: %w", err)
			}
			if currentResources.Limits == nil {
				currentResources.Limits = make(corev1.ResourceList)
			}
			currentResources.Limits[getGPUResourceName(resources.Limits)] = q
		}
		for k, v := range resources.Limits.Custom {
			q, err := resource.ParseQuantity(v)
			if err != nil {
				return nil, fmt.Errorf("parse limits %s quantity: %w", k, err)
			}
			if currentResources.Limits == nil {
				currentResources.Limits = make(corev1.ResourceList)
			}
			currentResources.Limits[corev1.ResourceName(k)] = q
		}
	}
	if resources.Requests != nil {
		if resources.Requests.CPU != "" {
			q, err := resource.ParseQuantity(resources.Requests.CPU)
			if err != nil {
				return nil, fmt.Errorf("parse requests cpu quantity: %w", err)
			}
			if currentResources.Requests == nil {
				currentResources.Requests = make(corev1.ResourceList)
			}
			currentResources.Requests[corev1.ResourceCPU] = q
		}
		if resources.Requests.Memory != "" {
			q, err := resource.ParseQuantity(resources.Requests.Memory)
			if err != nil {
				return nil, fmt.Errorf("parse requests memory quantity: %w", err)
			}
			if currentResources.Requests == nil {
				currentResources.Requests = make(corev1.ResourceList)
			}
			currentResources.Requests[corev1.ResourceMemory] = q
		}
		for k, v := range resources.Requests.Custom {
			q, err := resource.ParseQuantity(v)
			if err != nil {
				return nil, fmt.Errorf("parse requests %s quantity: %w", k, err)
			}
			if currentResources.Requests == nil {
				currentResources.Requests = make(corev1.ResourceList)
			}
			currentResources.Requests[corev1.ResourceName(k)] = q
		}
	}
	if resources.Claims != nil {
		if currentResources.Claims == nil {
			currentResources.Claims = make([]corev1.ResourceClaim, 0)
		}
		currentResources.Claims = append(currentResources.Claims, resources.Claims...)
	}
	return currentResources, nil
}

func getGPUResourceName(resourceItem *v1alpha1.ResourceItem) corev1.ResourceName {
	if resourceItem == nil {
		return corev1.ResourceName(consts.KubeResourceGPUNvidia)
	}
	if resourceItem.GPUType != "" {
		return corev1.ResourceName(resourceItem.GPUType)
	}
	return corev1.ResourceName(consts.KubeResourceGPUNvidia)
}

// AppendUniqueImagePullSecrets appends secrets to existing, skipping any that already exist by name.
func AppendUniqueImagePullSecrets(existing, additional []corev1.LocalObjectReference) []corev1.LocalObjectReference {
	if len(additional) == 0 {
		return existing
	}
	seen := make(map[string]bool, len(existing))
	for _, s := range existing {
		seen[s.Name] = true
	}
	for _, s := range additional {
		if !seen[s.Name] {
			existing = append(existing, s)
			seen[s.Name] = true
		}
	}
	return existing
}

type Resource struct {
	object          client.Object
	isReady         bool
	readyReason     string
	serviceStatuses map[string]v1alpha1.ServiceReplicaStatus
}

func NewResource[T client.Object](resource T, isReady func() (bool, string)) (*Resource, error) {
	v := reflect.ValueOf(resource)
	// handles untype nil and typed nil
	if !v.IsValid() || v.IsNil() {
		return nil, fmt.Errorf("resource is nil")
	}
	ready, reason := isReady()
	return &Resource{
		object:      resource,
		isReady:     ready,
		readyReason: reason,
	}, nil
}

func NewResourceWithServiceStatuses[T client.Object](resource T, isReadyAndServiceStatuses func() (bool, string, map[string]v1alpha1.ServiceReplicaStatus)) (*Resource, error) {
	v := reflect.ValueOf(resource)
	// handles untype nil and typed nil
	if !v.IsValid() || v.IsNil() {
		return nil, fmt.Errorf("resource is nil")
	}
	ready, reason, serviceStatuses := isReadyAndServiceStatuses()
	return &Resource{
		object:          resource,
		isReady:         ready,
		readyReason:     reason,
		serviceStatuses: serviceStatuses,
	}, nil
}

func (r *Resource) IsReady() (bool, string) {
	return r.isReady, r.readyReason
}

func (r *Resource) GetName() string {
	return r.object.GetName()
}

func (r *Resource) GetServiceStatuses() map[string]v1alpha1.ServiceReplicaStatus {
	return r.serviceStatuses
}