package dynamo

import (
	"fmt"
	"reflect"
	"strconv"
	"testing"

	"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
	"github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
)

func TestVLLMBackend_UpdateContainer(t *testing.T) {
	backend := &VLLMBackend{}

	tests := []struct {
		name                string
		numberOfNodes       int32
		role                Role
		component           *v1alpha1.DynamoComponentDeploymentSharedSpec
		multinodeDeployer   MultinodeDeployer
		initialContainer    *corev1.Container
		gpuCount            int64 // GPU count for the test case
		expectedArgs        []string
		expectNotModified   bool // If true, container args should not change
		expectProbesRemoved bool // If true, probes should be nil
	}{
		{
			name:              "single node does not modify args",
			numberOfNodes:     1,
			role:              RoleMain,
			component:         &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm"}},
			gpuCount:          0,
			expectNotModified: true,
		},
		{
			name:                "multinode leader uses ray (no annotations = legacy)",
			numberOfNodes:       3,
			role:                RoleLeader,
			component:           &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer:   &GroveMultinodeDeployer{},
			initialContainer:    &corev1.Container{Command: []string{"python3", "-m", "dynamo.vllm"}, Args: []string{"--model", "test", tensorParallelSizeFlag, "8"}},
			gpuCount:            4,
			expectedArgs:        []string{fmt.Sprintf("ray start --head --port=%s && python3 -m dynamo.vllm --model test %s 8 --distributed-executor-backend ray", VLLMPort, tensorParallelSizeFlag)},
			expectProbesRemoved: true,
		},
		{
			name:              "multinode leader uses ray with JSON args (no annotations = legacy)",
			numberOfNodes:     3,
			role:              RoleLeader,
			component:         &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer: &corev1.Container{
				Command: []string{"python3", "-m", "dynamo.vllm"},
				Args: []string{
					"--model", "test", tensorParallelSizeFlag, "8",
					"--kv-transfer-config",
					`{"kv_connector": "NixlConnector", "kv_role": "kv_both"}`,
				},
			},
			gpuCount: 4,
			expectedArgs: []string{fmt.Sprintf(
				`ray start --head --port=%s && python3 -m dynamo.vllm --model test %s 8 --kv-transfer-config "{\"kv_connector\": \"NixlConnector\", \"kv_role\": \"kv_both\"}" --distributed-executor-backend ray`,
				VLLMPort, tensorParallelSizeFlag,
			)},
			expectProbesRemoved: true,
		},
		{
			name:                "multinode worker uses ray (no annotations = legacy)",
			numberOfNodes:       3,
			role:                RoleWorker,
			component:           &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer:   &GroveMultinodeDeployer{},
			initialContainer:    &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", "--model", "test", tensorParallelSizeFlag, "8"}},
			gpuCount:            4,
			expectedArgs:        []string{"ray start --address=$(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE):6379 --block"},
			expectProbesRemoved: true,
		},
		{
			name:                "multinode worker with LWS deployment type (no annotations = legacy ray)",
			numberOfNodes:       2,
			role:                RoleWorker,
			component:           &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer:   &LWSMultinodeDeployer{},
			initialContainer:    &corev1.Container{Args: []string{"python3", "-m", "dynamo.vllm", tensorParallelSizeFlag, "8"}},
			gpuCount:            4,
			expectedArgs:        []string{"ray start --address=$(LWS_LEADER_ADDRESS):6379 --block"},
			expectProbesRemoved: true,
		},
		{
			name:              "multinode leader with no initial args",
			numberOfNodes:     2,
			role:              RoleLeader,
			component:         &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Args: []string{}},
			gpuCount:          0,
			expectNotModified: true, // Should not modify empty args
		},
		{
			name:              "multinode main role (non-leader/worker) does not modify args",
			numberOfNodes:     3,
			role:              RoleMain,
			component:         &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Args: []string{"python3", "-m", "dynamo.frontend"}},
			gpuCount:          0,
			expectNotModified: true,
		},
		{
			name:          "multinode leader uses mp (origin version >= threshold)",
			numberOfNodes: 2,
			role:          RoleLeader,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
				},
			},
			multinodeDeployer:   &GroveMultinodeDeployer{},
			initialContainer:    &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:            8,
			expectedArgs:        []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16", "--distributed-executor-backend", "mp", "--nnodes", "2", "--master-addr", "$(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE)", "--master-port", commonconsts.VLLMMpMasterPort, "--node-rank", "0"},
			expectProbesRemoved: true,
		},
		{
			name:          "multinode worker uses mp (origin version >= threshold) Grove",
			numberOfNodes: 2,
			role:          RoleWorker,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
				},
			},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			expectedArgs: []string{fmt.Sprintf(
				"exec python3 -m dynamo.vllm %s 16 --distributed-executor-backend mp --nnodes 2 --master-addr $(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE) --master-port %s --node-rank $((GROVE_PCLQ_POD_INDEX + 1)) --headless",
				tensorParallelSizeFlag, commonconsts.VLLMMpMasterPort)},
			expectProbesRemoved: true,
		},
		{
			name:          "multinode leader uses ray (explicit override despite new version)",
			numberOfNodes: 2,
			role:          RoleLeader,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion:    "1.0.0",
					commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "ray",
				},
			},
			multinodeDeployer:   &GroveMultinodeDeployer{},
			initialContainer:    &corev1.Container{Command: []string{"python3", "-m", "dynamo.vllm"}, Args: []string{"--model", "test", tensorParallelSizeFlag, "8"}},
			gpuCount:            4,
			expectedArgs:        []string{fmt.Sprintf("ray start --head --port=%s && python3 -m dynamo.vllm --model test %s 8 --distributed-executor-backend ray", VLLMPort, tensorParallelSizeFlag)},
			expectProbesRemoved: true,
		},
		{
			name:          "multinode leader uses mp (explicit override on legacy DGD)",
			numberOfNodes: 2,
			role:          RoleLeader,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "mp",
				},
			},
			multinodeDeployer:   &GroveMultinodeDeployer{},
			initialContainer:    &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:            8,
			expectedArgs:        []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16", "--distributed-executor-backend", "mp", "--nnodes", "2", "--master-addr", "$(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE)", "--master-port", commonconsts.VLLMMpMasterPort, "--node-rank", "0"},
			expectProbesRemoved: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := gomega.NewGomegaWithT(t)

			initialContainerArgs := append([]string{}, tt.initialContainer.Args...)

			// Create resources from GPU count and set in component
			if tt.gpuCount > 0 {
				tt.component.Resources = &v1alpha1.Resources{
					Limits: &v1alpha1.ResourceItem{
						GPU: strconv.FormatInt(tt.gpuCount, 10),
					},
				}
			}

			// Call UpdateContainer
			backend.UpdateContainer(tt.initialContainer, tt.numberOfNodes, tt.role, tt.component, "test-service", tt.multinodeDeployer)

			if tt.expectNotModified {
				// Args should not have changed
				g.Expect(tt.initialContainer.Args).To(gomega.Equal(initialContainerArgs))
			} else if tt.expectedArgs != nil {
				// Check exact match
				g.Expect(tt.initialContainer.Args).To(gomega.Equal(tt.expectedArgs))
			}

			if tt.expectProbesRemoved {
				g.Expect(tt.initialContainer.LivenessProbe).To(gomega.BeNil())
				g.Expect(tt.initialContainer.ReadinessProbe).To(gomega.BeNil())
				g.Expect(tt.initialContainer.StartupProbe).To(gomega.BeNil())
			}
		})
	}
}

func TestVLLMBackend_ShellCommandInjection(t *testing.T) {
	backend := &VLLMBackend{}

	tests := []struct {
		name              string
		numberOfNodes     int32
		role              Role
		multinodeDeployer MultinodeDeployer
		initialContainer  *corev1.Container
		gpuCount          int64 // GPU count for the test case
		expectedArgs      []string
		description       string
	}{
		{
			name:              "single node shell command not modified",
			numberOfNodes:     1,
			role:              RoleMain,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"sh", "-c"}, Args: []string{"python3 -m dynamo.vllm"}},
			gpuCount:          0,
			expectedArgs:      []string{"python3 -m dynamo.vllm"},
			description:       "Single node should not modify shell commands",
		},
		{
			name:              "multinode shell command with regex injection",
			numberOfNodes:     2,
			role:              RoleLeader,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"sh", "-c"}, Args: []string{fmt.Sprintf("python3 -m dynamo.vllm %s 8", dataParallelSizeFlag)}},
			gpuCount:          4,
			expectedArgs:      []string{"python3 -m dynamo.vllm --data-parallel-hybrid-lb --data-parallel-size-local 4 --data-parallel-start-rank 0 --data-parallel-address $(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE) --data-parallel-rpc-port 13445 --data-parallel-size 8"},
			description:       "Shell commands should use regex injection for python commands",
		},
		{
			name:              "multinode shell command with complex pipeline",
			numberOfNodes:     2,
			role:              RoleLeader,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"sh", "-c"}, Args: []string{fmt.Sprintf("echo blah | wc -l && python3 -m dynamo.vllm %s 8 && ls -al", dataParallelSizeFlag)}},
			gpuCount:          4,
			expectedArgs:      []string{"echo blah | wc -l && python3 -m dynamo.vllm --data-parallel-hybrid-lb --data-parallel-size-local 4 --data-parallel-start-rank 0 --data-parallel-address $(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE) --data-parallel-rpc-port 13445 --data-parallel-size 8 && ls -al"},
			description:       "Complex shell commands should inject flags only into python part",
		},
		{
			name:              "shell command with LWS deployer",
			numberOfNodes:     2,
			role:              RoleLeader,
			multinodeDeployer: &LWSMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"sh", "-c"}, Args: []string{fmt.Sprintf("python3 -m dynamo.vllm %s 8", dataParallelSizeFlag)}},
			gpuCount:          4,
			expectedArgs:      []string{"python3 -m dynamo.vllm --data-parallel-hybrid-lb --data-parallel-size-local 4 --data-parallel-start-rank 0 --data-parallel-address $(LWS_LEADER_ADDRESS) --data-parallel-rpc-port 13445 --data-parallel-size 8"},
			description:       "LWS shell commands should use LWS variables",
		},
		{
			name:              "shell command with pipes",
			numberOfNodes:     2,
			role:              RoleLeader,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"sh", "-c"}, Args: []string{fmt.Sprintf("python3 -m dynamo.vllm %s 8 | tee /tmp/log", dataParallelSizeFlag)}},
			gpuCount:          4,
			expectedArgs:      []string{"python3 -m dynamo.vllm --data-parallel-hybrid-lb --data-parallel-size-local 4 --data-parallel-start-rank 0 --data-parallel-address $(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE) --data-parallel-rpc-port 13445 --data-parallel-size 8 | tee /tmp/log"},
			description:       "Shell commands with pipes should inject flags before pipe",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			expectedCommand := append([]string{}, tt.initialContainer.Command...)

			// Create component with resources from GPU count
			component := &v1alpha1.DynamoComponentDeploymentSharedSpec{}
			if tt.gpuCount > 0 {
				component.Resources = &v1alpha1.Resources{
					Limits: &v1alpha1.ResourceItem{
						GPU: strconv.FormatInt(tt.gpuCount, 10),
					},
				}
			}

			backend.UpdateContainer(tt.initialContainer, tt.numberOfNodes, tt.role, component, "test-service", tt.multinodeDeployer)

			if !reflect.DeepEqual(tt.initialContainer.Args, tt.expectedArgs) {
				t.Errorf("UpdateContainer() args = %v, want %v", tt.initialContainer.Args, tt.expectedArgs)
			}

			if !reflect.DeepEqual(tt.initialContainer.Command, expectedCommand) {
				t.Errorf("UpdateContainer() should preserve shell command, got: %v, want: %v", tt.initialContainer.Command, expectedCommand)
			}
		})
	}
}

func TestVLLMBackend_UpdateContainer_UseAsCompilationCache(t *testing.T) {
	backend := &VLLMBackend{}

	tests := []struct {
		name                  string
		component             *v1alpha1.DynamoComponentDeploymentSharedSpec
		volumeMounts          []corev1.VolumeMount
		expectCacheEnvVar     bool
		expectCacheEnvVarName string
		expectCacheEnvVarVal  string
	}{
		{
			name: "VLLM backend with useAsCompilationCache volume mount",
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				VolumeMounts: []v1alpha1.VolumeMount{
					{
						Name:                  "vllm-cache",
						MountPoint:            "/root/.cache/vllm",
						UseAsCompilationCache: true,
					},
				},
			},
			volumeMounts:          []corev1.VolumeMount{},
			expectCacheEnvVar:     true,
			expectCacheEnvVarName: "VLLM_CACHE_ROOT",
			expectCacheEnvVarVal:  "/root/.cache/vllm",
		},
		{
			name: "VLLM backend with useAsCompilationCache at custom mount point",
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				VolumeMounts: []v1alpha1.VolumeMount{
					{
						Name:                  "custom-cache",
						MountPoint:            "/custom/cache/path",
						UseAsCompilationCache: true,
					},
				},
			},
			volumeMounts:          []corev1.VolumeMount{},
			expectCacheEnvVar:     true,
			expectCacheEnvVarName: "VLLM_CACHE_ROOT",
			expectCacheEnvVarVal:  "/custom/cache/path",
		},
		{
			name: "VLLM backend without useAsCompilationCache",
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				VolumeMounts: []v1alpha1.VolumeMount{
					{
						Name:       "regular-volume",
						MountPoint: "/data",
					},
				},
			},
			volumeMounts:      []corev1.VolumeMount{},
			expectCacheEnvVar: false,
		},
		{
			name: "VLLM backend with no volume mounts",
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				VolumeMounts: nil,
			},
			volumeMounts:      []corev1.VolumeMount{},
			expectCacheEnvVar: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := gomega.NewGomegaWithT(t)

			// Create a container with initial state including volume mounts
			container := &corev1.Container{
				Env:          []corev1.EnvVar{},
				VolumeMounts: tt.volumeMounts,
			}

			// Call UpdateContainer
			backend.UpdateContainer(container, 1, RoleMain, tt.component, "test-service", &GroveMultinodeDeployer{})

			if tt.expectCacheEnvVar {
				// Check that the VLLM_CACHE_ROOT environment variable is set
				found := false
				for _, env := range container.Env {
					if env.Name == tt.expectCacheEnvVarName {
						found = true
						g.Expect(env.Value).To(gomega.Equal(tt.expectCacheEnvVarVal))
						break
					}
				}
				if !found {
					t.Errorf("Expected environment variable %s not found in container", tt.expectCacheEnvVarName)
				}
			} else {
				// Check that no cache environment variable is set
				for _, env := range container.Env {
					if env.Name == "VLLM_CACHE_ROOT" {
						t.Errorf("Unexpected environment variable VLLM_CACHE_ROOT found: %s", env.Value)
					}
				}
			}
		})
	}
}

func TestUpdateVLLMMultinodeArgs(t *testing.T) {
	tests := []struct {
		name              string
		role              Role
		multinodeDeployer MultinodeDeployer
		initialContainer  *corev1.Container
		gpuCount          int64
		annotations       map[string]string // nil = legacy (no annotations)
		expectedArgs      []string
		expectNotModified bool
	}{
		{
			name:              "leader uses ray (nil annotations = legacy)",
			role:              RoleLeader,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations:       nil,
			expectedArgs:      []string{fmt.Sprintf("ray start --head --port=%s && python3 -m dynamo.vllm %s 16 --distributed-executor-backend ray", VLLMPort, tensorParallelSizeFlag)},
		},
		{
			name:              "leader uses mp (origin version >= threshold)",
			role:              RoleLeader,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
			},
			expectedArgs: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16", "--distributed-executor-backend", "mp", "--nnodes", "2", "--master-addr", "$(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE)", "--master-port", commonconsts.VLLMMpMasterPort, "--node-rank", "0"},
		},
		{
			name:              "worker uses mp (origin version >= threshold) Grove",
			role:              RoleWorker,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
			},
			expectedArgs: []string{fmt.Sprintf(
				"exec python3 -m dynamo.vllm %s 16 --distributed-executor-backend mp --nnodes 2 --master-addr $(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE) --master-port %s --node-rank $((GROVE_PCLQ_POD_INDEX + 1)) --headless",
				tensorParallelSizeFlag, commonconsts.VLLMMpMasterPort)},
		},
		{
			// LWS worker: $(LWS_LEADER_ADDRESS) and $(LWS_WORKER_INDEX) are both
			// kubelet-expanded, so flags are appended directly to Args without an
			// sh -c wrapper.
			name:              "worker uses mp (origin version >= threshold) LWS",
			role:              RoleWorker,
			multinodeDeployer: &LWSMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
			},
			expectedArgs: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16", "--distributed-executor-backend", "mp", "--nnodes", "2", "--master-addr", "$(LWS_LEADER_ADDRESS)", "--master-port", commonconsts.VLLMMpMasterPort, "--node-rank", "$(LWS_WORKER_INDEX)", "--headless"},
		},
		{
			// Regression test: LWS leader with direct python command must emit
			// Kubernetes $(LWS_LEADER_ADDRESS) syntax so the kubelet expands it
			// from the LWS-injected env var. Emitting the bare shell $VAR causes
			// vLLM to receive the literal string and fail to resolve the leader.
			name:              "leader uses mp (origin version >= threshold) LWS",
			role:              RoleLeader,
			multinodeDeployer: &LWSMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
			},
			expectedArgs: []string{"-m", "dynamo.vllm", tensorParallelSizeFlag, "16", "--distributed-executor-backend", "mp", "--nnodes", "2", "--master-addr", "$(LWS_LEADER_ADDRESS)", "--master-port", commonconsts.VLLMMpMasterPort, "--node-rank", "0"},
		},
		{
			// Regression test: LWS leader on the data-parallel path. Same bug
			// class as the MP leader case above - bare $LWS_LEADER_ADDRESS would
			// not be expanded by K8s, so we emit $(LWS_LEADER_ADDRESS) instead.
			name:              "leader with data parallel launch LWS",
			role:              RoleLeader,
			multinodeDeployer: &LWSMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", dataParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations:       nil,
			expectedArgs:      []string{"-m", "dynamo.vllm", dataParallelSizeFlag, "16", "--data-parallel-hybrid-lb", "--data-parallel-size-local", "8", "--data-parallel-start-rank", "0", "--data-parallel-address", "$(LWS_LEADER_ADDRESS)", "--data-parallel-rpc-port", "13445"},
		},
		{
			name:              "leader prepends distributed data parallel flags (annotations don't affect DP path)",
			role:              RoleLeader,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", dataParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations:       nil,
			expectedArgs:      []string{"-m", "dynamo.vllm", dataParallelSizeFlag, "16", "--data-parallel-hybrid-lb", "--data-parallel-size-local", "8", "--data-parallel-start-rank", "0", "--data-parallel-address", "$(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE)", "--data-parallel-rpc-port", "13445"},
		},
		{
			name:              "leader with empty args does not modify",
			role:              RoleLeader,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Args: []string{}},
			gpuCount:          0,
			annotations:       nil,
			expectNotModified: true,
		},
		{
			name:              "worker with ray distributed launch Grove (nil annotations)",
			role:              RoleWorker,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Args: []string{"python3", "-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations:       nil,
			expectedArgs:      []string{"ray start --address=$(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE):6379 --block"},
		},
		{
			name:              "worker with data parallel launch Grove",
			role:              RoleWorker,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", dataParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations:       nil,
			expectedArgs:      []string{fmt.Sprintf("exec python3 -m dynamo.vllm %s 16 --data-parallel-hybrid-lb --data-parallel-size-local 8 --data-parallel-start-rank $(( 8 * $((GROVE_PCLQ_POD_INDEX + 1)) )) --data-parallel-address $(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE) --data-parallel-rpc-port 13445", dataParallelSizeFlag)},
		},
		{
			name:              "worker with data parallel launch Grove, tp > 1",
			role:              RoleWorker,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm", dataParallelSizeFlag, "8", tensorParallelSizeFlag, "2"}},
			gpuCount:          8,
			annotations:       nil,
			expectedArgs:      []string{fmt.Sprintf("exec python3 -m dynamo.vllm %s 8 %s 2 --data-parallel-hybrid-lb --data-parallel-size-local 4 --data-parallel-start-rank $(( 4 * $((GROVE_PCLQ_POD_INDEX + 1)) )) --data-parallel-address $(GROVE_PCSG_NAME)-$(GROVE_PCSG_INDEX)-test-service-ldr-0.$(GROVE_HEADLESS_SERVICE) --data-parallel-rpc-port 13445", dataParallelSizeFlag, tensorParallelSizeFlag)},
		},
		{
			name:              "worker with ray distributed launch LWS (nil annotations)",
			role:              RoleWorker,
			multinodeDeployer: &LWSMultinodeDeployer{},
			initialContainer:  &corev1.Container{Args: []string{"python3", "-m", "dynamo.vllm", tensorParallelSizeFlag, "16"}},
			gpuCount:          8,
			annotations:       nil,
			expectedArgs:      []string{"ray start --address=$(LWS_LEADER_ADDRESS):6379 --block"},
		},
		{
			name:              "main role does not modify args",
			role:              RoleMain,
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialContainer:  &corev1.Container{Args: []string{"python3", "-m", "dynamo.frontend"}},
			gpuCount:          0,
			annotations:       nil,
			expectNotModified: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := gomega.NewGomegaWithT(t)

			initialContainerArgs := append([]string{}, tt.initialContainer.Args...)

			// Create resources from GPU count
			var resources *v1alpha1.Resources
			if tt.gpuCount > 0 {
				resources = &v1alpha1.Resources{
					Limits: &v1alpha1.ResourceItem{
						GPU: strconv.FormatInt(tt.gpuCount, 10),
					},
				}
			}

			// Call updateVLLMMultinodeArgs with annotations
			updateVLLMMultinodeArgs(tt.initialContainer, tt.role, "test-service", tt.multinodeDeployer, resources, 2, tt.annotations)

			if tt.expectNotModified {
				// Args should not have changed
				g.Expect(tt.initialContainer.Args).To(gomega.Equal(initialContainerArgs))
			} else if tt.expectedArgs != nil {
				// Check exact match
				g.Expect(tt.initialContainer.Args).To(gomega.Equal(tt.expectedArgs))
			}
		})
	}
}

func TestVLLMBackend_UpdatePodSpec(t *testing.T) {
	backend := &VLLMBackend{ParentGraphDeploymentName: "test-dgd"}

	tests := []struct {
		name                string
		numberOfNodes       int32
		role                Role
		component           *v1alpha1.DynamoComponentDeploymentSharedSpec
		multinodeDeployer   MultinodeDeployer
		initialPodSpec      *corev1.PodSpec
		expectInitContainer bool
		expectedInitImage   string
		expectedLeaderHost  string
	}{
		{
			name:          "mp worker with Grove deployer injects init container",
			numberOfNodes: 2,
			role:          RoleWorker,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
				},
			},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialPodSpec: &corev1.PodSpec{
				Containers: []corev1.Container{
					{Name: "main", Image: "vllm:latest"},
				},
			},
			expectInitContainer: true,
			expectedInitImage:   "vllm:latest",
			expectedLeaderHost:  "${GROVE_PCSG_NAME}-${GROVE_PCSG_INDEX}-test-service-ldr-0.${GROVE_HEADLESS_SERVICE}",
		},
		{
			name:          "mp worker with LWS deployer injects init container",
			numberOfNodes: 2,
			role:          RoleWorker,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
				},
			},
			multinodeDeployer: &LWSMultinodeDeployer{},
			initialPodSpec: &corev1.PodSpec{
				Containers: []corev1.Container{
					{Name: "main", Image: "vllm:v2"},
				},
			},
			expectInitContainer: true,
			expectedInitImage:   "vllm:v2",
			expectedLeaderHost:  "${LWS_LEADER_ADDRESS}",
		},
		{
			name:          "mp leader does not inject init container",
			numberOfNodes: 2,
			role:          RoleLeader,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
				},
			},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialPodSpec: &corev1.PodSpec{
				Containers: []corev1.Container{
					{Name: "main", Image: "vllm:latest"},
				},
			},
			expectInitContainer: false,
		},
		{
			name:              "ray worker does not inject init container (legacy)",
			numberOfNodes:     2,
			role:              RoleWorker,
			component:         &v1alpha1.DynamoComponentDeploymentSharedSpec{},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialPodSpec: &corev1.PodSpec{
				Containers: []corev1.Container{
					{Name: "main", Image: "vllm:latest"},
				},
			},
			expectInitContainer: false,
		},
		{
			name:          "single node does not inject init container",
			numberOfNodes: 1,
			role:          RoleMain,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
				},
			},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialPodSpec: &corev1.PodSpec{
				Containers: []corev1.Container{
					{Name: "main", Image: "vllm:latest"},
				},
			},
			expectInitContainer: false,
		},
		{
			name:          "mp worker preserves existing init containers",
			numberOfNodes: 2,
			role:          RoleWorker,
			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{
				Annotations: map[string]string{
					commonconsts.KubeAnnotationDynamoOperatorOriginVersion: "1.0.0",
				},
			},
			multinodeDeployer: &GroveMultinodeDeployer{},
			initialPodSpec: &corev1.PodSpec{
				InitContainers: []corev1.Container{
					{Name: "existing-init", Image: "busybox"},
				},
				Containers: []corev1.Container{
					{Name: "main", Image: "vllm:latest"},
				},
			},
			expectInitContainer: true,
			expectedInitImage:   "vllm:latest",
			expectedLeaderHost:  "${GROVE_PCSG_NAME}-${GROVE_PCSG_INDEX}-test-service-ldr-0.${GROVE_HEADLESS_SERVICE}",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := gomega.NewGomegaWithT(t)

			initialInitCount := len(tt.initialPodSpec.InitContainers)
			initialVolCount := len(tt.initialPodSpec.Volumes)
			backend.UpdatePodSpec(tt.initialPodSpec, tt.numberOfNodes, tt.role, tt.component, "test-service", tt.multinodeDeployer)

			if tt.expectInitContainer {
				g.Expect(tt.initialPodSpec.InitContainers).To(gomega.HaveLen(initialInitCount + 1))
				g.Expect(tt.initialPodSpec.Volumes).To(gomega.HaveLen(initialVolCount + 1))

				injected := tt.initialPodSpec.InitContainers[len(tt.initialPodSpec.InitContainers)-1]
				g.Expect(injected.Name).To(gomega.Equal("wait-for-leader-mp"))
				g.Expect(injected.Image).To(gomega.Equal(tt.expectedInitImage))

				expectedCmd := fmt.Sprintf(
					`export LEADER_HOST="%s" LEADER_PORT="%s" && exec python3 /scripts/wait-for-leader.py`,
					tt.expectedLeaderHost, commonconsts.VLLMMpMasterPort)
				g.Expect(injected.Command).To(gomega.Equal([]string{"sh", "-c", expectedCmd}))
				g.Expect(injected.Env).To(gomega.BeEmpty())

				g.Expect(injected.VolumeMounts).To(gomega.HaveLen(1))
				g.Expect(injected.VolumeMounts[0].Name).To(gomega.Equal("wait-leader-script"))
				g.Expect(injected.VolumeMounts[0].MountPath).To(gomega.Equal("/scripts"))
				g.Expect(injected.VolumeMounts[0].ReadOnly).To(gomega.BeTrue())

				vol := tt.initialPodSpec.Volumes[len(tt.initialPodSpec.Volumes)-1]
				g.Expect(vol.Name).To(gomega.Equal("wait-leader-script"))
				g.Expect(vol.ConfigMap).ToNot(gomega.BeNil())
				g.Expect(vol.ConfigMap.Name).To(gomega.Equal("test-dgd-wait-leader-script"))
			} else {
				g.Expect(tt.initialPodSpec.InitContainers).To(gomega.HaveLen(initialInitCount))
				g.Expect(tt.initialPodSpec.Volumes).To(gomega.HaveLen(initialVolCount))
			}
		})
	}
}

func TestGenerateWaitLeaderConfigMap(t *testing.T) {
	g := gomega.NewGomegaWithT(t)

	cm := GenerateWaitLeaderConfigMap("my-dgd", "my-ns")

	g.Expect(cm.Name).To(gomega.Equal("my-dgd-wait-leader-script"))
	g.Expect(cm.Namespace).To(gomega.Equal("my-ns"))
	g.Expect(cm.Labels).To(gomega.HaveKeyWithValue(commonconsts.KubeLabelDynamoGraphDeploymentName, "my-dgd"))
	g.Expect(cm.Data).To(gomega.HaveKey("wait-for-leader.py"))

	script := cm.Data["wait-for-leader.py"]
	g.Expect(script).To(gomega.ContainSubstring(`os.environ["LEADER_HOST"]`))
	g.Expect(script).To(gomega.ContainSubstring(`os.environ["LEADER_PORT"]`))
	g.Expect(script).To(gomega.ContainSubstring("leader_pod_is_healthy"))
	g.Expect(script).To(gomega.ContainSubstring("kubernetes.default.svc"))
	g.Expect(script).To(gomega.ContainSubstring("fieldSelector=status.podIP="))
	g.Expect(script).To(gomega.ContainSubstring("deletionTimestamp"))
	g.Expect(script).To(gomega.ContainSubstring("socket.create_connection"))
	g.Expect(script).To(gomega.ContainSubstring("time.sleep(5)"))
}

func TestGetWaitLeaderConfigMapName(t *testing.T) {
	g := gomega.NewGomegaWithT(t)
	g.Expect(GetWaitLeaderConfigMapName("foo")).To(gomega.Equal("foo-wait-leader-script"))
}

func TestShouldUseMpBackend(t *testing.T) {
	// Version-based gate behavior is tested in featuregate.TestOperatorOriginFeatureGate_IsEnabled.
	// These tests focus on the explicit override logic and its interaction with the feature gate.
	tests := []struct {
		name        string
		annotations map[string]string
		want        bool
	}{
		{
			name:        "nil annotations = legacy = ray (delegates to feature gate)",
			annotations: nil,
			want:        false,
		},
		{
			name:        "empty annotations = legacy = ray (delegates to feature gate)",
			annotations: map[string]string{},
			want:        false,
		},
		{
			name: "explicit override mp takes priority over version",
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion:    "0.1.0",
				commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "mp",
			},
			want: true,
		},
		{
			name: "explicit override ray takes priority over version",
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion:    "1.0.0",
				commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "ray",
			},
			want: false,
		},
		{
			name: "explicit override mp (no origin version)",
			annotations: map[string]string{
				commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "mp",
			},
			want: true,
		},
		{
			name: "explicit override with invalid value falls through to feature gate",
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion:    "1.0.0",
				commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "invalid",
			},
			want: true, // invalid override ignored, version >= threshold via feature gate
		},
		{
			name: "explicit override case insensitive MP",
			annotations: map[string]string{
				commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "MP",
			},
			want: true,
		},
		{
			name: "explicit override case insensitive Ray",
			annotations: map[string]string{
				commonconsts.KubeAnnotationDynamoOperatorOriginVersion:    "1.0.0",
				commonconsts.KubeAnnotationVLLMDistributedExecutorBackend: "RAY",
			},
			want: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := shouldUseMpBackend(tt.annotations)
			if got != tt.want {
				t.Errorf("shouldUseMpBackend() = %v, want %v", got, tt.want)
			}
		})
	}
}
