package commontasks

import (
	"context"
	"errors"
	"testing"

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

	clusterserverv1 "github.com/tidbcloud/observability/agent-control-server/api/clusterserver/v1"
	"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/testkits"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

const testSpecs = `
- name: default
  maxNodes: 200
  instanceType: t3.medium
  vmagent:
    limits:
      cpu: "2"
      memory: "3Gi"
    requests:
      cpu: "1500m"
      memory: "3Gi"
- name: medium
  maxNodes: 999999
  instanceType: c5.xlarge
  vmagent:
    limits:
      cpu: "4"
      memory: "7Gi"
    requests:
      cpu: "3"
      memory: "7Gi"
`

func TestDetermineSpec(t *testing.T) {
	var specs []config.CustomerAgentSpec
	err := yaml.Unmarshal([]byte(testSpecs), &specs)
	assert.NoError(t, err)

	tests := []struct {
		name        string
		desiredSpec string
		minSpec     string
		maxSpec     string
		nodeCount   uint32
		want        string
	}{
		{
			name:      "node count 0",
			nodeCount: 0,
			want:      "default",
		},
		{
			name:      "node count 50",
			nodeCount: 50,
			want:      "default",
		},
		{
			name:      "node count 51",
			nodeCount: 51,
			want:      "default",
		},
		{
			name:      "node count 101",
			nodeCount: 101,
			want:      "default",
		},
		{
			name:      "node count 201",
			nodeCount: 201,
			want:      "medium",
		},
		{
			name:      "custom min spec",
			minSpec:   "custom-min",
			maxSpec:   "custom-max",
			nodeCount: 201,
			want:      "custom-min",
		},
		{
			name:      "respect max spec",
			minSpec:   "default",
			maxSpec:   "default",
			nodeCount: 999,
			want:      "default",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := determineSpec(specs, tt.desiredSpec, tt.minSpec, tt.maxSpec, tt.nodeCount)
			assert.Equal(t, tt.want, got)
		})
	}
}

func TestReconcileAutoResourcing(t *testing.T) {
	var specs config.CustomerAgentSpecs
	err := yaml.Unmarshal([]byte(testSpecs), &specs)
	assert.NoError(t, err)

	matchingLabels := map[string]string{"app": "test"}
	tests := []struct {
		name           string
		agent          client.Object
		defaultMaxSpec string
		clusters       []clusterserverv1.Cluster
		listErr        error
		wantErr        bool
		validate       func(t *testing.T, agent client.Object)
	}{
		{
			name: "list clusters error",
			agent: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-agent",
				},
			},
			defaultMaxSpec: "medium",
			listErr:        errors.New("list clusters error"),
			wantErr:        true,
			validate: func(t *testing.T, agent client.Object) {
				// Defaults should be set even if list fails
				assert.Equal(t, "default", util.GetAnnotation(agent, consts.AnnotationKeyO11yMinSpec))
				assert.Equal(t, "medium", util.GetAnnotation(agent, consts.AnnotationKeyO11yMaxSpec))
			},
		},
		{
			name: "no change needed",
			agent: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-agent",
					Annotations: map[string]string{
						consts.AnnotationKeyO11yDesiredSpec: "default",
						consts.AnnotationKeyO11yMinSpec:     "default",
						consts.AnnotationKeyO11yMaxSpec:     "medium",
					},
				},
			},
			defaultMaxSpec: "medium",
			clusters: []clusterserverv1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:   "test-cluster-1",
						Labels: matchingLabels,
					},
					Spec: clusterserverv1.ClusterSpec{
						Cluster: clusterserverv1.ClusterTopology{
							TiDB:    clusterserverv1.TiDB{Replicas: 2},
							PD:      clusterserverv1.PD{Replicas: 3},
							TiKV:    clusterserverv1.TiKV{Replicas: 3},
							TiFlash: clusterserverv1.TiFlash{Replicas: 2},
						},
					},
				},
			}, // Total 10 nodes
			wantErr: false,
			validate: func(t *testing.T, agent client.Object) {
				assert.Equal(t, "default", util.GetAnnotation(agent, consts.AnnotationKeyO11yDesiredSpec))
				assert.Equal(t, "", util.GetAnnotation(agent, consts.AnnotationKeyO11ySpecScaleStatus))
			},
		},
		{
			name: "scale up needed",
			agent: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-agent",
					Annotations: map[string]string{
						consts.AnnotationKeyO11yDesiredSpec: "default",
						consts.AnnotationKeyO11yMinSpec:     "default",
						consts.AnnotationKeyO11yMaxSpec:     "medium",
					},
				},
			},
			defaultMaxSpec: "medium",
			clusters: []clusterserverv1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:   "test-cluster-1",
						Labels: matchingLabels,
					},
					Spec: clusterserverv1.ClusterSpec{
						Cluster: clusterserverv1.ClusterTopology{
							TiDB:    clusterserverv1.TiDB{Replicas: 50},
							PD:      clusterserverv1.PD{Replicas: 50},
							TiKV:    clusterserverv1.TiKV{Replicas: 50},
							TiFlash: clusterserverv1.TiFlash{Replicas: 51},
						},
					},
				},
			}, // Total 201 nodes
			wantErr: false,
			validate: func(t *testing.T, agent client.Object) {
				assert.Equal(t, "medium", util.GetAnnotation(agent, consts.AnnotationKeyO11yDesiredSpec))
				assert.Equal(t, "default", util.GetAnnotation(agent, consts.AnnotationKeyO11yPreviousSpec))
				assert.Equal(t, consts.SpecScaleStatusPendingScale, util.GetAnnotation(agent, consts.AnnotationKeyO11ySpecScaleStatus))
			},
		},
		{
			name: "scale down needed",
			agent: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-agent",
					Annotations: map[string]string{
						consts.AnnotationKeyO11yDesiredSpec: "medium",
						consts.AnnotationKeyO11yMinSpec:     "default",
						consts.AnnotationKeyO11yMaxSpec:     "medium",
					},
				},
			},
			defaultMaxSpec: "medium",
			clusters: []clusterserverv1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:   "test-cluster-1",
						Labels: matchingLabels,
					},
					Spec: clusterserverv1.ClusterSpec{
						Cluster: clusterserverv1.ClusterTopology{
							TiDB:    clusterserverv1.TiDB{Replicas: 2},
							PD:      clusterserverv1.PD{Replicas: 3},
							TiKV:    clusterserverv1.TiKV{Replicas: 3},
							TiFlash: clusterserverv1.TiFlash{Replicas: 2},
						},
					},
				},
			}, // Total 10 nodes
			wantErr: false,
			validate: func(t *testing.T, agent client.Object) {
				assert.Equal(t, "default", util.GetAnnotation(agent, consts.AnnotationKeyO11yDesiredSpec))
				assert.Equal(t, "medium", util.GetAnnotation(agent, consts.AnnotationKeyO11yPreviousSpec))
				assert.Equal(t, consts.SpecScaleStatusPendingScale, util.GetAnnotation(agent, consts.AnnotationKeyO11ySpecScaleStatus))
			},
		},
		{
			name: "initialize defaults",
			agent: &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-agent",
				},
			},
			defaultMaxSpec: "medium",
			clusters: []clusterserverv1.Cluster{
				{
					ObjectMeta: metav1.ObjectMeta{
						Name:   "test-cluster-1",
						Labels: matchingLabels,
					},
					Spec: clusterserverv1.ClusterSpec{
						Cluster: clusterserverv1.ClusterTopology{
							TiDB:    clusterserverv1.TiDB{Replicas: 2},
							PD:      clusterserverv1.PD{Replicas: 3},
							TiKV:    clusterserverv1.TiKV{Replicas: 3},
							TiFlash: clusterserverv1.TiFlash{Replicas: 2},
						},
					},
				},
			}, // Total 10 nodes
			wantErr: false,
			validate: func(t *testing.T, agent client.Object) {
				assert.Equal(t, "default", util.GetAnnotation(agent, consts.AnnotationKeyO11yMinSpec))
				assert.Equal(t, "medium", util.GetAnnotation(agent, consts.AnnotationKeyO11yMaxSpec))
				// Since desired was empty, and determined is "default", it sets to default.
				assert.Equal(t, "default", util.GetAnnotation(agent, consts.AnnotationKeyO11yDesiredSpec))
				assert.Equal(t, "", util.GetAnnotation(agent, consts.AnnotationKeyO11yPreviousSpec))
				assert.Equal(t, "", util.GetAnnotation(agent, consts.AnnotationKeyO11ySpecScaleStatus))
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			objs := []client.Object{tt.agent}
			for i := range tt.clusters {
				objs = append(objs, &tt.clusters[i])
			}
			fakeCli := testkits.BuildClient(objs...)
			var c client.Client = fakeCli
			if tt.listErr != nil {
				c = testkits.NewClientGetter(fakeCli).WithListFunc(func(ctx context.Context, obj client.ObjectList, opts ...client.ListOption) error {
					return tt.listErr
				})
			}

			err := ReconcileAutoResourcing(context.Background(), specs, tt.agent, "test-cluster", "test-agent", tt.defaultMaxSpec, matchingLabels, c)
			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
			if tt.validate != nil {
				tt.validate(t, tt.agent)
			}
		})
	}
}