sandbox_controller.go

// Copyright 2025 The Kubernetes Authors.
//
// 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 controllers

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"reflect"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	corev1 "k8s.io/api/core/v1"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/log"
	ctrlMetrics "sigs.k8s.io/controller-runtime/pkg/metrics"
	"sigs.k8s.io/controller-runtime/pkg/predicate"

	sandboxv1alpha1 "sigs.k8s.io/agent-sandbox/api/v1alpha1"
	extensionsv1alpha1 "sigs.k8s.io/agent-sandbox/extensions/api/v1alpha1"
	asmetrics "sigs.k8s.io/agent-sandbox/internal/metrics"
)

const (
	sandboxLabel                 = "agents.x-k8s.io/sandbox-name-hash"
	SandboxPodNameAnnotation     = "agents.x-k8s.io/pod-name"
	SandboxTemplateRefAnnotation = "agents.x-k8s.io/sandbox-template-ref"
	sandboxControllerFieldOwner  = "sandbox-controller"
	metricsCollectTimeout        = 5 * time.Second
)

var (
	// Scheme for use by sandbox controllers. Registers required types for client.
	Scheme = runtime.NewScheme()
)

func init() {
	utilruntime.Must(clientgoscheme.AddToScheme(Scheme))
	utilruntime.Must(sandboxv1alpha1.AddToScheme(Scheme))
}

// SandboxReconciler reconciles a Sandbox object
type SandboxReconciler struct {
	client.Client
	Scheme *runtime.Scheme
	Tracer asmetrics.Instrumenter
}

//+kubebuilder:rbac:groups=agents.x-k8s.io,resources=sandboxes,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=agents.x-k8s.io,resources=sandboxes/finalizers,verbs=get;update;patch
//+kubebuilder:rbac:groups=agents.x-k8s.io,resources=sandboxes/status,verbs=get;update;patch
//+kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=core,resources=services,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=core,resources=persistentvolumeclaims,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=coordination.k8s.io,resources=leases,verbs=get;list;watch;create;update;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.
// TODO(user): Modify the Reconcile function to compare the state specified by
// the Sandbox object against the actual cluster state, and then
// perform operations to make the cluster state reflect the state specified by
// the user.
//
// For more details, check Reconcile and its Result here:
// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.14.1/pkg/reconcile
func (r *SandboxReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := log.FromContext(ctx)

	sandbox := &sandboxv1alpha1.Sandbox{}
	if err := r.Get(ctx, req.NamespacedName, sandbox); err != nil {
		if k8serrors.IsNotFound(err) {
			log.Info("sandbox resource not found. Ignoring since object must be deleted")
			return ctrl.Result{}, nil
		}
		return ctrl.Result{}, err
	}

	// Start Tracing Span
	initialAttrs := map[string]string{
		"sandbox.name":      sandbox.Name,
		"sandbox.namespace": sandbox.Namespace,
	}
	ctx, end := r.Tracer.StartSpan(ctx, sandbox, "ReconcileSandbox", initialAttrs)
	defer end()

	// If the sandbox is being deleted, do nothing
	if !sandbox.DeletionTimestamp.IsZero() {
		log.Info("Sandbox is being deleted")
		return ctrl.Result{}, nil
	}

	// Check if already marked as expired to avoid repeated operations, including cleanups
	if sandboxMarkedExpired(sandbox) {
		log.Info("Sandbox is already marked as expired")
		// Note: The sandbox won't be deleted if shutdown policy is changed to delete after expiration.
		//       To delete an expired sandbox, the user should delete the sandbox instead of updating it.
		//       This keeps the controller code simple.
		return ctrl.Result{}, nil
	}

	// Initialize trace ID for active resources missing an ID
	tc := r.Tracer.GetTraceContext(ctx)
	if tc != "" && (sandbox.Annotations == nil || sandbox.Annotations[asmetrics.TraceContextAnnotation] == "") {
		patch := client.MergeFrom(sandbox.DeepCopy())
		if sandbox.Annotations == nil {
			sandbox.Annotations = make(map[string]string)
		}
		sandbox.Annotations[asmetrics.TraceContextAnnotation] = tc

		if err := r.Patch(ctx, sandbox, patch); err != nil {
			return ctrl.Result{}, err
		}
		// Return to ensure the next loop uses the persisted ID
		return ctrl.Result{}, nil
	}

	if sandbox.Spec.Replicas == nil {
		replicas := int32(1)
		sandbox.Spec.Replicas = &replicas
	}

	oldStatus := sandbox.Status.DeepCopy()
	var err error
	sandboxDeleted := false

	expired, requeueAfter := checkSandboxExpiry(sandbox)

	// Check if sandbox has expired
	if expired {
		log.Info("Sandbox has expired, deleting child resources and checking shutdown policy")
		sandboxDeleted, err = r.handleSandboxExpiry(ctx, sandbox)
	} else {
		err = r.reconcileChildResources(ctx, sandbox)
	}

	if !sandboxDeleted {
		// Update status
		if statusUpdateErr := r.updateStatus(ctx, oldStatus, sandbox); statusUpdateErr != nil {
			// Surface update error
			err = errors.Join(err, statusUpdateErr)
		}
	}
	// return errors seen
	return ctrl.Result{RequeueAfter: requeueAfter}, err
}

func (r *SandboxReconciler) reconcileChildResources(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox) error {
	// Create a hash from the sandbox.Name and use it as label value
	nameHash := NameHash(sandbox.Name)

	var allErrors error

	// Reconcile PVCs
	err := r.reconcilePVCs(ctx, sandbox)
	allErrors = errors.Join(allErrors, err)

	// Reconcile Pod
	pod, err := r.reconcilePod(ctx, sandbox, nameHash)
	allErrors = errors.Join(allErrors, err)
	if pod == nil {
		sandbox.Status.Replicas = 0
		sandbox.Status.LabelSelector = ""
	} else {
		sandbox.Status.Replicas = 1
		sandbox.Status.LabelSelector = fmt.Sprintf("%s=%s", sandboxLabel, NameHash(sandbox.Name))
	}

	// Reconcile Service
	svc, err := r.reconcileService(ctx, sandbox, nameHash)
	allErrors = errors.Join(allErrors, err)

	// compute and set overall Ready condition
	readyCondition := r.computeReadyCondition(sandbox, allErrors, svc, pod)
	meta.SetStatusCondition(&sandbox.Status.Conditions, readyCondition)

	return allErrors
}

func (r *SandboxReconciler) computeReadyCondition(sandbox *sandboxv1alpha1.Sandbox, err error, svc *corev1.Service, pod *corev1.Pod) metav1.Condition {
	readyCondition := metav1.Condition{
		Type:               string(sandboxv1alpha1.SandboxConditionReady),
		ObservedGeneration: sandbox.Generation,
		Message:            "",
		Status:             metav1.ConditionFalse,
		Reason:             "DependenciesNotReady",
	}

	if err != nil {
		readyCondition.Reason = "ReconcilerError"
		readyCondition.Message = "Error seen: " + err.Error()
		return readyCondition
	}

	message := ""
	podReady := false
	if pod != nil {
		message = "Pod exists with phase: " + string(pod.Status.Phase)
		// Check if pod Ready condition is true
		if pod.Status.Phase == corev1.PodRunning {
			message = "Pod is Running but not Ready"
			for _, condition := range pod.Status.Conditions {
				if condition.Type == corev1.PodReady {
					if condition.Status == corev1.ConditionTrue {
						message = "Pod is Ready"
						podReady = true
					}
					break
				}
			}
		}
	} else {
		if sandbox.Spec.Replicas != nil && *sandbox.Spec.Replicas == 0 {
			message = "Pod does not exist, replicas is 0"
			// This is intended behaviour. So marking it ready.
			podReady = true
		} else {
			message = "Pod does not exist"
		}
	}

	svcReady := false
	if svc != nil {
		message += "; Service Exists"
		svcReady = true
	} else {
		message += "; Service does not exist"
	}

	readyCondition.Message = message
	if podReady && svcReady {
		readyCondition.Status = metav1.ConditionTrue
		readyCondition.Reason = "DependenciesReady"
	}

	return readyCondition
}

func (r *SandboxReconciler) updateStatus(ctx context.Context, oldStatus *sandboxv1alpha1.SandboxStatus, sandbox *sandboxv1alpha1.Sandbox) error {
	log := log.FromContext(ctx)

	if reflect.DeepEqual(oldStatus, &sandbox.Status) {
		return nil
	}

	if err := r.Status().Update(ctx, sandbox); err != nil {
		log.Error(err, "Failed to update sandbox status")
		return err
	}

	// Surface error
	return nil
}

// NameHash generates an FNV-1a hash from a string and returns
// it as a fixed-length hexadecimal string.
func NameHash(objectName string) string {
	h := fnv.New32a()
	h.Write([]byte(objectName))
	hashValue := h.Sum32()

	// Convert the uint32 to a hexadecimal string.
	// This results in an 8-character string (e.g., "a5b3c2d1").
	return fmt.Sprintf("%08x", hashValue)
}

func (r *SandboxReconciler) reconcileService(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox, nameHash string) (*corev1.Service, error) {
	log := log.FromContext(ctx)
	service := &corev1.Service{}
	if err := r.Get(ctx, types.NamespacedName{Name: sandbox.Name, Namespace: sandbox.Namespace}, service); err != nil {
		if !k8serrors.IsNotFound(err) {
			log.Error(err, "Failed to get Service")
			return nil, fmt.Errorf("service get failed: %w", err)
		}
	} else {
		log.Info("Found Service", "Service.Namespace", service.Namespace, "Service.Name", service.Name)
		setServiceStatus(sandbox, service)
		return service, nil
	}

	log.Info("Creating a new Headless Service", "Service.Namespace", sandbox.Namespace, "Service.Name", sandbox.Name)
	service = &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      sandbox.Name,
			Namespace: sandbox.Namespace,
			Labels: map[string]string{
				sandboxLabel: nameHash,
			},
		},
		Spec: corev1.ServiceSpec{
			ClusterIP: "None",
			Selector: map[string]string{
				sandboxLabel: nameHash,
			},
		},
	}
	service.SetGroupVersionKind(corev1.SchemeGroupVersion.WithKind("Service"))
	if err := ctrl.SetControllerReference(sandbox, service, r.Scheme); err != nil {
		log.Error(err, "Failed to set controller reference")
		return nil, fmt.Errorf("SetControllerReference for Service failed: %w", err)
	}

	err := r.Create(ctx, service, client.FieldOwner(sandboxControllerFieldOwner))
	if err != nil {
		log.Error(err, "Failed to create", "Service.Namespace", service.Namespace, "Service.Name", service.Name)
		return nil, err
	}

	setServiceStatus(sandbox, service)
	return service, nil
}

// setServiceStatus updates the sandbox status with the service name and FQDN.
// TODO(barney-s): hardcoded to svc.cluster.local which is the default. Need a way to change it.
func setServiceStatus(sandbox *sandboxv1alpha1.Sandbox, service *corev1.Service) {
	sandbox.Status.Service = service.Name
	sandbox.Status.ServiceFQDN = service.Name + "." + service.Namespace + ".svc.cluster.local"
}

func (r *SandboxReconciler) reconcilePod(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox, nameHash string) (*corev1.Pod, error) {
	log := log.FromContext(ctx)

	// Start a child span of ReconcileSandbox
	ctx, end := r.Tracer.StartSpan(ctx, nil, "reconcilePod", nil)
	defer end()

	// List all pods with the pool label matching the warm pool name hash
	// TODO: find a better way to make sure one sandbox has at most one pod
	podList := &corev1.PodList{}
	labelSelector := labels.SelectorFromSet(labels.Set{
		sandboxLabel: nameHash,
	})

	if err := r.List(ctx, podList, &client.ListOptions{
		LabelSelector: labelSelector,
		Namespace:     sandbox.Namespace,
	}); err != nil {
		log.Error(err, "Failed to list pods")
	}

	if len(podList.Items) > 1 {
		log.Info("Multiple pods found for sandbox, this should not happen", "Sandbox", sandbox.Name, "PodCount", len(podList.Items))
	}

	// Determine the pod name to look up
	podName := sandbox.Name
	var trackedPodName string
	var podNameAnnotationExists bool
	if trackedPodName, podNameAnnotationExists = sandbox.Annotations[SandboxPodNameAnnotation]; podNameAnnotationExists && trackedPodName != "" {
		podName = trackedPodName
		log.Info("Using tracked pod name from sandbox annotation", "podName", podName)
	}

	pod := &corev1.Pod{}
	err := r.Get(ctx, types.NamespacedName{Name: podName, Namespace: sandbox.Namespace}, pod)
	if err != nil {
		if !k8serrors.IsNotFound(err) {
			log.Error(err, "Failed to get Pod")
			return nil, fmt.Errorf("pod get failed: %w", err)
		}
		if podNameAnnotationExists {
			log.Error(err, "Pod not found")
			return nil, fmt.Errorf("pod in annotation get failed: %w", err)
		}
		pod = nil
	}

	// 1. PATH: Logic for deleting Pod when replicas is 0
	if *sandbox.Spec.Replicas == 0 {
		if pod != nil {
			if pod.DeletionTimestamp.IsZero() {
				log.Info("Deleting Pod because .Spec.Replicas is 0", "Pod.Namespace", pod.Namespace, "Pod.Name", pod.Name)
				if err := r.Delete(ctx, pod); err != nil {
					return nil, fmt.Errorf("failed to delete pod: %w", err)
				}
			} else {
				log.Info("Pod is already being deleted", "Pod.Namespace", pod.Namespace, "Pod.Name", pod.Name)
			}
		}

		// Remove the pod name annotation from the sandbox if it exists
		if _, exists := sandbox.Annotations[SandboxPodNameAnnotation]; exists {
			log.Info("Removing pod name annotation from sandbox", "Sandbox.Name", sandbox.Name)
			// Create a patch to update only the annotations
			patch := client.MergeFrom(sandbox.DeepCopy())
			delete(sandbox.Annotations, SandboxPodNameAnnotation)

			if err := r.Patch(ctx, sandbox, patch); err != nil {
				return nil, fmt.Errorf("failed to remove pod name annotation: %w", err)
			}
		}

		return nil, nil
	}

	// 2. PATH: Existing Pod found (e.g., adopted from WarmPool or already exists)
	if pod != nil {
		log.Info("Found Pod", "Pod.Namespace", pod.Namespace, "Pod.Name", pod.Name)

		if r.Tracer.IsRecording(ctx) {
			r.Tracer.AddEvent(ctx, "ExistingPodStatusObserved", map[string]string{
				"pod.Name":  pod.Name,
				"pod.Phase": string(pod.Status.Phase),
			})
		}

		if pod.Labels == nil {
			pod.Labels = make(map[string]string)
		}
		pod.Labels[sandboxLabel] = nameHash

		// Set controller reference if the pod is not controlled by anything.
		if controllerRef := metav1.GetControllerOf(pod); controllerRef == nil {
			if err := ctrl.SetControllerReference(sandbox, pod, r.Scheme); err != nil {
				return nil, fmt.Errorf("SetControllerReference for Pod failed: %w", err)
			}
		}

		if err := r.Update(ctx, pod); err != nil {
			return nil, fmt.Errorf("failed to update pod: %w", err)
		}

		// TODO - Do we enfore (change) spec if a pod exists ?
		// r.Patch(ctx, pod, client.Apply, client.ForceOwnership, client.FieldOwner("sandbox-controller"))
		return pod, nil
	}

	// 3. PATH: Create new Pod
	log.Info("Creating a new Pod", "Pod.Namespace", sandbox.Namespace, "Pod.Name", sandbox.Name)
	labels := map[string]string{
		sandboxLabel: nameHash,
	}
	for k, v := range sandbox.Spec.PodTemplate.ObjectMeta.Labels {
		labels[k] = v
	}
	annotations := map[string]string{}
	for k, v := range sandbox.Spec.PodTemplate.ObjectMeta.Annotations {
		annotations[k] = v
	}

	mutatedSpec := sandbox.Spec.PodTemplate.Spec.DeepCopy()

	for _, pvcTemplate := range sandbox.Spec.VolumeClaimTemplates {
		pvcName := pvcTemplate.Name + "-" + sandbox.Name
		mutatedSpec.Volumes = append(mutatedSpec.Volumes, corev1.Volume{
			Name: pvcTemplate.Name,
			VolumeSource: corev1.VolumeSource{
				PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
					ClaimName: pvcName,
				},
			},
		})
	}
	pod = &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:        sandbox.Name,
			Namespace:   sandbox.Namespace,
			Labels:      labels,
			Annotations: annotations,
		},
		Spec: *mutatedSpec,
	}
	pod.SetGroupVersionKind(corev1.SchemeGroupVersion.WithKind("Pod"))
	if err := ctrl.SetControllerReference(sandbox, pod, r.Scheme); err != nil {
		return nil, fmt.Errorf("SetControllerReference for Pod failed: %w", err)
	}
	if err := r.Create(ctx, pod, client.FieldOwner(sandboxControllerFieldOwner)); err != nil {
		log.Error(err, "Failed to create", "Pod.Namespace", pod.Namespace, "Pod.Name", pod.Name)
		return nil, err
	}

	if r.Tracer.IsRecording(ctx) {
		r.Tracer.AddEvent(ctx, "NewPodStatusObserved", map[string]string{
			"pod.Name":  pod.Name,
			"pod.Phase": string(pod.Status.Phase),
		})
	}

	return pod, nil
}

func (r *SandboxReconciler) reconcilePVCs(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox) error {
	log := log.FromContext(ctx)

	// Start a child span of ReconcileSandbox
	ctx, end := r.Tracer.StartSpan(ctx, nil, "reconcilePVCs", nil)
	defer end()

	for _, pvcTemplate := range sandbox.Spec.VolumeClaimTemplates {
		pvc := &corev1.PersistentVolumeClaim{}
		pvcName := pvcTemplate.Name + "-" + sandbox.Name
		err := r.Get(ctx, types.NamespacedName{Name: pvcName, Namespace: sandbox.Namespace}, pvc)
		if err == nil {
			continue
		}

		if !k8serrors.IsNotFound(err) {
			log.Error(err, "Failed to get PVC")
			return fmt.Errorf("PVC Get Failed: %w", err)
		}

		log.Info("Creating a new PVC", "PVC.Namespace", sandbox.Namespace, "PVC.Name", pvcName)
		pvc = &corev1.PersistentVolumeClaim{
			ObjectMeta: metav1.ObjectMeta{
				Name:      pvcName,
				Namespace: sandbox.Namespace,
			},
			Spec: pvcTemplate.Spec,
		}
		if err := ctrl.SetControllerReference(sandbox, pvc, r.Scheme); err != nil {
			return fmt.Errorf("SetControllerReference for PVC failed: %w", err)
		}
		if err := r.Create(ctx, pvc, client.FieldOwner(sandboxControllerFieldOwner)); err != nil {
			log.Error(err, "Failed to create PVC", "PVC.Namespace", sandbox.Namespace, "PVC.Name", pvcName)
			return err
		}
	}
	return nil
}

// handles sandbox expiry by deleting child resources and the sandbox itself if needed
func (r *SandboxReconciler) handleSandboxExpiry(ctx context.Context, sandbox *sandboxv1alpha1.Sandbox) (bool, error) {
	var allErrors error
	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      sandbox.Name,
			Namespace: sandbox.Namespace,
		},
	}
	if err := r.Delete(ctx, pod); err != nil && !k8serrors.IsNotFound(err) {
		allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete pod: %w", err))
	}

	service := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      sandbox.Name,
			Namespace: sandbox.Namespace,
		},
	}
	if err := r.Delete(ctx, service); err != nil && !k8serrors.IsNotFound(err) {
		allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete service: %w", err))
	}

	if sandbox.Spec.ShutdownPolicy != nil && *sandbox.Spec.ShutdownPolicy == sandboxv1alpha1.ShutdownPolicyDelete {
		if err := r.Delete(ctx, sandbox); err != nil && !k8serrors.IsNotFound(err) {
			allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete sandbox: %w", err))
		} else {
			return true, nil
		}
	}

	// If we reach here, sandbox is not deleted
	// Only update "expired" status if cleanup was successful
	if allErrors == nil {
		// Clear all status fields explicitly
		sandbox.Status = sandboxv1alpha1.SandboxStatus{}
		// Update status to mark as expired
		meta.SetStatusCondition(&sandbox.Status.Conditions, metav1.Condition{
			Type:               string(sandboxv1alpha1.SandboxConditionReady),
			Status:             metav1.ConditionFalse,
			ObservedGeneration: sandbox.Generation,
			Reason:             sandboxv1alpha1.SandboxReasonExpired,
			Message:            "Sandbox has expired",
		})
	}

	return false, allErrors
}

// checks if the sandbox has expired
// returns true if expired, false otherwise
// if not expired, also returns the duration to requeue after
func checkSandboxExpiry(sandbox *sandboxv1alpha1.Sandbox) (bool, time.Duration) {
	if sandbox.Spec.ShutdownTime == nil {
		return false, 0
	}

	expiryTime := sandbox.Spec.ShutdownTime.Time
	if time.Now().After(expiryTime) {
		return true, 0
	}

	// Calculate remaining time
	remainingTime := time.Until(expiryTime)

	// TODO(barney-s): Do we need a inverse exponential backoff here ?
	//requeueAfter := max(remainingTime/2, 2*time.Second)

	// Requeue at expiry time or in 2 seconds whichever is later
	requeueAfter := max(remainingTime, 2*time.Second)
	return false, requeueAfter
}

// sandboxMarkedExpired checks if the sandbox is already marked as expired
func sandboxMarkedExpired(sandbox *sandboxv1alpha1.Sandbox) bool {
	cond := meta.FindStatusCondition(sandbox.Status.Conditions, string(sandboxv1alpha1.SandboxConditionReady))
	return cond != nil && cond.Reason == sandboxv1alpha1.SandboxReasonExpired
}

// SandboxCollector is a custom Prometheus collector that dynamically fetches sandbox counts.
type SandboxCollector struct {
	client             client.Client
	agentSandboxesDesc *prometheus.Desc
}

// NewSandboxCollector initializes a SandboxCollector.
func NewSandboxCollector(c client.Client) *SandboxCollector {
	return &SandboxCollector{
		client:             c,
		agentSandboxesDesc: asmetrics.AgentSandboxesDesc,
	}
}

// Describe sends the metric descriptor to the channel.
func (c *SandboxCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.agentSandboxesDesc
}

// Collect fetches sandboxes, calculates labels, and sends metrics to the channel.
func (c *SandboxCollector) Collect(ch chan<- prometheus.Metric) {
	var sandboxList sandboxv1alpha1.SandboxList
	ctx, cancel := context.WithTimeout(context.Background(), metricsCollectTimeout)
	defer cancel()
	if err := c.client.List(ctx, &sandboxList); err != nil {
		log.FromContext(ctx).Error(err, "Failed to list sandboxes for metrics collection")
		return
	}
	counts := make(map[asmetrics.AgentSandboxesMetricKey]int)
	for _, sandbox := range sandboxList.Items {
		readyConditionStr := "false"
		expiredStr := "false"
		readyCond := meta.FindStatusCondition(sandbox.Status.Conditions, string(sandboxv1alpha1.SandboxConditionReady))
		if readyCond != nil {
			if readyCond.Status == metav1.ConditionTrue {
				readyConditionStr = "true"
			}
			if readyCond.Reason == sandboxv1alpha1.SandboxReasonExpired || readyCond.Reason == extensionsv1alpha1.ClaimExpiredReason {
				expiredStr = "true"
			}
		}
		launchTypeStr := asmetrics.LaunchTypeCold
		if _, ok := sandbox.Annotations[SandboxPodNameAnnotation]; ok && sandbox.Annotations[SandboxPodNameAnnotation] != "" {
			launchTypeStr = asmetrics.LaunchTypeWarm
		}
		sandboxTemplateStr := "unknown"
		if template, ok := sandbox.Annotations[SandboxTemplateRefAnnotation]; ok && template != "" {
			sandboxTemplateStr = template
		}
		key := asmetrics.AgentSandboxesMetricKey{
			ReadyCondition: readyConditionStr,
			Expired:        expiredStr,
			LaunchType:     launchTypeStr,
			Template:       sandboxTemplateStr,
		}
		counts[key]++
	}
	for key, count := range counts {
		ch <- asmetrics.NewAgentSandboxesConstMetric(count, key)
	}
}

// SetupWithManager sets up the controller with the Manager.
func (r *SandboxReconciler) SetupWithManager(mgr ctrl.Manager, concurrentWorkers int) error {
	labelSelectorPredicate, err := predicate.LabelSelectorPredicate(metav1.LabelSelector{
		MatchExpressions: []metav1.LabelSelectorRequirement{
			{
				Key:      sandboxLabel,
				Operator: metav1.LabelSelectorOpExists,
				Values:   []string{},
			},
		},
	})
	if err != nil {
		return err
	}

	// Register the custom Sandbox metric collector
	ctrlMetrics.Registry.MustRegister(NewSandboxCollector(mgr.GetClient()))

	return ctrl.NewControllerManagedBy(mgr).
		For(&sandboxv1alpha1.Sandbox{}).
		Owns(&corev1.Pod{}, builder.WithPredicates(labelSelectorPredicate)).
		Owns(&corev1.Service{}, builder.WithPredicates(labelSelectorPredicate)).
		WithOptions(controller.Options{MaxConcurrentReconciles: concurrentWorkers}).
		Complete(r)
}