package task

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	appsv1 "k8s.io/api/apps/v1"
	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"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/pkg/config"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/testkits"
)

func makeVectorTask(agent *agentv2.O11yK8sAgent, opts testkits.TaskOptions) *VectorTask {
	objs := []client.Object{agent}
	if len(opts.LocalObjects) > 0 {
		objs = append(objs, opts.LocalObjects...)
	}
	localCli := testkits.BuildClient(objs...)

	ctx := &reconciler.Context[*agentv2.O11yK8sAgent]{
		Agent:     agent,
		Local:     localCli,
		Remote:    testkits.NewClientGetterWithOpts(opts),
		Config:    &opts.Config,
		Templates: opts.Templates,
	}
	return NewVectorTask(ctx)
}

func baseVectorConfig() config.RootConfig {
	return config.RootConfig{
		O11yConfig: config.ObservabilityInfra{
			Vendor: "aws",
			CustomerConfigsV2: config.CustomerConfigsV2{
				Dedicated: config.CustomerDedicated{
					K8s: config.CustomerDedicatedK8s{
						Default: config.CustomerAgents{
							Vector: &config.VectorConfig{
								Chart: &config.CustomerHelmChart{
									Name:    "vector",
									Version: "1.0.0",
								},
								Image:     "vectorio/vector:latest",
								Resources: &corev1.ResourceRequirements{},
							},
						},
					},
				},
			},
		},
	}
}

func TestVectorTask(t *testing.T) {
	tests := []struct {
		name   string
		verify func(t *testing.T)
	}{
		{
			name: "Name",
			verify: func(t *testing.T) {
				task := &VectorTask{}
				assert.Equal(t, "Vector", task.Name())
			},
		},
		{
			name: "ClusterMeta",
			verify: func(t *testing.T) {
				agent := testkits.NewO11yK8sAgentBuilder().Build()
				task := makeVectorTask(agent, testkits.TaskOptions{})
				meta := task.ClusterMeta()

				assert.Equal(t, consts.AgentO11yK8s, meta.Agent)
				assert.Equal(t, "test-cluster", meta.K8sClusterName)
				assert.Equal(t, "tenant-1", meta.TenantID)
				assert.Equal(t, "app-1", meta.ApplicationID)
				assert.Equal(t, "tidbcloud", meta.BizType)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tt.verify(t)
		})
	}
}

func TestVectorTask_Create(t *testing.T) {
	tests := []struct {
		name    string
		agent   *agentv2.O11yK8sAgent
		opts    testkits.TaskOptions
		wantErr bool
		verify  func(t *testing.T, opts testkits.TaskOptions)
	}{
		{
			name:    "create fails with remote client error",
			agent:   testkits.NewO11yK8sAgentBuilder().Build(),
			opts:    testkits.TaskOptions{ClientError: true, Config: baseVectorConfig()},
			wantErr: true,
		},
		{
			name:    "create fails with missing template",
			agent:   testkits.NewO11yK8sAgentBuilder().Build(),
			opts:    testkits.TaskOptions{Client: testkits.BuildClient(), Config: baseVectorConfig()},
			wantErr: true,
		},
		{
			name: "create handles service account annotation update",
			agent: testkits.NewO11yK8sAgentBuilder().
				WithAnnotations(map[string]string{consts.AnnotationKeyO11yAgentRoleArn: "arn:aws:iam::123456789:role/writer-role"}).
				Build(),
			opts: func() testkits.TaskOptions {
				ns := &corev1.Namespace{ObjectMeta: metav1.ObjectMeta{Name: consts.O11ySystemNamespace}}
				sa := &corev1.ServiceAccount{ObjectMeta: metav1.ObjectMeta{
					Name:      consts.O11yNameVector,
					Namespace: consts.O11ySystemNamespace,
					Annotations: map[string]string{
						EKSRoleAnnotation: "arn:aws:iam::123456789:role/old-role",
					},
				}}
				return testkits.TaskOptions{
					Client: testkits.BuildClient(ns, sa),
					Config: baseVectorConfig(),
				}
			}(),
			wantErr: true,
			verify: func(t *testing.T, opts testkits.TaskOptions) {
				updatedSA := &corev1.ServiceAccount{}
				err := opts.Client.Get(context.Background(), types.NamespacedName{
					Name:      consts.O11yNameVector,
					Namespace: consts.O11ySystemNamespace,
				}, updatedSA)
				assert.NoError(t, err)
				_, exists := updatedSA.Annotations[EKSRoleAnnotation]
				assert.False(t, exists)
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			task := makeVectorTask(tt.agent, tt.opts)
			result := task.Create(context.Background())
			if tt.wantErr {
				assert.Error(t, result.Error())
			} else {
				assert.NoError(t, result.Error())
			}
			if tt.verify != nil {
				tt.verify(t, tt.opts)
			}
		})
	}
}

func TestVectorTask_PostCreate(t *testing.T) {
	tests := []struct {
		name    string
		agent   *agentv2.O11yK8sAgent
		opts    testkits.TaskOptions
		wantErr bool
	}{
		{
			name:    "post create fails with remote client error",
			agent:   testkits.NewO11yK8sAgentBuilder().Build(),
			opts:    testkits.TaskOptions{ClientError: true},
			wantErr: true,
		},
		{
			name:    "post create fails when daemonset not ready",
			agent:   testkits.NewO11yK8sAgentBuilder().Build(),
			opts:    testkits.TaskOptions{Client: testkits.BuildClient()},
			wantErr: true,
		},
		{
			name:  "post create succeeds when daemonset ready",
			agent: testkits.NewO11yK8sAgentBuilder().Build(),
			opts: testkits.TaskOptions{
				Client: testkits.BuildClient(
					&appsv1.DaemonSet{
						ObjectMeta: metav1.ObjectMeta{
							Name:      consts.O11yNameVector,
							Namespace: consts.O11ySystemNamespace,
						},
						Spec:   appsv1.DaemonSetSpec{},
						Status: appsv1.DaemonSetStatus{DesiredNumberScheduled: 1, NumberReady: 1},
					},
				),
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			task := makeVectorTask(tt.agent, tt.opts)
			result := task.PostCreate(context.Background())
			if tt.wantErr {
				assert.Error(t, result.Error())
			} else {
				assert.NoError(t, result.Error())
			}
		})
	}
}

func TestVectorTask_Delete(t *testing.T) {
	tests := []struct {
		name    string
		agent   *agentv2.O11yK8sAgent
		opts    testkits.TaskOptions
		wantErr bool
	}{
		{
			name:    "delete fails with helm error",
			agent:   testkits.NewO11yK8sAgentBuilder().Build(),
			opts:    testkits.TaskOptions{HelmError: true, Config: baseVectorConfig()},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			task := makeVectorTask(tt.agent, tt.opts)
			result := task.Delete(context.Background())
			if tt.wantErr {
				assert.Error(t, result.Error())
			} else {
				assert.NoError(t, result.Error())
			}
		})
	}
}

func TestVectorTask_PostDelete(t *testing.T) {
	tests := []struct {
		name    string
		agent   *agentv2.O11yK8sAgent
		opts    testkits.TaskOptions
		wantErr bool
	}{
		{
			name:    "post delete fails with remote client error",
			agent:   testkits.NewO11yK8sAgentBuilder().Build(),
			opts:    testkits.TaskOptions{ClientError: true},
			wantErr: true,
		},
		{
			name:    "post delete succeeds when daemonset deleted",
			agent:   testkits.NewO11yK8sAgentBuilder().Build(),
			opts:    testkits.TaskOptions{Client: testkits.BuildClient()},
			wantErr: false,
		},
		{
			name:  "post delete fails when daemonset still exists",
			agent: testkits.NewO11yK8sAgentBuilder().Build(),
			opts: testkits.TaskOptions{
				Client: testkits.BuildClient(
					&appsv1.DaemonSet{
						ObjectMeta: metav1.ObjectMeta{
							Name:      consts.O11yNameVector,
							Namespace: consts.O11ySystemNamespace,
						},
					},
				),
			},
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			task := makeVectorTask(tt.agent, tt.opts)
			result := task.PostDelete(context.Background())
			if tt.wantErr {
				assert.Error(t, result.Error())
			} else {
				assert.NoError(t, result.Error())
			}
		})
	}
}