package dynamo

import (
	"strings"
	"testing"

	"github.com/ai-dynamo/dynamo/deploy/cloud/operator/api/v1alpha1"
	"github.com/ai-dynamo/dynamo/deploy/cloud/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.DynamoComponentDeploymentOverridesSpec
		multinodeDeploymentType consts.MultinodeDeploymentType
		initialArgs             []string
		initialLivenessProbe    *corev1.Probe
		initialReadinessProbe   *corev1.Probe
		initialStartupProbe     *corev1.Probe
		expectedArgs            []string
		expectContains          []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.DynamoComponentDeploymentOverridesSpec{},
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{"python3", "-m", "dynamo.vllm"},
			expectNotModified:       true,
		},
		{
			name:                    "multinode leader prepends ray start --head",
			numberOfNodes:           3,
			role:                    RoleLeader,
			component:               &v1alpha1.DynamoComponentDeploymentOverridesSpec{},
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{"python3", "-m", "dynamo.vllm", "--model", "test"},
			expectContains:          []string{"ray start --head --port=6379 &&", "python3", "-m", "dynamo.vllm", "--model", "test"},
			expectProbesRemoved:     true,
		},
		{
			name:                    "multinode worker replaces args with ray start --block",
			numberOfNodes:           3,
			role:                    RoleWorker,
			component:               &v1alpha1.DynamoComponentDeploymentOverridesSpec{},
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{"python3", "-m", "dynamo.vllm", "--model", "test"},
			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",
			numberOfNodes:           2,
			role:                    RoleWorker,
			component:               &v1alpha1.DynamoComponentDeploymentOverridesSpec{},
			multinodeDeploymentType: consts.MultinodeDeploymentTypeLWS,
			initialArgs:             []string{"python3", "-m", "dynamo.vllm"},
			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.DynamoComponentDeploymentOverridesSpec{},
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{},
			expectNotModified:       true, // Should not modify empty args
		},
		{
			name:                    "multinode main role (non-leader/worker) does not modify args",
			numberOfNodes:           3,
			role:                    RoleMain,
			component:               &v1alpha1.DynamoComponentDeploymentOverridesSpec{},
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{"python3", "-m", "dynamo.frontend"},
			expectNotModified:       true,
		},
	}

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

			// Create a container with initial state
			container := &corev1.Container{
				Args:           append([]string{}, tt.initialArgs...), // Copy slice to avoid modifying original
				LivenessProbe:  tt.initialLivenessProbe,
				ReadinessProbe: tt.initialReadinessProbe,
				StartupProbe:   tt.initialStartupProbe,
			}

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

			if tt.expectNotModified {
				// Args should not have changed
				g.Expect(container.Args).To(gomega.Equal(tt.initialArgs))
			} else if tt.expectedArgs != nil {
				// Check exact match
				g.Expect(container.Args).To(gomega.Equal(tt.expectedArgs))
			} else if tt.expectContains != nil {
				// Check that expected strings are contained in the result
				argsStr := strings.Join(container.Args, " ")
				for _, expected := range tt.expectContains {
					if !strings.Contains(argsStr, expected) {
						t.Errorf("UpdateContainer() args = %v, should contain %s", container.Args, expected)
					}
				}
			}

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

func TestUpdateVLLMMultinodeArgs(t *testing.T) {
	tests := []struct {
		name                    string
		role                    Role
		multinodeDeploymentType consts.MultinodeDeploymentType
		initialArgs             []string
		expectedArgs            []string
		expectContains          []string
		expectNotModified       bool
	}{
		{
			name:                    "leader prepends ray start --head",
			role:                    RoleLeader,
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{"python3", "-m", "dynamo.vllm"},
			expectContains:          []string{"ray start --head --port=6379 &&", "python3", "-m", "dynamo.vllm"},
		},
		{
			name:                    "leader with empty args does not modify",
			role:                    RoleLeader,
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{},
			expectNotModified:       true,
		},
		{
			name:                    "worker with Grove deployment",
			role:                    RoleWorker,
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{"python3", "-m", "dynamo.vllm"},
			expectedArgs:            []string{"ray start --address=${GROVE_PCSG_NAME}-${GROVE_PCSG_INDEX}-test-service-ldr-0.${GROVE_HEADLESS_SERVICE}:6379 --block"},
		},
		{
			name:                    "worker with LWS deployment",
			role:                    RoleWorker,
			multinodeDeploymentType: consts.MultinodeDeploymentTypeLWS,
			initialArgs:             []string{"python3", "-m", "dynamo.vllm"},
			expectedArgs:            []string{"ray start --address=${LWS_LEADER_ADDRESS}:6379 --block"},
		},
		{
			name:                    "main role does not modify args",
			role:                    RoleMain,
			multinodeDeploymentType: consts.MultinodeDeploymentTypeGrove,
			initialArgs:             []string{"python3", "-m", "dynamo.frontend"},
			expectNotModified:       true,
		},
	}

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

			// Create a container with initial args
			container := &corev1.Container{
				Args: append([]string{}, tt.initialArgs...), // Copy slice to avoid modifying original
			}

			// Call updateVLLMMultinodeArgs
			updateVLLMMultinodeArgs(container, tt.role, tt.multinodeDeploymentType, "test-service")

			if tt.expectNotModified {
				// Args should not have changed
				g.Expect(container.Args).To(gomega.Equal(tt.initialArgs))
			} else if tt.expectedArgs != nil {
				// Check exact match
				g.Expect(container.Args).To(gomega.Equal(tt.expectedArgs))
			} else if tt.expectContains != nil {
				// Check that expected strings are contained in the result
				argsStr := strings.Join(container.Args, " ")
				for _, expected := range tt.expectContains {
					if !strings.Contains(argsStr, expected) {
						t.Errorf("updateVLLMMultinodeArgs() args = %v, should contain %s", container.Args, expected)
					}
				}
			}
		})
	}
}
