package commontasks

import (
	"context"
	"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 TestDrainNodeTask_Create(t *testing.T) {
	tests := []struct {
		name           string
		annotations    map[string]string
		setupResources []client.Object
		wantError      bool
		wantDone       bool
		wantAnnotation string // empty string means annotation should be removed
		verify         func(t *testing.T, cli client.Client)
	}{
		{
			name: "no drain annotation, return Done",
			annotations: map[string]string{
				"other-key": "other-value",
			},
			wantError:      false,
			wantDone:       true,
			wantAnnotation: "",
		},
		{
			name: "invalid annotation format - no slash",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "invalid-format",
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "invalid annotation format - too many parts",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "ns/pod/extra",
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "invalid annotation - empty namespace",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "/pod-name",
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "invalid annotation - empty pod name",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "namespace/",
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "invalid annotation - both empty",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "/",
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "pod not found",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "default/non-existent-pod",
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "pod not scheduled to any node",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "default/pending-pod",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "pending-pod",
						Namespace: "default",
					},
					Spec: corev1.PodSpec{
						NodeName: "", // Pod not scheduled
					},
				},
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "node not found",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "default/orphaned-pod",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "orphaned-pod",
						Namespace: "default",
					},
					Spec: corev1.PodSpec{
						NodeName: "non-existent-node",
					},
				},
			},
			wantError:      true,
			wantDone:       false,
			wantAnnotation: "",
		},
		{
			name: "successfully apply out-of-service taint to node",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "default/test-pod",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "test-pod",
						Namespace: "default",
					},
					Spec: corev1.PodSpec{
						NodeName: "test-node",
					},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-node",
					},
				},
			},
			wantError:      false,
			wantDone:       true,
			wantAnnotation: "",
			verify: func(t *testing.T, cli client.Client) {
				node := &corev1.Node{}
				err := cli.Get(context.Background(), types.NamespacedName{Name: "test-node"}, node)
				assert.NoError(t, err)

				found := false
				for _, taint := range node.Spec.Taints {
					if taint.Key == "node.kubernetes.io/out-of-service" &&
						taint.Effect == corev1.TaintEffectNoExecute {
						found = true
						break
					}
				}
				assert.True(t, found, "out-of-service taint should be applied")
			},
		},
		{
			name: "node already has out-of-service taint",
			annotations: map[string]string{
				consts.AnnotationKeyO11yDrainNode: "default/test-pod-2",
			},
			setupResources: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "test-pod-2",
						Namespace: "default",
					},
					Spec: corev1.PodSpec{
						NodeName: "already-tainted-node",
					},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "already-tainted-node",
					},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{
							{
								Key:    "node.kubernetes.io/out-of-service",
								Effect: corev1.TaintEffectNoExecute,
							},
						},
					},
				},
			},
			wantError:      false,
			wantDone:       true,
			wantAnnotation: "",
			verify: func(t *testing.T, cli client.Client) {
				node := &corev1.Node{}
				err := cli.Get(context.Background(), types.NamespacedName{Name: "already-tainted-node"}, node)
				assert.NoError(t, err)

				// Count taints - should still be just 1
				assert.Len(t, node.Spec.Taints, 1, "should not duplicate taint")
			},
		},
	}

	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),
			}

			cm := deploy.ClusterMeta{Agent: "test-agent"}
			task := NewDrainNodeTask(reconcilerCtx, cm)
			assert.Equal(t, "DrainNode", task.Name())
			assert.Equal(t, cm, task.ClusterMeta())
			assert.True(t, task.Delete(context.Background()).IsDoneState())

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

			_, exists := agent.Annotations[consts.AnnotationKeyO11yDrainNode]
			assert.False(t, exists, "drain annotation should be removed after processing")
			if tt.verify != nil {
				tt.verify(t, fakeCli)
			}
		})
	}
}