package commontasks

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/o11y"
	"github.com/tidbcloud/observability/agent-control-server/pkg/o11y/api"
	"github.com/tidbcloud/observability/agent-control-server/pkg/o11y/mockapi"
)

func makeBackendTask(peeringNLB string, annotations map[string]string) *BackendTask {
	agent := &agentv2.KubernetesAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name:        "test-agent",
			Namespace:   "default",
			Annotations: annotations,
		},
	}
	return &BackendTask{Agent: agent, PeeringNLB: peeringNLB}
}

func makeBackendTaskWithAPI(peeringNLB string, annotations map[string]string, mockAPI *mockapi.MockO11yAPI) *BackendTask {
	backend := makeBackendTask(peeringNLB, annotations)
	backend.O11yAPI = mockAPI
	return backend
}

func TestResolveServiceDNSName(t *testing.T) {
	tests := []struct {
		name       string
		peeringNLB string
		backendDNS *string
		wantDNS    string
		wantOK     bool
	}{
		{
			name:       "peering mode overrides backend DNS name",
			peeringNLB: "internal-nlb.example.com",
			backendDNS: ptr.To("backend.example.com"),
			wantDNS:    "internal-nlb.example.com",
			wantOK:     true,
		},
		{
			name:       "peering mode resolves even when backend DNS is nil",
			peeringNLB: "internal-nlb.example.com",
			backendDNS: nil,
			wantDNS:    "internal-nlb.example.com",
			wantOK:     true,
		},
		{
			name:       "peering mode resolves even when backend DNS is empty",
			peeringNLB: "internal-nlb.example.com",
			backendDNS: ptr.To(""),
			wantDNS:    "internal-nlb.example.com",
			wantOK:     true,
		},
		{
			name:       "non-peering mode uses backend DNS name",
			peeringNLB: "",
			backendDNS: ptr.To("backend.example.com"),
			wantDNS:    "backend.example.com",
			wantOK:     true,
		},
		{
			name:       "non-peering mode fails when backend DNS is nil",
			peeringNLB: "",
			backendDNS: nil,
			wantDNS:    "",
			wantOK:     false,
		},
		{
			name:       "non-peering mode fails when backend DNS is empty",
			peeringNLB: "",
			backendDNS: ptr.To(""),
			wantDNS:    "",
			wantOK:     false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			bt := makeBackendTask(tt.peeringNLB, nil)
			gotDNS, gotOK := bt.resolveServiceDNSName(tt.backendDNS)
			assert.Equal(t, tt.wantOK, gotOK)
			assert.Equal(t, tt.wantDNS, gotDNS)
		})
	}
}

func TestGetLogStorageInfo(t *testing.T) {
	tests := []struct {
		name               string
		storages           []api.StorageInfo
		wantBucket         string
		wantStorageAccount string
		wantErr            bool
	}{
		{
			name: "found log storage",
			storages: []api.StorageInfo{
				{Category: api.METRICS, BucketName: "metrics-bucket"},
				{Category: api.LOGS, BucketName: "log-bucket"},
			},
			wantBucket: "log-bucket",
			wantErr:    false,
		},
		{
			name: "found log storage at first position",
			storages: []api.StorageInfo{
				{Category: api.LOGS, BucketName: "log-bucket"},
				{Category: api.METRICS, BucketName: "metrics-bucket"},
			},
			wantBucket: "log-bucket",
			wantErr:    false,
		},
		{
			name: "found log storage with storage account",
			storages: []api.StorageInfo{
				{Category: api.LOGS, BucketName: "log-container", StorageAccount: "mystorageaccount"},
			},
			wantBucket:         "log-container",
			wantStorageAccount: "mystorageaccount",
			wantErr:            false,
		},
		{
			name:       "no log storage",
			storages:   []api.StorageInfo{{Category: api.METRICS, BucketName: "metrics-bucket"}},
			wantBucket: "",
			wantErr:    true,
		},
		{
			name:       "empty storage list",
			storages:   []api.StorageInfo{},
			wantBucket: "",
			wantErr:    true,
		},
		{
			name:       "nil storage list",
			storages:   nil,
			wantBucket: "",
			wantErr:    true,
		},
		{
			name: "multiple log storages returns first",
			storages: []api.StorageInfo{
				{Category: api.LOGS, BucketName: "log-bucket-1"},
				{Category: api.LOGS, BucketName: "log-bucket-2"},
			},
			wantBucket: "log-bucket-1",
			wantErr:    false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotBucket, gotStorageAccount, err := getLogStorageInfo(tt.storages)
			if tt.wantErr {
				assert.Error(t, err)
				assert.Equal(t, "no log storage found", err.Error())
			} else {
				assert.NoError(t, err)
			}
			assert.Equal(t, tt.wantBucket, gotBucket)
			assert.Equal(t, tt.wantStorageAccount, gotStorageAccount)
		})
	}
}

func TestIsRunning(t *testing.T) {
	tests := []struct {
		name   string
		status string
		want   bool
	}{
		{"CREATING is running", string(api.CREATING), true},
		{"UPDATING is running", string(api.UPDATING), true},
		{"RESTARTING is running", string(api.RESTARTING), true},
		{"ACTIVE is running", string(api.ACTIVE), true},
		{"PAUSED is not running", string(api.PAUSED), false},
		{"PAUSING is not running", string(api.PAUSING), false},
		{"TO_BE_PAUSED is not running", string(api.TO_BE_PAUSED), false},
		{"DELETING is not running", string(api.DELETING), false},
		{"TO_BE_DELETED is not running", string(api.TO_BE_DELETED), false},
		{"CREATION_FAILED is not running", string(api.CREATION_FAILED), false},
		{"empty status is not running", "", false},
		{"unknown status is not running", "UNKNOWN", false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := isRunning(tt.status)
			assert.Equal(t, tt.want, got)
		})
	}
}

func TestIsPausing(t *testing.T) {
	tests := []struct {
		name   string
		status string
		want   bool
	}{
		{"PAUSING is pausing", string(api.PAUSING), true},
		{"TO_BE_PAUSED is pausing", string(api.TO_BE_PAUSED), true},
		{"PAUSED is pausing", string(api.PAUSED), true},
		{"ACTIVE is not pausing", string(api.ACTIVE), false},
		{"CREATING is not pausing", string(api.CREATING), false},
		{"UPDATING is not pausing", string(api.UPDATING), false},
		{"RESTARTING is not pausing", string(api.RESTARTING), false},
		{"DELETING is not pausing", string(api.DELETING), false},
		{"empty status is not pausing", "", false},
		{"unknown status is not pausing", "UNKNOWN", false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := isPausing(tt.status)
			assert.Equal(t, tt.want, got)
		})
	}
}

func makeAwsParams(agentAccessAccount string) o11y.Params {
	return o11y.CreateParamsNextGen(context.Background(), "aws", "us-east-1", "tidbcloud", "tenant-1", "app-1", "", agentAccessAccount, "alert-id", false)
}

func makeAzureSkipDNSParams(agentAccessAccount string) o11y.Params {
	return o11y.CreateParamsNextGen(context.Background(), "azure", "westus3", "dev-tidbcloud-nextgen", "tenant-1", "app-1", "", agentAccessAccount, "alert-id", true)
}

func TestCouldSkipCreate_PeeringNLB(t *testing.T) {
	baseAnnotations := map[string]string{
		consts.AnnotationKeyO11yID:             "cluster-1",
		consts.AnnotationKeyO11yServiceDNSName: "internal-nlb.example.com",
		consts.AnnotationKeyO11yVMAccountID:    "vm-acc",
		consts.AnnotationKeyO11yVMProjectID:    "vm-proj",
		consts.AnnotationKeyO11yLogBucketName:  "log-bucket",
		consts.AnnotationKeyO11yAgentRoleArn:   "role-arn",
	}

	tests := []struct {
		name        string
		peeringNLB  string
		annotations map[string]string
		params      o11y.Params
		want        bool
	}{
		{
			name:        "peering mode skips without endpoint service name annotation",
			peeringNLB:  "internal-nlb.example.com",
			annotations: baseAnnotations,
			params:      makeAwsParams("aws-123"),
			want:        true,
		},
		{
			name:        "peering mode skips even without agent access account",
			peeringNLB:  "internal-nlb.example.com",
			annotations: baseAnnotations,
			params:      makeAwsParams(""),
			want:        true,
		},
		{
			name:        "non-peering mode requires endpoint service name when account is set",
			peeringNLB:  "",
			annotations: baseAnnotations,
			params:      makeAwsParams("aws-123"),
			want:        false,
		},
		{
			name:       "non-peering mode skips when endpoint service name is present",
			peeringNLB: "",
			annotations: func() map[string]string {
				annos := make(map[string]string)
				for k, v := range baseAnnotations {
					annos[k] = v
				}
				annos[consts.AnnotationKeyO11yEndpointServiceName] = "epsrv-123"
				return annos
			}(),
			params: makeAwsParams("aws-123"),
			want:   true,
		},
		{
			name:       "missing o11y ID annotation should not skip",
			peeringNLB: "internal-nlb.example.com",
			annotations: func() map[string]string {
				m := make(map[string]string)
				for k, v := range baseAnnotations {
					m[k] = v
				}
				delete(m, consts.AnnotationKeyO11yID)
				return m
			}(),
			params: makeAwsParams("aws-123"),
			want:   false,
		},
		{
			name:       "missing service DNS name annotation should not skip",
			peeringNLB: "internal-nlb.example.com",
			annotations: func() map[string]string {
				m := make(map[string]string)
				for k, v := range baseAnnotations {
					m[k] = v
				}
				delete(m, consts.AnnotationKeyO11yServiceDNSName)
				return m
			}(),
			params: makeAwsParams("aws-123"),
			want:   false,
		},
		{
			name:       "missing VM account ID annotation should not skip",
			peeringNLB: "internal-nlb.example.com",
			annotations: func() map[string]string {
				m := make(map[string]string)
				for k, v := range baseAnnotations {
					m[k] = v
				}
				delete(m, consts.AnnotationKeyO11yVMAccountID)
				return m
			}(),
			params: makeAwsParams("aws-123"),
			want:   false,
		},
		{
			name:       "missing VM project ID annotation should not skip",
			peeringNLB: "internal-nlb.example.com",
			annotations: func() map[string]string {
				m := make(map[string]string)
				for k, v := range baseAnnotations {
					m[k] = v
				}
				delete(m, consts.AnnotationKeyO11yVMProjectID)
				return m
			}(),
			params: makeAwsParams("aws-123"),
			want:   false,
		},
		{
			name:       "missing log bucket name annotation should not skip",
			peeringNLB: "internal-nlb.example.com",
			annotations: func() map[string]string {
				m := make(map[string]string)
				for k, v := range baseAnnotations {
					m[k] = v
				}
				delete(m, consts.AnnotationKeyO11yLogBucketName)
				return m
			}(),
			params: makeAwsParams("aws-123"),
			want:   false,
		},
		{
			name:       "missing agent role ARN annotation should not skip",
			peeringNLB: "internal-nlb.example.com",
			annotations: func() map[string]string {
				m := make(map[string]string)
				for k, v := range baseAnnotations {
					m[k] = v
				}
				delete(m, consts.AnnotationKeyO11yAgentRoleArn)
				return m
			}(),
			params: makeAwsParams("aws-123"),
			want:   false,
		},
		{
			name:        "empty annotations should not skip",
			peeringNLB:  "internal-nlb.example.com",
			annotations: nil,
			params:      makeAwsParams("aws-123"),
			want:        false,
		},
		{
			name:        "non-peering mode with skip endpoint service name flag skips",
			peeringNLB:  "",
			annotations: baseAnnotations,
			params:      o11y.CreateParamsNextGen(context.Background(), "aws", "us-east-1", "tidbcloud", "tenant-1", "app-1", "", "aws-123", "alert-id", true),
			want:        true,
		},
		{
			name:        "non-peering mode without agent access account should not skip",
			peeringNLB:  "",
			annotations: baseAnnotations,
			params:      makeAwsParams(""),
			want:        false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			task := makeBackendTask(tt.peeringNLB, tt.annotations)
			got := task.couldSkipCreate(tt.params)
			assert.Equal(t, tt.want, got)
		})
	}
}

func TestBackendTask_Create(t *testing.T) {
	tests := []struct {
		name        string
		peeringNLB  string
		annotations map[string]string
		pause       bool
		couldSkip   bool
		mockSetup   func(*mockapi.MockO11yAPI)
		wantDone    bool
	}{
		{
			name:        "create success",
			peeringNLB:  "",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:             ptr.To("test.dns.com"),
						EndpointServiceName: "epsrv-123",
						VMAccountID:         "vm-acc-1",
						VMProjectID:         "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.ACTIVE,
				}, nil)
			},
			wantDone: true,
		},
		{
			name:        "create with peering NLB",
			peeringNLB:  "internal-nlb.example.com",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:     ptr.To("test.dns.com"),
						VMAccountID: "vm-acc-1",
						VMProjectID: "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.ACTIVE,
				}, nil)
			},
			wantDone: true,
		},
		{
			name:        "create failed",
			peeringNLB:  "",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(nil, errors.New("create failed"))
			},
			wantDone: false,
		},
		{
			name:       "skip create when couldSkip is true and annotations exist",
			peeringNLB: "internal-nlb.example.com",
			annotations: map[string]string{
				consts.AnnotationKeyO11yID:             "cluster-1",
				consts.AnnotationKeyO11yServiceDNSName: "internal-nlb.example.com",
				consts.AnnotationKeyO11yVMAccountID:    "vm-acc",
				consts.AnnotationKeyO11yVMProjectID:    "vm-proj",
				consts.AnnotationKeyO11yLogBucketName:  "log-bucket",
				consts.AnnotationKeyO11yAgentRoleArn:   "role-arn",
			},
			pause:     false,
			couldSkip: true,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				// No API call expected
			},
			wantDone: true,
		},
		{
			name:        "DNS resolution fails when backend DNS is nil",
			peeringNLB:  "",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:     nil,
						VMAccountID: "vm-acc-1",
						VMProjectID: "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.ACTIVE,
				}, nil)
			},
			wantDone: false,
		},
		{
			name:        "no log storage fails",
			peeringNLB:  "",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:             ptr.To("test.dns.com"),
						EndpointServiceName: "epsrv-123",
						VMAccountID:         "vm-acc-1",
						VMProjectID:         "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.METRICS, BucketName: "metrics-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.ACTIVE,
				}, nil)
			},
			wantDone: false,
		},
		{
			name:        "missing endpoint service name with agent access account fails",
			peeringNLB:  "",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:             ptr.To("test.dns.com"),
						EndpointServiceName: "",
						VMAccountID:         "vm-acc-1",
						VMProjectID:         "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.ACTIVE,
				}, nil)
			},
			wantDone: false,
		},
		{
			name:        "resume cluster when status is pausing and pause is false",
			peeringNLB:  "",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:             ptr.To("test.dns.com"),
						EndpointServiceName: "epsrv-123",
						VMAccountID:         "vm-acc-1",
						VMProjectID:         "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.PAUSING,
				}, nil)
				m.On("Resume", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					Status:          api.ACTIVE,
				}, nil)
			},
			wantDone: true,
		},
		{
			name:        "pause cluster when status is active and pause is true",
			peeringNLB:  "",
			annotations: nil,
			pause:       true,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:             ptr.To("test.dns.com"),
						EndpointServiceName: "epsrv-123",
						VMAccountID:         "vm-acc-1",
						VMProjectID:         "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.ACTIVE,
				}, nil)
				m.On("Pause", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					Status:          api.PAUSING,
				}, nil)
			},
			wantDone: true,
		},
		{
			name:        "resume cluster fails",
			peeringNLB:  "",
			annotations: nil,
			pause:       false,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:             ptr.To("test.dns.com"),
						EndpointServiceName: "epsrv-123",
						VMAccountID:         "vm-acc-1",
						VMProjectID:         "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.PAUSING,
				}, nil)
				m.On("Resume", mock.Anything).Return(nil, errors.New("resume failed"))
			},
			wantDone: false,
		},
		{
			name:        "pause cluster fails",
			peeringNLB:  "",
			annotations: nil,
			pause:       true,
			couldSkip:   false,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Create", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					MonitorInfo: api.MonitorInfo{
						DnsName:             ptr.To("test.dns.com"),
						EndpointServiceName: "epsrv-123",
						VMAccountID:         "vm-acc-1",
						VMProjectID:         "vm-proj-1",
					},
					StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
					AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
					Status:      api.ACTIVE,
				}, nil)
				m.On("Pause", mock.Anything).Return(nil, errors.New("pause failed"))
			},
			wantDone: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockAPI := new(mockapi.MockO11yAPI)
			if tt.mockSetup != nil {
				tt.mockSetup(mockAPI)
			}

			task := makeBackendTaskWithAPI(tt.peeringNLB, tt.annotations, mockAPI)
			params := makeAwsParams("aws-123")
			res := task.Create(context.Background(), params, tt.pause, tt.couldSkip)

			assert.Equal(t, tt.wantDone, res.IsDoneState())
			mockAPI.AssertExpectations(t)
		})
	}
}

func TestBackendTask_Create_AzureSkipsBackendDNSButStillRequiresEndpointService(t *testing.T) {
	tests := []struct {
		name     string
		cluster  *api.ObservabilityCluster
		wantDone bool
	}{
		{
			name: "azure succeeds when endpoint service exists and backend DNS is empty",
			cluster: &api.ObservabilityCluster{
				ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-azure"},
				MonitorInfo: api.MonitorInfo{
					DnsName:             nil,
					EndpointServiceName: "azure-pls-id",
					VMAccountID:         "vm-acc-1",
					VMProjectID:         "vm-proj-1",
				},
				StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
				AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
				Status:      api.ACTIVE,
			},
			wantDone: true,
		},
		{
			name: "azure still fails when endpoint service is missing",
			cluster: &api.ObservabilityCluster{
				ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-azure"},
				MonitorInfo: api.MonitorInfo{
					DnsName:             nil,
					EndpointServiceName: "",
					VMAccountID:         "vm-acc-1",
					VMProjectID:         "vm-proj-1",
				},
				StorageInfo: []api.StorageInfo{{Category: api.LOGS, BucketName: "log-bucket"}},
				AgentInfo:   api.AgentInfo{RoleARN: "role-arn"},
				Status:      api.ACTIVE,
			},
			wantDone: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockAPI := new(mockapi.MockO11yAPI)
			mockAPI.On("Create", mock.MatchedBy(func(p o11y.Params) bool {
				return p.Vendor() == "azure" && p.SkipEndpointServiceName() && p.AgentAccessAccount() == "sub-123"
			})).Return(tt.cluster, nil)

			task := makeBackendTaskWithAPI("", nil, mockAPI)
			res := task.Create(context.Background(), makeAzureSkipDNSParams("sub-123"), false, false)

			assert.Equal(t, tt.wantDone, res.IsDoneState())
			mockAPI.AssertExpectations(t)
		})
	}
}

func TestBackendTask_PostCreate(t *testing.T) {
	tests := []struct {
		name        string
		annotations map[string]string
		pause       bool
		wantDone    bool
	}{
		{
			name: "active status without pause",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.ACTIVE),
			},
			pause:    false,
			wantDone: true,
		},
		{
			name: "paused status with pause",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.PAUSED),
			},
			pause:    true,
			wantDone: true,
		},
		{
			name: "pausing status with pause",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.PAUSING),
			},
			pause:    true,
			wantDone: true,
		},
		{
			name: "to_be_paused status with pause",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.TO_BE_PAUSED),
			},
			pause:    true,
			wantDone: true,
		},
		{
			name: "paused status without pause should error",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.PAUSED),
			},
			pause:    false,
			wantDone: false,
		},
		{
			name: "pausing status without pause should error",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.PAUSING),
			},
			pause:    false,
			wantDone: false,
		},
		{
			name: "active status with pause should error",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.ACTIVE),
			},
			pause:    true,
			wantDone: false,
		},
		{
			name: "creating status without pause is ok",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.CREATING),
			},
			pause:    false,
			wantDone: true,
		},
		{
			name: "updating status without pause is ok",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.UPDATING),
			},
			pause:    false,
			wantDone: true,
		},
		{
			name: "restarting status without pause is ok",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.RESTARTING),
			},
			pause:    false,
			wantDone: true,
		},
		{
			name: "deleting status without pause should error",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: string(api.DELETING),
			},
			pause:    false,
			wantDone: false,
		},
		{
			name: "empty status without pause should error",
			annotations: map[string]string{
				consts.AnnotationKeyO11yClusterStatus: "",
			},
			pause:    false,
			wantDone: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			task := makeBackendTask("", tt.annotations)
			res := task.PostCreate(context.Background(), tt.pause)
			assert.Equal(t, tt.wantDone, res.IsDoneState())
		})
	}
}

func TestBackendTask_Delete(t *testing.T) {
	tests := []struct {
		name        string
		annotations map[string]string
		mockSetup   func(*mockapi.MockO11yAPI)
		wantDone    bool
	}{
		{
			name: "delete with o11y ID annotation",
			annotations: map[string]string{
				consts.AnnotationKeyO11yID: "cluster-1",
			},
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Delete", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					Status:          api.DELETING,
				}, nil)
			},
			wantDone: true,
		},
		{
			name: "delete with o11y ID annotation fails",
			annotations: map[string]string{
				consts.AnnotationKeyO11yID: "cluster-1",
			},
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("Delete", mock.Anything).Return(nil, errors.New("delete failed"))
			},
			wantDone: false,
		},
		{
			name:        "delete without o11y ID annotation - list returns one cluster",
			annotations: nil,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("List", mock.Anything).Return([]api.ClusterBaseInfo{
					{ID: "cluster-1"},
				}, nil)
				m.On("Delete", mock.Anything).Return(&api.ObservabilityCluster{
					ClusterBaseInfo: api.ClusterBaseInfo{ID: "cluster-1"},
					Status:          api.DELETING,
				}, nil)
			},
			wantDone: true,
		},
		{
			name:        "delete without o11y ID annotation - list returns zero clusters - skip and continue",
			annotations: nil,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("List", mock.Anything).Return([]api.ClusterBaseInfo{}, nil)
			},
			wantDone: true,
		},
		{
			name:        "delete without o11y ID annotation - list returns multiple clusters",
			annotations: nil,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("List", mock.Anything).Return([]api.ClusterBaseInfo{
					{ID: "cluster-1"},
					{ID: "cluster-2"},
				}, nil)
			},
			wantDone: false,
		},
		{
			name:        "delete without o11y ID annotation - list fails",
			annotations: nil,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("List", mock.Anything).Return(nil, errors.New("list failed"))
			},
			wantDone: false,
		},
		{
			name:        "delete without o11y ID annotation - delete after list fails",
			annotations: nil,
			mockSetup: func(m *mockapi.MockO11yAPI) {
				m.On("List", mock.Anything).Return([]api.ClusterBaseInfo{
					{ID: "cluster-1"},
				}, nil)
				m.On("Delete", mock.Anything).Return(nil, errors.New("delete failed"))
			},
			wantDone: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mockAPI := new(mockapi.MockO11yAPI)
			if tt.mockSetup != nil {
				tt.mockSetup(mockAPI)
			}

			task := makeBackendTaskWithAPI("", tt.annotations, mockAPI)
			res := task.Delete(context.Background(), "aws", "us-east-1", "tenant-1", "project-1", "app-1", "tidbcloud")

			assert.Equal(t, tt.wantDone, res.IsDoneState())
			mockAPI.AssertExpectations(t)
		})
	}
}

func TestBackendTask_GetAgentAccessARN(t *testing.T) {
	tests := []struct {
		name              string
		withAgentAccess   bool
		agentAccessARNKey string
		annotations       map[string]string
		want              string
	}{
		{
			name:            "disabled access arn returns empty",
			withAgentAccess: false,
			annotations: map[string]string{
				"custom-key": "custom-value",
			},
			want: "",
		},
		{
			name:              "use custom key when provided",
			withAgentAccess:   true,
			agentAccessARNKey: "custom-key",
			annotations: map[string]string{
				"custom-key": "custom-value",
			},
			want: "custom-value",
		},
		{
			name:              "custom key missing does not fallback",
			withAgentAccess:   true,
			agentAccessARNKey: "custom-key",
			annotations: map[string]string{
				consts.AnnotationKeyAKSKubeletObjectID: "kubelet-obj-id",
			},
			want: "",
		},
		{
			name:            "fallback to aks kubelet object id",
			withAgentAccess: true,
			annotations: map[string]string{
				consts.AnnotationKeyAKSKubeletObjectID: "kubelet-obj-id",
			},
			want: "kubelet-obj-id",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			task := makeBackendTask("", tt.annotations)
			task.WithAgentAccessARN = tt.withAgentAccess
			task.AgentAccessARNKey = tt.agentAccessARNKey

			got := task.getAgentAccessARN()
			assert.Equal(t, tt.want, got)
		})
	}
}