embeddingserver_controller.go

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

// Package controllers contains the reconciliation logic for the EmbeddingServer custom resource.
// It handles the creation, update, and deletion of HuggingFace embedding inference servers in Kubernetes.
package controllers

import (
	"context"
	"fmt"
	"maps"
	"reflect"
	"time"

	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/client-go/tools/events"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/controller-runtime/pkg/log"

	mcpv1beta1 "github.com/stacklok/toolhive/cmd/thv-operator/api/v1beta1"
	ctrlutil "github.com/stacklok/toolhive/cmd/thv-operator/pkg/controllerutil"
	"github.com/stacklok/toolhive/cmd/thv-operator/pkg/imagepullsecrets"
)

// EmbeddingServerReconciler reconciles a EmbeddingServer object
type EmbeddingServerReconciler struct {
	client.Client
	Scheme           *runtime.Scheme
	Recorder         events.EventRecorder
	PlatformDetector *ctrlutil.SharedPlatformDetector
	// ImagePullSecretsDefaults are cluster-wide defaults sourced from the
	// operator chart, applied to the StatefulSet's PodSpec before the
	// user-provided PodTemplateSpec strategic-merge patch runs. The strategic
	// merge with the user PTS continues to additively merge the user's
	// imagePullSecrets entries on top, with the user's entries winning on
	// name collisions per Kubernetes' strategic-merge semantics.
	ImagePullSecretsDefaults imagepullsecrets.Defaults
}

const (
	// embeddingContainerName is the name of the embedding container used in pod templates
	embeddingContainerName = "embedding"

	// embeddingFinalizerName is the finalizer name for EmbeddingServer resources
	embeddingFinalizerName = "embeddingserver.toolhive.stacklok.dev/finalizer"

	// modelCacheMountPath is the mount path for the model cache volume
	modelCacheMountPath = "/data"
)

//+kubebuilder:rbac:groups=toolhive.stacklok.dev,resources=embeddingservers,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=toolhive.stacklok.dev,resources=embeddingservers/status,verbs=get;update;patch
//+kubebuilder:rbac:groups=toolhive.stacklok.dev,resources=embeddingservers/finalizers,verbs=update
//+kubebuilder:rbac:groups=apps,resources=statefulsets,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups="",resources=services,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups="",resources=persistentvolumeclaims,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups="",resources=secrets,verbs=get;list;watch
//+kubebuilder:rbac:groups=events.k8s.io,resources=events,verbs=create;patch

// Reconcile is part of the main kubernetes reconciliation loop which aims to
// move the current state of the cluster closer to the desired state.
//
//nolint:gocyclo // Reconciliation logic complexity is acceptable
func (r *EmbeddingServerReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	ctxLogger := log.FromContext(ctx)

	// Fetch the EmbeddingServer instance
	embedding := &mcpv1beta1.EmbeddingServer{}
	err := r.Get(ctx, req.NamespacedName, embedding)
	if err != nil {
		if errors.IsNotFound(err) {
			ctxLogger.Info("EmbeddingServer resource not found. Ignoring since object must be deleted")
			return ctrl.Result{}, nil
		}
		ctxLogger.Error(err, "Failed to get EmbeddingServer")
		return ctrl.Result{}, err
	}

	// Perform early validations
	if result, err := r.performValidations(ctx, embedding); err != nil || result.RequeueAfter > 0 {
		return result, err
	}

	// Handle deletion
	if result, done, err := r.handleDeletion(ctx, embedding); done {
		return result, err
	}

	// Add finalizer if needed
	if result, done, err := r.ensureFinalizer(ctx, embedding); done {
		return result, err
	}

	// Track if we need to requeue after status update
	var requeueResult ctrl.Result

	// Ensure statefulset exists and is up to date
	if result, err := r.ensureStatefulSet(ctx, embedding); err != nil {
		return ctrl.Result{}, err
	} else if result.RequeueAfter > 0 {
		requeueResult = result
	}

	// Ensure service exists
	if result, err := r.ensureService(ctx, embedding); err != nil {
		return ctrl.Result{}, err
	} else if result.RequeueAfter > 0 {
		// If we already have a requeue scheduled, keep the shorter duration
		if requeueResult.RequeueAfter == 0 || (result.RequeueAfter > 0 && result.RequeueAfter < requeueResult.RequeueAfter) {
			requeueResult = result
		}
	}

	// Always update the EmbeddingServer status before returning
	if err := r.updateEmbeddingServerStatus(ctx, embedding); err != nil {
		ctxLogger.Error(err, "Failed to update EmbeddingServer status")
		return ctrl.Result{}, err
	}

	return requeueResult, nil
}

// performValidations performs all early validations for the EmbeddingServer
//
//nolint:unparam // error return kept for consistency with reconciler pattern
func (r *EmbeddingServerReconciler) performValidations(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) (ctrl.Result, error) {
	ctxLogger := log.FromContext(ctx)

	// Validate PodTemplateSpec early
	if !r.validateAndUpdatePodTemplateStatus(ctx, embedding) {
		// Status fields were set by validateAndUpdatePodTemplateStatus, now update
		if err := r.Status().Update(ctx, embedding); err != nil {
			ctxLogger.Error(err, "Failed to update EmbeddingServer status after PodTemplateSpec validation failure")
			return ctrl.Result{}, err
		}
		return ctrl.Result{}, nil
	}

	return ctrl.Result{}, nil
}

// handleDeletion handles the deletion of EmbeddingServer resources
//
//nolint:unparam // ctrl.Result return kept for consistency with reconciler pattern
func (r *EmbeddingServerReconciler) handleDeletion(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) (ctrl.Result, bool, error) {
	if embedding.GetDeletionTimestamp() == nil {
		return ctrl.Result{}, false, nil
	}

	if controllerutil.ContainsFinalizer(embedding, embeddingFinalizerName) {
		r.finalizeEmbeddingServer(ctx, embedding)

		controllerutil.RemoveFinalizer(embedding, embeddingFinalizerName)
		err := r.Update(ctx, embedding)
		if err != nil {
			return ctrl.Result{}, true, err
		}
	}
	return ctrl.Result{}, true, nil
}

// ensureFinalizer ensures the finalizer is added to the EmbeddingServer
//
//nolint:unparam // ctrl.Result return kept for consistency with reconciler pattern
func (r *EmbeddingServerReconciler) ensureFinalizer(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) (ctrl.Result, bool, error) {
	if controllerutil.ContainsFinalizer(embedding, embeddingFinalizerName) {
		return ctrl.Result{}, false, nil
	}

	controllerutil.AddFinalizer(embedding, embeddingFinalizerName)
	err := r.Update(ctx, embedding)
	if err != nil {
		return ctrl.Result{}, true, err
	}
	return ctrl.Result{}, false, nil
}

// ensureStatefulSet ensures the statefulset exists and is up to date
func (r *EmbeddingServerReconciler) ensureStatefulSet(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) (ctrl.Result, error) {
	ctxLogger := log.FromContext(ctx)

	statefulSet := &appsv1.StatefulSet{}
	err := r.Get(ctx, types.NamespacedName{Name: embedding.Name, Namespace: embedding.Namespace}, statefulSet)
	if err != nil && errors.IsNotFound(err) {
		sts := r.statefulSetForEmbedding(ctx, embedding)
		if sts == nil {
			ctxLogger.Error(nil, "Failed to create StatefulSet object")
			return ctrl.Result{}, fmt.Errorf("failed to create StatefulSet object")
		}
		ctxLogger.Info("Creating a new StatefulSet", "StatefulSet.Namespace", sts.Namespace, "StatefulSet.Name", sts.Name)
		err = r.Create(ctx, sts)
		if err != nil {
			ctxLogger.Error(err, "Failed to create new StatefulSet", "StatefulSet.Namespace", sts.Namespace, "StatefulSet.Name", sts.Name)
			return ctrl.Result{}, err
		}
		// StatefulSet created successfully, continue to ensure service
		return ctrl.Result{}, nil
	} else if err != nil {
		ctxLogger.Error(err, "Failed to get StatefulSet")
		return ctrl.Result{}, err
	}

	// Ensure the statefulset size matches the spec
	desiredReplicas := embedding.GetReplicas()
	if *statefulSet.Spec.Replicas != desiredReplicas {
		statefulSet.Spec.Replicas = &desiredReplicas
		if err := r.Update(ctx, statefulSet); err != nil {
			ctxLogger.Error(err, "Failed to update StatefulSet replicas",
				"StatefulSet.Namespace", statefulSet.Namespace,
				"StatefulSet.Name", statefulSet.Name)
			return ctrl.Result{}, err
		}
		return ctrl.Result{RequeueAfter: time.Second}, nil
	}

	// Check if the statefulset spec changed
	if r.statefulSetNeedsUpdate(ctx, statefulSet, embedding) {
		newStatefulSet := r.statefulSetForEmbedding(ctx, embedding)
		statefulSet.Spec = newStatefulSet.Spec
		statefulSet.Annotations = newStatefulSet.Annotations
		statefulSet.Labels = newStatefulSet.Labels
		if err := r.Update(ctx, statefulSet); err != nil {
			ctxLogger.Error(err, "Failed to update StatefulSet",
				"StatefulSet.Namespace", statefulSet.Namespace,
				"StatefulSet.Name", statefulSet.Name)
			return ctrl.Result{}, err
		}
		return ctrl.Result{RequeueAfter: time.Second}, nil
	}

	return ctrl.Result{}, nil
}

// ensureService ensures the service exists and is up to date
//
//nolint:unparam // ctrl.Result return kept for consistency with reconciler pattern
func (r *EmbeddingServerReconciler) ensureService(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) (ctrl.Result, error) {
	ctxLogger := log.FromContext(ctx)

	service := &corev1.Service{}
	err := r.Get(ctx, types.NamespacedName{Name: embedding.Name, Namespace: embedding.Namespace}, service)
	if err != nil && errors.IsNotFound(err) {
		svc := r.serviceForEmbedding(ctx, embedding)
		if svc == nil {
			ctxLogger.Error(nil, "Failed to create Service object")
			return ctrl.Result{}, fmt.Errorf("failed to create Service object")
		}
		ctxLogger.Info("Creating a new Service", "Service.Namespace", svc.Namespace, "Service.Name", svc.Name)
		err = r.Create(ctx, svc)
		if err != nil {
			ctxLogger.Error(err, "Failed to create new Service", "Service.Namespace", svc.Namespace, "Service.Name", svc.Name)
			return ctrl.Result{}, err
		}
		// Service created successfully, continue to update status
		return ctrl.Result{}, nil
	} else if err != nil {
		ctxLogger.Error(err, "Failed to get Service")
		return ctrl.Result{}, err
	}

	// Check if the service needs to be updated
	if r.serviceNeedsUpdate(service, embedding) {
		desiredService := r.serviceForEmbedding(ctx, embedding)
		service.Spec.Ports = desiredService.Spec.Ports
		// Merge (not replace) Labels/Annotations so keys written by external controllers
		// (e.g. GKE NEG's cloud.google.com/* annotations) are preserved while the
		// operator-owned values are applied.
		service.Labels = ctrlutil.MergeLabels(desiredService.Labels, service.Labels)
		service.Annotations = ctrlutil.MergeAnnotations(desiredService.Annotations, service.Annotations)
		// Preserve ClusterIP as it's immutable
		if err := r.Update(ctx, service); err != nil {
			ctxLogger.Error(err, "Failed to update Service",
				"Service.Namespace", service.Namespace,
				"Service.Name", service.Name)
			return ctrl.Result{}, err
		}
		ctxLogger.Info("Updated Service", "Service.Namespace", service.Namespace, "Service.Name", service.Name)
		return ctrl.Result{RequeueAfter: time.Second}, nil
	}

	return ctrl.Result{}, nil
}

// serviceNeedsUpdate checks if the service needs to be updated based on the embedding spec
func (*EmbeddingServerReconciler) serviceNeedsUpdate(
	service *corev1.Service,
	embedding *mcpv1beta1.EmbeddingServer,
) bool {
	desiredPort := embedding.GetPort()

	// Check if any port has changed
	for _, port := range service.Spec.Ports {
		if port.Name == "http" && port.Port != desiredPort {
			return true
		}
	}

	// Check ResourceOverrides (annotations and labels)
	expectedAnnotations := make(map[string]string)
	expectedLabels := make(map[string]string)

	if embedding.Spec.ResourceOverrides != nil && embedding.Spec.ResourceOverrides.Service != nil {
		if embedding.Spec.ResourceOverrides.Service.Annotations != nil {
			maps.Copy(expectedAnnotations, embedding.Spec.ResourceOverrides.Service.Annotations)
		}
		if embedding.Spec.ResourceOverrides.Service.Labels != nil {
			maps.Copy(expectedLabels, embedding.Spec.ResourceOverrides.Service.Labels)
		}
	}

	// Subset check rather than exact equality: the Service is co-owned by external
	// controllers (e.g. GKE NEG/Gateway writes cloud.google.com/* annotations), so only
	// the operator-owned keys must match. Exact equality would treat those external
	// annotations as drift and hot-loop Update against the concurrent writer.
	if !ctrlutil.MapIsSubset(expectedAnnotations, service.Annotations) {
		return true
	}

	if !ctrlutil.MapIsSubset(expectedLabels, service.Labels) {
		return true
	}

	return false
}

// validateAndUpdatePodTemplateStatus validates the PodTemplateSpec and sets the status condition
// Status is not updated here - it will be updated at the end of reconciliation
func (r *EmbeddingServerReconciler) validateAndUpdatePodTemplateStatus(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) bool {
	ctxLogger := log.FromContext(ctx)

	if embedding.Spec.PodTemplateSpec == nil {
		meta.SetStatusCondition(&embedding.Status.Conditions, metav1.Condition{
			Type:               mcpv1beta1.ConditionPodTemplateValid,
			Status:             metav1.ConditionTrue,
			Reason:             mcpv1beta1.ConditionReasonPodTemplateValid,
			Message:            "No PodTemplateSpec provided",
			ObservedGeneration: embedding.Generation,
		})
		return true
	}

	// Parse and validate PodTemplateSpec using builder
	_, err := ctrlutil.NewPodTemplateSpecBuilder(embedding.Spec.PodTemplateSpec, embeddingContainerName)
	if err != nil {
		ctxLogger.Error(err, "Invalid PodTemplateSpec")
		embedding.Status.Phase = mcpv1beta1.EmbeddingServerPhaseFailed
		embedding.Status.Message = fmt.Sprintf("Invalid PodTemplateSpec: %v", err)
		meta.SetStatusCondition(&embedding.Status.Conditions, metav1.Condition{
			Type:               mcpv1beta1.ConditionPodTemplateValid,
			Status:             metav1.ConditionFalse,
			Reason:             mcpv1beta1.ConditionReasonPodTemplateInvalid,
			Message:            fmt.Sprintf("Invalid PodTemplateSpec: %v", err),
			ObservedGeneration: embedding.Generation,
		})
		r.Recorder.Eventf(
			embedding,
			nil,
			corev1.EventTypeWarning,
			"ValidationFailed",
			"ValidatePodTemplateSpec",
			"Invalid PodTemplateSpec: %v",
			err,
		)
		return false
	}

	meta.SetStatusCondition(&embedding.Status.Conditions, metav1.Condition{
		Type:               mcpv1beta1.ConditionPodTemplateValid,
		Status:             metav1.ConditionTrue,
		Reason:             mcpv1beta1.ConditionReasonPodTemplateValid,
		Message:            "PodTemplateSpec is valid",
		ObservedGeneration: embedding.Generation,
	})

	return true
}

// statefulSetForEmbedding creates a StatefulSet for the embedding server
func (r *EmbeddingServerReconciler) statefulSetForEmbedding(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) *appsv1.StatefulSet {
	replicas := embedding.GetReplicas()
	labels := r.labelsForEmbedding(embedding)

	// Build container
	container := r.buildEmbeddingContainer(embedding)

	// Build pod template
	podTemplate := r.buildPodTemplate(labels, container)

	// Apply statefulset overrides
	stsAnnotations, stsLabels := r.applyStatefulSetOverrides(embedding, &podTemplate)

	// Merge ResourceOverrides labels into base labels
	finalLabels := make(map[string]string)
	maps.Copy(finalLabels, labels)
	maps.Copy(finalLabels, stsLabels)

	statefulSet := &appsv1.StatefulSet{
		ObjectMeta: metav1.ObjectMeta{
			Name:        embedding.Name,
			Namespace:   embedding.Namespace,
			Labels:      finalLabels,
			Annotations: stsAnnotations,
		},
		Spec: appsv1.StatefulSetSpec{
			Replicas:    &replicas,
			ServiceName: embedding.Name, // Required for StatefulSet
			Selector: &metav1.LabelSelector{
				MatchLabels: labels,
			},
			Template: podTemplate,
		},
	}

	// Add volumeClaimTemplates if model caching is enabled
	if embedding.IsModelCacheEnabled() {
		statefulSet.Spec.VolumeClaimTemplates = r.buildVolumeClaimTemplates(embedding)
	}

	// Apply user-provided PodTemplateSpec customizations via strategic merge patch.
	// This must happen after the controller-generated template is fully populated so
	// that user fields override controller defaults rather than the other way around.
	// The merge is soft-fail: invalid input is logged and the StatefulSet is built
	// from controller defaults. See applyPodTemplateSpecToStatefulSet's godoc.
	r.applyPodTemplateSpecToStatefulSet(ctx, embedding, statefulSet)

	if err := ctrl.SetControllerReference(embedding, statefulSet, r.Scheme); err != nil {
		return nil
	}
	return statefulSet
}

// buildVolumeClaimTemplates builds the volumeClaimTemplates for the StatefulSet
func (r *EmbeddingServerReconciler) buildVolumeClaimTemplates(
	embedding *mcpv1beta1.EmbeddingServer,
) []corev1.PersistentVolumeClaim {
	size := "10Gi"
	if embedding.Spec.ModelCache.Size != "" {
		size = embedding.Spec.ModelCache.Size
	}

	accessMode := corev1.ReadWriteOnce
	if embedding.Spec.ModelCache.AccessMode != "" {
		accessMode = corev1.PersistentVolumeAccessMode(embedding.Spec.ModelCache.AccessMode)
	}

	pvc := corev1.PersistentVolumeClaim{
		ObjectMeta: metav1.ObjectMeta{
			Name:   "model-cache",
			Labels: r.labelsForEmbedding(embedding),
		},
		Spec: corev1.PersistentVolumeClaimSpec{
			AccessModes: []corev1.PersistentVolumeAccessMode{accessMode},
			Resources: corev1.VolumeResourceRequirements{
				Requests: corev1.ResourceList{
					corev1.ResourceStorage: resource.MustParse(size),
				},
			},
		},
	}

	if embedding.Spec.ModelCache.StorageClassName != nil {
		pvc.Spec.StorageClassName = embedding.Spec.ModelCache.StorageClassName
	}

	// Apply resource overrides if specified
	if embedding.Spec.ResourceOverrides != nil && embedding.Spec.ResourceOverrides.PersistentVolumeClaim != nil {
		if pvc.Annotations == nil && embedding.Spec.ResourceOverrides.PersistentVolumeClaim.Annotations != nil {
			pvc.Annotations = make(map[string]string)
		}
		if embedding.Spec.ResourceOverrides.PersistentVolumeClaim.Annotations != nil {
			maps.Copy(pvc.Annotations, embedding.Spec.ResourceOverrides.PersistentVolumeClaim.Annotations)
		}
		if embedding.Spec.ResourceOverrides.PersistentVolumeClaim.Labels != nil {
			maps.Copy(pvc.Labels, embedding.Spec.ResourceOverrides.PersistentVolumeClaim.Labels)
		}
	}

	return []corev1.PersistentVolumeClaim{pvc}
}

// buildEmbeddingContainer builds the container spec for the embedding server
func (r *EmbeddingServerReconciler) buildEmbeddingContainer(embedding *mcpv1beta1.EmbeddingServer) corev1.Container {
	// Build container args
	args := []string{
		"--model-id", embedding.Spec.Model,
		"--port", fmt.Sprintf("%d", embedding.GetPort()),
	}
	args = append(args, embedding.Spec.Args...)

	// Build environment variables
	envVars := r.buildEnvVars(embedding)

	// Build container
	container := corev1.Container{
		Name:            embeddingContainerName,
		Image:           embedding.Spec.Image,
		Args:            args,
		Env:             envVars,
		ImagePullPolicy: embedding.GetImagePullPolicy(),
		Ports: []corev1.ContainerPort{
			{
				Name:          "http",
				ContainerPort: embedding.GetPort(),
				Protocol:      corev1.ProtocolTCP,
			},
		},
		LivenessProbe:  r.buildLivenessProbe(embedding),
		ReadinessProbe: r.buildReadinessProbe(embedding),
	}

	// Add volume mount and HF_HOME for model cache if enabled
	if embedding.IsModelCacheEnabled() {
		container.VolumeMounts = []corev1.VolumeMount{
			{
				Name:      "model-cache",
				MountPath: modelCacheMountPath,
			},
		}
		container.Env = append(container.Env, corev1.EnvVar{
			Name:  "HF_HOME",
			Value: modelCacheMountPath,
		})
	}

	// Add resources if specified
	r.applyResourceRequirements(embedding, &container)

	return container
}

// buildEnvVars builds environment variables for the container
func (*EmbeddingServerReconciler) buildEnvVars(embedding *mcpv1beta1.EmbeddingServer) []corev1.EnvVar {
	envVars := []corev1.EnvVar{
		{
			Name:  "MODEL_ID",
			Value: embedding.Spec.Model,
		},
	}

	// Add HuggingFace token from secret if provided
	if embedding.Spec.HFTokenSecretRef != nil {
		envVars = append(envVars, corev1.EnvVar{
			Name: "HF_TOKEN",
			ValueFrom: &corev1.EnvVarSource{
				SecretKeyRef: &corev1.SecretKeySelector{
					LocalObjectReference: corev1.LocalObjectReference{
						Name: embedding.Spec.HFTokenSecretRef.Name,
					},
					Key: embedding.Spec.HFTokenSecretRef.Key,
				},
			},
		})
	}

	for _, env := range embedding.Spec.Env {
		envVars = append(envVars, corev1.EnvVar{
			Name:  env.Name,
			Value: env.Value,
		})
	}
	return envVars
}

// buildLivenessProbe builds the liveness probe for the container
func (*EmbeddingServerReconciler) buildLivenessProbe(embedding *mcpv1beta1.EmbeddingServer) *corev1.Probe {
	return &corev1.Probe{
		ProbeHandler: corev1.ProbeHandler{
			HTTPGet: &corev1.HTTPGetAction{
				Path: "/health",
				Port: intstr.FromInt(int(embedding.GetPort())),
			},
		},
		InitialDelaySeconds: 60,
		PeriodSeconds:       30,
		TimeoutSeconds:      10,
		FailureThreshold:    3,
	}
}

// buildReadinessProbe builds the readiness probe for the container
func (*EmbeddingServerReconciler) buildReadinessProbe(embedding *mcpv1beta1.EmbeddingServer) *corev1.Probe {
	return &corev1.Probe{
		ProbeHandler: corev1.ProbeHandler{
			HTTPGet: &corev1.HTTPGetAction{
				Path: "/health",
				Port: intstr.FromInt(int(embedding.GetPort())),
			},
		},
		InitialDelaySeconds: 30,
		PeriodSeconds:       10,
		TimeoutSeconds:      5,
		FailureThreshold:    3,
	}
}

// applyResourceRequirements applies resource requirements to the container
func (*EmbeddingServerReconciler) applyResourceRequirements(embedding *mcpv1beta1.EmbeddingServer, container *corev1.Container) {
	if embedding.Spec.Resources.Limits.CPU == "" && embedding.Spec.Resources.Limits.Memory == "" &&
		embedding.Spec.Resources.Requests.CPU == "" && embedding.Spec.Resources.Requests.Memory == "" {
		return
	}

	container.Resources = corev1.ResourceRequirements{
		Limits:   corev1.ResourceList{},
		Requests: corev1.ResourceList{},
	}

	if embedding.Spec.Resources.Limits.CPU != "" {
		container.Resources.Limits[corev1.ResourceCPU] = resource.MustParse(embedding.Spec.Resources.Limits.CPU)
	}
	if embedding.Spec.Resources.Limits.Memory != "" {
		container.Resources.Limits[corev1.ResourceMemory] = resource.MustParse(embedding.Spec.Resources.Limits.Memory)
	}
	if embedding.Spec.Resources.Requests.CPU != "" {
		container.Resources.Requests[corev1.ResourceCPU] = resource.MustParse(embedding.Spec.Resources.Requests.CPU)
	}
	if embedding.Spec.Resources.Requests.Memory != "" {
		container.Resources.Requests[corev1.ResourceMemory] = resource.MustParse(embedding.Spec.Resources.Requests.Memory)
	}
}

// buildPodTemplate builds the pod template for the statefulset.
// User-provided PodTemplateSpec customizations are applied later in
// statefulSetForEmbedding via strategic merge patch.
//
// Cluster-wide chart defaults for imagePullSecrets are placed on the base
// PodSpec here so that a subsequent strategic-merge with the user PTS
// additively unions the lists (Kubernetes treats PodSpec.ImagePullSecrets
// as a merge list keyed on Name; user entries win on name collisions).
func (r *EmbeddingServerReconciler) buildPodTemplate(
	labels map[string]string,
	container corev1.Container,
) corev1.PodTemplateSpec {
	// Note: Volumes for model cache are managed by StatefulSet volumeClaimTemplates
	// and will be automatically mounted with the name "model-cache"
	return corev1.PodTemplateSpec{
		ObjectMeta: metav1.ObjectMeta{
			Labels: labels,
		},
		Spec: corev1.PodSpec{
			ImagePullSecrets: r.ImagePullSecretsDefaults.List(),
			Containers:       []corev1.Container{container},
		},
	}
}

// applyPodTemplateSpecToStatefulSet applies user-provided PodTemplateSpec customizations
// to the StatefulSet's pod template using strategic merge patch. This preserves every
// user-supplied PodSpec field (imagePullSecrets, additional volumes, priorityClassName,
// topologySpreadConstraints, init containers, sidecars, etc.) while keeping controller
// defaults for fields the user did not set.
//
// The merge itself is delegated to ctrlutil.ApplyPodTemplateSpecPatch — which is
// policy-neutral. Invalid user input is treated here as a soft failure: the merge is
// skipped and the StatefulSet is built from controller defaults. The user-facing signal
// lives on the EmbeddingServer status (set by validateAndUpdatePodTemplateStatus):
// Phase=Failed and ConditionPodTemplateValid=False. This mirrors the pre-existing
// tolerant behavior — refusing to create the StatefulSet would leave the resource stuck
// with no pod and no observable controller-side state, while the validation condition
// already tells the user exactly why their input was rejected. The vMCP controller
// makes the opposite choice (hard-fail) for the same helper; both are documented on
// ApplyPodTemplateSpecPatch's godoc.
//
// The function does not return an error: every failure mode is converted to a log line
// plus controller-default fallback at this call site.
func (*EmbeddingServerReconciler) applyPodTemplateSpecToStatefulSet(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
	statefulSet *appsv1.StatefulSet,
) {
	if embedding.Spec.PodTemplateSpec == nil || len(embedding.Spec.PodTemplateSpec.Raw) == 0 {
		return
	}

	logger := log.FromContext(ctx)

	// Validate the user-provided PodTemplateSpec is well-formed.
	// We don't check builder.Build() == nil for "empty" customizations: that helper
	// only enumerates a subset of PodSpec fields and would skip the patch for
	// fields like runtimeClassName or topologySpreadConstraints. Strategic merge
	// patch is a no-op for `{}` anyway, so always running it is safe.
	if _, err := ctrlutil.NewPodTemplateSpecBuilder(embedding.Spec.PodTemplateSpec, embeddingContainerName); err != nil {
		logger.Info("Skipping PodTemplateSpec merge: input is invalid; StatefulSet will use controller defaults",
			"error", err.Error(),
			"embeddingserver", embedding.Name,
			"namespace", embedding.Namespace)
		return
	}

	merged, err := ctrlutil.ApplyPodTemplateSpecPatch(statefulSet.Spec.Template, embedding.Spec.PodTemplateSpec.Raw)
	if err != nil {
		// Soft failure: log and fall back to controller defaults. See function
		// godoc above for the rationale and the contrast with the vMCP caller.
		logger.Info("Skipping PodTemplateSpec merge: strategic merge patch failed; StatefulSet will use controller defaults",
			"error", err.Error(),
			"embeddingserver", embedding.Name,
			"namespace", embedding.Namespace)
		return
	}

	statefulSet.Spec.Template = merged

	logger.V(1).Info("Applied PodTemplateSpec customizations to StatefulSet",
		"embeddingserver", embedding.Name,
		"namespace", embedding.Namespace)
}

// applyStatefulSetOverrides applies statefulset-level overrides and returns annotations and labels
func (*EmbeddingServerReconciler) applyStatefulSetOverrides(
	embedding *mcpv1beta1.EmbeddingServer,
	podTemplate *corev1.PodTemplateSpec,
) (map[string]string, map[string]string) {
	annotations := make(map[string]string)
	labels := make(map[string]string)

	if embedding.Spec.ResourceOverrides == nil || embedding.Spec.ResourceOverrides.StatefulSet == nil {
		return annotations, labels
	}

	if embedding.Spec.ResourceOverrides.StatefulSet.Annotations != nil {
		maps.Copy(annotations, embedding.Spec.ResourceOverrides.StatefulSet.Annotations)
	}

	if embedding.Spec.ResourceOverrides.StatefulSet.Labels != nil {
		maps.Copy(labels, embedding.Spec.ResourceOverrides.StatefulSet.Labels)
	}

	if embedding.Spec.ResourceOverrides.StatefulSet.PodTemplateMetadataOverrides != nil {
		if podTemplate.Annotations == nil {
			podTemplate.Annotations = make(map[string]string)
		}
		if embedding.Spec.ResourceOverrides.StatefulSet.PodTemplateMetadataOverrides.Annotations != nil {
			maps.Copy(
				podTemplate.Annotations,
				embedding.Spec.ResourceOverrides.StatefulSet.PodTemplateMetadataOverrides.Annotations,
			)
		}
		if embedding.Spec.ResourceOverrides.StatefulSet.PodTemplateMetadataOverrides.Labels != nil {
			maps.Copy(podTemplate.Labels, embedding.Spec.ResourceOverrides.StatefulSet.PodTemplateMetadataOverrides.Labels)
		}
	}

	return annotations, labels
}

// serviceForEmbedding creates a Service for the embedding server
func (r *EmbeddingServerReconciler) serviceForEmbedding(
	_ context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) *corev1.Service {
	labels := r.labelsForEmbedding(embedding)
	annotations := make(map[string]string)

	// Apply service overrides if specified
	finalLabels := make(map[string]string)
	maps.Copy(finalLabels, labels)

	if embedding.Spec.ResourceOverrides != nil && embedding.Spec.ResourceOverrides.Service != nil {
		if embedding.Spec.ResourceOverrides.Service.Annotations != nil {
			maps.Copy(annotations, embedding.Spec.ResourceOverrides.Service.Annotations)
		}
		if embedding.Spec.ResourceOverrides.Service.Labels != nil {
			maps.Copy(finalLabels, embedding.Spec.ResourceOverrides.Service.Labels)
		}
	}

	service := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:        embedding.Name,
			Namespace:   embedding.Namespace,
			Labels:      finalLabels,
			Annotations: annotations,
		},
		Spec: corev1.ServiceSpec{
			Selector: labels,
			Ports: []corev1.ServicePort{
				{
					Name:       "http",
					Port:       embedding.GetPort(),
					TargetPort: intstr.FromInt(int(embedding.GetPort())),
					Protocol:   corev1.ProtocolTCP,
				},
			},
		},
	}

	if err := ctrl.SetControllerReference(embedding, service, r.Scheme); err != nil {
		return nil
	}
	return service
}

// labelsForEmbedding returns the labels for the embedding resources
func (*EmbeddingServerReconciler) labelsForEmbedding(embedding *mcpv1beta1.EmbeddingServer) map[string]string {
	return map[string]string{
		"app.kubernetes.io/name":       "embeddingserver",
		"app.kubernetes.io/instance":   embedding.Name,
		"app.kubernetes.io/component":  "embedding-server",
		"app.kubernetes.io/managed-by": "toolhive-operator",
	}
}

// statefulSetNeedsUpdate checks if the statefulset needs to be updated
func (r *EmbeddingServerReconciler) statefulSetNeedsUpdate(
	ctx context.Context,
	currentSts *appsv1.StatefulSet,
	embedding *mcpv1beta1.EmbeddingServer,
) bool {
	// Generate the expected StatefulSet from the current spec
	newSts := r.statefulSetForEmbedding(ctx, embedding)
	if newSts == nil {
		// If we can't generate a new StatefulSet, assume update is needed
		return true
	}

	// Check StatefulSet-level fields
	if r.statefulSetMetadataChanged(currentSts, newSts) {
		return true
	}

	// Check container-level fields
	existingContainer, newContainer := r.findEmbeddingContainers(currentSts, newSts)
	if existingContainer == nil || newContainer == nil {
		return true
	}

	if r.containerNeedsUpdate(existingContainer, newContainer) {
		return true
	}

	// Check pod template metadata
	if r.podTemplateMetadataChanged(currentSts, newSts) {
		return true
	}

	return false
}

// statefulSetMetadataChanged checks if StatefulSet-level metadata has changed
func (*EmbeddingServerReconciler) statefulSetMetadataChanged(currentSts, newSts *appsv1.StatefulSet) bool {
	if *currentSts.Spec.Replicas != *newSts.Spec.Replicas {
		return true
	}
	if !reflect.DeepEqual(newSts.Annotations, currentSts.Annotations) {
		return true
	}
	if !reflect.DeepEqual(newSts.Labels, currentSts.Labels) {
		return true
	}
	return false
}

// findEmbeddingContainers finds the embedding container in both StatefulSets
func (*EmbeddingServerReconciler) findEmbeddingContainers(
	currentSts, newSts *appsv1.StatefulSet,
) (*corev1.Container, *corev1.Container) {
	var existingContainer *corev1.Container
	for i := range currentSts.Spec.Template.Spec.Containers {
		if currentSts.Spec.Template.Spec.Containers[i].Name == embeddingContainerName {
			existingContainer = &currentSts.Spec.Template.Spec.Containers[i]
			break
		}
	}

	var newContainer *corev1.Container
	for i := range newSts.Spec.Template.Spec.Containers {
		if newSts.Spec.Template.Spec.Containers[i].Name == embeddingContainerName {
			newContainer = &newSts.Spec.Template.Spec.Containers[i]
			break
		}
	}

	return existingContainer, newContainer
}

// containerNeedsUpdate checks if the container spec has changed
func (*EmbeddingServerReconciler) containerNeedsUpdate(existingContainer, newContainer *corev1.Container) bool {
	if existingContainer.Image != newContainer.Image {
		return true
	}
	if !reflect.DeepEqual(existingContainer.Args, newContainer.Args) {
		return true
	}
	if !reflect.DeepEqual(existingContainer.Env, newContainer.Env) {
		return true
	}
	if !reflect.DeepEqual(existingContainer.Ports, newContainer.Ports) {
		return true
	}
	if existingContainer.ImagePullPolicy != newContainer.ImagePullPolicy {
		return true
	}
	if !reflect.DeepEqual(existingContainer.Resources, newContainer.Resources) {
		return true
	}
	return false
}

// podTemplateMetadataChanged checks if pod template metadata has changed
func (*EmbeddingServerReconciler) podTemplateMetadataChanged(currentSts, newSts *appsv1.StatefulSet) bool {
	if !reflect.DeepEqual(currentSts.Spec.Template.Annotations, newSts.Spec.Template.Annotations) {
		return true
	}
	if !reflect.DeepEqual(currentSts.Spec.Template.Labels, newSts.Spec.Template.Labels) {
		return true
	}
	return false
}

// updateEmbeddingServerStatus updates the status based on statefulset state
func (r *EmbeddingServerReconciler) updateEmbeddingServerStatus(
	ctx context.Context,
	embedding *mcpv1beta1.EmbeddingServer,
) error {
	ctxLogger := log.FromContext(ctx)

	// Set the service URL if not already set
	if embedding.Status.URL == "" {
		embedding.Status.URL = fmt.Sprintf("http://%s.%s.svc.cluster.local:%d",
			embedding.Name, embedding.Namespace, embedding.GetPort())
	}

	statefulSet := &appsv1.StatefulSet{}
	err := r.Get(ctx, types.NamespacedName{Name: embedding.Name, Namespace: embedding.Namespace}, statefulSet)
	if err != nil {
		if errors.IsNotFound(err) {
			embedding.Status.Phase = mcpv1beta1.EmbeddingServerPhasePending
			embedding.Status.ReadyReplicas = 0
		} else {
			return err
		}
	} else {
		embedding.Status.ReadyReplicas = statefulSet.Status.ReadyReplicas
		embedding.Status.ObservedGeneration = embedding.Generation

		// Determine phase and message based on statefulset status using immutable assignment
		type phaseInfo struct {
			phase   mcpv1beta1.EmbeddingServerPhase
			message string
		}

		info := func() phaseInfo {
			if statefulSet.Status.ReadyReplicas > 0 {
				return phaseInfo{
					phase:   mcpv1beta1.EmbeddingServerPhaseReady,
					message: "Embedding server is running",
				}
			}
			if statefulSet.Status.Replicas > 0 && statefulSet.Status.ReadyReplicas == 0 {
				// Check if pods are downloading the model
				return phaseInfo{
					phase:   mcpv1beta1.EmbeddingServerPhaseDownloading,
					message: "Downloading embedding model",
				}
			}
			return phaseInfo{
				phase:   mcpv1beta1.EmbeddingServerPhasePending,
				message: "Waiting for statefulset",
			}
		}()

		embedding.Status.Phase = info.phase
		embedding.Status.Message = info.message
	}

	err = r.Status().Update(ctx, embedding)
	if err != nil {
		ctxLogger.Error(err, "Failed to update EmbeddingServer status")
		return err
	}

	return nil
}

// finalizeEmbeddingServer performs cleanup before the EmbeddingServer is deleted
func (r *EmbeddingServerReconciler) finalizeEmbeddingServer(ctx context.Context, embedding *mcpv1beta1.EmbeddingServer) {
	ctxLogger := log.FromContext(ctx)
	ctxLogger.Info("Finalizing EmbeddingServer", "name", embedding.Name)

	// Update status to Terminating
	embedding.Status.Phase = mcpv1beta1.EmbeddingServerPhaseTerminating
	if err := r.Status().Update(ctx, embedding); err != nil {
		ctxLogger.Error(err, "Failed to update EmbeddingServer status to Terminating")
	}

	// Cleanup logic here if needed
	// For now, Kubernetes will handle cascade deletion of owned resources

	r.Recorder.Eventf(embedding, nil, corev1.EventTypeNormal, "Deleted", "Finalize", "EmbeddingServer has been finalized")
}

// SetupWithManager sets up the controller with the Manager.
func (r *EmbeddingServerReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&mcpv1beta1.EmbeddingServer{}).
		Owns(&appsv1.StatefulSet{}).
		Owns(&corev1.Service{}).
		Owns(&corev1.PersistentVolumeClaim{}).
		Complete(r)
}