package commontasks

import (
	"context"
	"fmt"
	"testing"

	"github.com/go-logr/logr"
	"github.com/stretchr/testify/assert"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/deploy"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/testkits"
)

func TestRestartTask_Create(t *testing.T) {
	tests := []struct {
		name            string
		annotations     map[string]string
		setupResources  []client.Object
		enableVector    bool
		enableVMAgent   bool
		namespace       string
		wantError       bool
		wantAnnotations map[string]string // expected annotations after processing
		verify          func(t *testing.T, cli client.Client)
	}{
		{
			name: "no restart annotations, return Done",
			annotations: map[string]string{
				"other-key": "other-value",
			},
			enableVector:  true,
			enableVMAgent: true,
			wantError:     false,
			wantAnnotations: map[string]string{
				"other-key": "other-value",
			},
		},
		{
			name: "only Vector restart annotation, delete Vector pods",
			annotations: map[string]string{
				consts.AnnotationKeyO11yVectorRestart: "true",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vector-pod-1",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vector",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vector-pod-2",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vector",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
			},
			enableVector:    true,
			enableVMAgent:   false,
			wantError:       false,
			wantAnnotations: map[string]string{},
			verify: func(t *testing.T, cli client.Client) {
				pod1 := &corev1.Pod{}
				err := cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: "vector-pod-1"}, pod1)
				assert.Error(t, err, "pod should be deleted")

				pod2 := &corev1.Pod{}
				err = cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: "vector-pod-2"}, pod2)
				assert.Error(t, err, "pod should be deleted")
			},
		},
		{
			name: "only VMAgent restart annotation, delete VMAgent pods",
			annotations: map[string]string{
				consts.AnnotationKeyO11yVMAgentRestart: "true",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vmagent-pod-1",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vmagent",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
			},
			enableVector:    false,
			enableVMAgent:   true,
			wantError:       false,
			wantAnnotations: map[string]string{},
			verify: func(t *testing.T, cli client.Client) {
				pod := &corev1.Pod{}
				err := cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: "vmagent-pod-1"}, pod)
				assert.Error(t, err, "pod should be deleted")
			},
		},
		{
			name: "both Vector and VMAgent restart annotations, delete all pods",
			annotations: map[string]string{
				consts.AnnotationKeyO11yVectorRestart:  "true",
				consts.AnnotationKeyO11yVMAgentRestart: "true",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vector-pod",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vector",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vmagent-pod",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vmagent",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
			},
			enableVector:    true,
			enableVMAgent:   true,
			namespace:       "monitoring",
			wantError:       false,
			wantAnnotations: map[string]string{},
			verify: func(t *testing.T, cli client.Client) {
				vectorPod := &corev1.Pod{}
				err := cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: "vector-pod"}, vectorPod)
				assert.Error(t, err, "pod should be deleted")

				vmagentPod := &corev1.Pod{}
				err = cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: "vmagent-pod"}, vmagentPod)
				assert.Error(t, err, "pod should be deleted")
			},
		},
		{
			name: "multiple pods with same labels",
			annotations: map[string]string{
				consts.AnnotationKeyO11yVectorRestart: "true",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vector-pod-1",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vector",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vector-pod-2",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vector",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vector-pod-3",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vector",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
			},
			enableVector:    true,
			enableVMAgent:   false,
			namespace:       "monitoring",
			wantError:       false,
			wantAnnotations: map[string]string{},
			verify: func(t *testing.T, cli client.Client) {
				for i := 1; i <= 3; i++ {
					pod := &corev1.Pod{}
					name := fmt.Sprintf("vector-pod-%d", i)
					err := cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: name}, pod)
					assert.Error(t, err, fmt.Sprintf("pod %s should be deleted", name))
				}
			},
		},
		{
			name: "Vector annotation present but Vector disabled, should not process",
			annotations: map[string]string{
				consts.AnnotationKeyO11yVectorRestart: "true",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vector-pod",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vector",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
			},
			enableVector:  false, // Vector disabled
			enableVMAgent: true,
			wantError:     false,
			wantAnnotations: map[string]string{
				consts.AnnotationKeyO11yVectorRestart: "true", // annotation should remain
			},
			verify: func(t *testing.T, cli client.Client) {
				pod := &corev1.Pod{}
				err := cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: "vector-pod"}, pod)
				assert.NoError(t, err, "pod should not be deleted when Vector is disabled")
			},
		},
		{
			name: "VMAgent annotation present but VMAgent disabled, should not process",
			annotations: map[string]string{
				consts.AnnotationKeyO11yVMAgentRestart: "true",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "vmagent-pod",
						Namespace: "monitoring",
						Labels: map[string]string{
							"controller/component": "vmagent",
							"controller/agent":     consts.AgentKubernetes,
						},
					},
				},
			},
			enableVector:  true,
			enableVMAgent: false, // VMAgent disabled
			wantError:     false,
			wantAnnotations: map[string]string{
				consts.AnnotationKeyO11yVMAgentRestart: "true", // annotation should remain
			},
			verify: func(t *testing.T, cli client.Client) {
				pod := &corev1.Pod{}
				err := cli.Get(context.Background(), types.NamespacedName{Namespace: "monitoring", Name: "vmagent-pod"}, pod)
				assert.NoError(t, err, "pod should not be deleted when VMAgent is disabled")
			},
		},
		{
			name: "pods do not exist, should succeed without error",
			annotations: map[string]string{
				consts.AnnotationKeyO11yVectorRestart:  "true",
				consts.AnnotationKeyO11yVMAgentRestart: "true",
			},
			setupResources:  []client.Object{}, // No pods
			enableVector:    true,
			enableVMAgent:   true,
			wantError:       false,
			wantAnnotations: map[string]string{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			fakeCli := testkits.BuildClient(tt.setupResources...)
			agent := &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name:        "test-agent",
					Namespace:   "default",
					Annotations: tt.annotations,
				},
			}

			reconcilerCtx := &reconciler.Context[client.Object]{
				Agent:  agent,
				L:      logr.Discard(),
				Remote: testkits.NewClientGetter(fakeCli),
			}

			namespace := tt.namespace
			if namespace == "" {
				namespace = "monitoring"
			}
			task := &RestartTask[client.Object]{
				BaseTask:      &reconciler.BaseTask[client.Object]{Context: reconcilerCtx},
				cm:            deploy.ClusterMeta{Agent: "test-agent"},
				agentType:     consts.AgentKubernetes,
				namespace:     namespace,
				enableVector:  tt.enableVector,
				enableVMAgent: tt.enableVMAgent,
			}

			assert.Equal(t, "Restart", task.Name())
			assert.Equal(t, deploy.ClusterMeta{Agent: "test-agent"}, task.ClusterMeta())

			result := task.Create(context.Background())
			if tt.wantError {
				assert.Error(t, result.Error())
			} else {
				assert.NoError(t, result.Error())
			}

			assert.Equal(t, tt.wantAnnotations, agent.Annotations, "annotations should match expected")
			if tt.verify != nil {
				tt.verify(t, fakeCli)
			}
		})
	}
}