package common

import (
	"context"
	"fmt"
	"testing"

	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"go.uber.org/mock/gomock"
	appsv1 "k8s.io/api/apps/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/utils/ptr"

	v1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/config"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
)

var instanceMap = map[string]v1.InstanceType{
	"2c4g": {
		Name:   "2c4g",
		CPU:    resource.MustParse("2"),
		Memory: resource.MustParse("4Gi"),
	},
	"2c8g": {
		Name:   "2c8g",
		CPU:    resource.MustParse("2"),
		Memory: resource.MustParse("8Gi"),
	},
	"4c8g": {
		Name:   "4c8g",
		CPU:    resource.MustParse("4"),
		Memory: resource.MustParse("8Gi"),
	},
	"4c8g-2": {
		Name:   "4c8g-2",
		CPU:    resource.MustParse("4"),
		Memory: resource.MustParse("8Gi"),
	},
}

func Test_isComponentScalingUp(t *testing.T) {
	ctl := gomock.NewController(t)
	provider := gitops.NewMockProvider(ctl)
	for name, instanceType := range instanceMap {
		provider.EXPECT().FindInstanceType(name).Return(instanceType, nil).AnyTimes()
	}
	provider.EXPECT().FindInstanceType("not-exist").Return(v1.InstanceType{}, fmt.Errorf("failed to find the instance type")).AnyTimes()

	tests := []struct {
		name        string
		spec        v1.ClusterComponent
		status      v1.ClusterComponentStatus
		isScalingUp bool
		wantErr     bool
	}{
		{
			name: "can not find spec instance type",
			spec: &v1.TiFlash{
				InstanceType: "not-exist",
			},
			status: &v1.TiFlashStatus{
				Members: []v1.StoreMember{},
			},
			isScalingUp: false,
			wantErr:     true,
		},
		{
			name: "can not find status instance type",
			spec: &v1.TiFlash{
				InstanceType: "2c4g",
			},
			status: &v1.TiFlashStatus{
				Members: []v1.StoreMember{
					{InstanceType: "not-exist"},
				},
			},
			isScalingUp: false,
			wantErr:     true,
		},
		{
			name: "adding tiflash is not scaling up",
			spec: &v1.TiFlash{
				InstanceType: "2c4g",
			},
			status: &v1.TiFlashStatus{
				Members: []v1.StoreMember{},
			},
			isScalingUp: false,
		},
		{
			name: "add tiproxy is not scaling up",
			spec: &v1.TiProxy{
				InstanceType: "2c4g",
			},
			status: &v1.TiProxyStatus{
				Members: []v1.TiProxyMember{},
			},
			isScalingUp: false,
		},
		{
			name: "did not change instance type",
			spec: &v1.PD{
				InstanceType: "2c4g",
			},
			status: &v1.PDStatus{
				Members: []v1.PDMember{
					{InstanceType: "2c4g"},
					{InstanceType: "2c4g"},
				},
			},
			isScalingUp: false,
		},
		{
			name: "change instance type but not scaling up",
			spec: &v1.PD{
				InstanceType: "4c8g",
			},
			status: &v1.PDStatus{
				Members: []v1.PDMember{
					{InstanceType: "4c8g-2"},
					{InstanceType: "4c8g"},
				},
			},
			isScalingUp: false,
		},
		{
			name: "scaling up memory",
			spec: &v1.PD{
				InstanceType: "2c8g",
			},
			status: &v1.PDStatus{
				Members: []v1.PDMember{
					{InstanceType: "2c4g"},
					{InstanceType: "2c8g"},
					{InstanceType: "2c8g"},
				},
			},
			isScalingUp: true,
		},
		{
			name: "scaling up cpu and memory",
			spec: &v1.PD{
				InstanceType: "4c8g",
			},
			status: &v1.PDStatus{
				Members: []v1.PDMember{
					{InstanceType: "2c4g"},
					{InstanceType: "2c4g"},
					{InstanceType: "4c8g"},
				},
			},
			isScalingUp: true,
		},
		{
			name: "scaling up cpu",
			spec: &v1.PD{
				InstanceType: "4c8g",
			},
			status: &v1.PDStatus{
				Members: []v1.PDMember{
					{InstanceType: "4c8g"},
					{InstanceType: "2c8g"},
					{InstanceType: "2c8g"},
				},
			},
			isScalingUp: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := isComponentScalingUp(provider, tt.spec, tt.status)
			if (err != nil) != tt.wantErr {
				t.Errorf("isComponentScalingUp() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if got != tt.isScalingUp {
				t.Errorf("isComponentScalingUp() = %v, isScalingUp %v", got, tt.isScalingUp)
			}
		})
	}
}

func makeTidbClusterStatusWithReplicas(pd, tidb, tikv int32, tiflash, tiproxy *appsv1.StatefulSetStatus) *opv1alpha1.TidbClusterStatus {
	status := &opv1alpha1.TidbClusterStatus{
		PD:      opv1alpha1.PDStatus{},
		TiDB:    opv1alpha1.TiDBStatus{},
		TiKV:    opv1alpha1.TiKVStatus{},
		TiFlash: opv1alpha1.TiFlashStatus{},
		TiProxy: opv1alpha1.TiProxyStatus{},
	}

	if pd > 0 {
		status.PD.SetStatefulSet(&appsv1.StatefulSetStatus{
			Replicas:        pd,
			UpdatedReplicas: pd,
		})
	}
	if tidb > 0 {
		status.TiDB.SetStatefulSet(&appsv1.StatefulSetStatus{
			Replicas:        tidb,
			UpdatedReplicas: tidb,
		})
	}
	if tikv > 0 {
		status.TiKV.SetStatefulSet(&appsv1.StatefulSetStatus{
			Replicas:        tikv,
			UpdatedReplicas: tikv,
		})
	}

	status.TiFlash.SetStatefulSet(tiflash)
	status.TiProxy.SetStatefulSet(tiproxy)

	return status
}

func makeTidbClusterSpecWithDesiredReplicas(pd, tidb, tikv int32, tiflash *opv1alpha1.TiFlashSpec, tiproxy *opv1alpha1.TiProxySpec) *opv1alpha1.TidbClusterSpec {
	return &opv1alpha1.TidbClusterSpec{
		PD:      &opv1alpha1.PDSpec{Replicas: pd},
		TiDB:    &opv1alpha1.TiDBSpec{Replicas: tidb},
		TiKV:    &opv1alpha1.TiKVSpec{Replicas: tikv},
		TiFlash: tiflash,
		TiProxy: tiproxy,
	}
}

func Test_isModifyingComponentReplicas(t *testing.T) {
	tests := []struct {
		name   string
		spec   *opv1alpha1.TidbClusterSpec
		status *opv1alpha1.TidbClusterStatus
		want   bool
	}{
		{
			name:   "no change",
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, nil, nil),
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, nil, nil),
			want:   false,
		},
		{
			name:   "scale out tidb",
			status: makeTidbClusterStatusWithReplicas(3, 1, 3, nil, nil),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, nil, nil),
			want:   true,
		},
		{
			name:   "scale in tidb",
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, nil, nil),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 1, 3, nil, nil),
			want:   true,
		},
		{
			name:   "scale tiflash from 0 (no asts)",
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, nil, nil),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, &opv1alpha1.TiFlashSpec{Replicas: 1}, nil),
			want:   false,
		},
		{
			name:   "scale tiflash from 0 (has asts)",
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, &appsv1.StatefulSetStatus{Replicas: 0, UpdatedReplicas: 0}, nil),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, &opv1alpha1.TiFlashSpec{Replicas: 1}, nil),
			want:   false,
		},
		{
			name:   "scale tiflash to 0",
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, &appsv1.StatefulSetStatus{Replicas: 2, UpdatedReplicas: 2}, nil),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, &opv1alpha1.TiFlashSpec{Replicas: 0}, nil),
			want:   false,
		},
		{
			name:   "scale tiproxy from 0 (no asts)",
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, nil, nil),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, nil, &opv1alpha1.TiProxySpec{Replicas: 3}),
			want:   false,
		},
		{
			name:   "scale tiproxy from 0 (has asts)",
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, nil, &appsv1.StatefulSetStatus{Replicas: 0, UpdatedReplicas: 0}),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, nil, &opv1alpha1.TiProxySpec{Replicas: 3}),
			want:   false,
		},
		{
			name:   "scale tiproxy to 0",
			status: makeTidbClusterStatusWithReplicas(3, 2, 3, nil, &appsv1.StatefulSetStatus{Replicas: 3, UpdatedReplicas: 3}),
			spec:   makeTidbClusterSpecWithDesiredReplicas(3, 2, 3, nil, &opv1alpha1.TiProxySpec{Replicas: 0}),
			want:   false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tc := &opv1alpha1.TidbCluster{
				Spec:   *tt.spec,
				Status: *tt.status,
			}
			if got := isModifyingComponentReplicas(tc); got != tt.want {
				t.Errorf("isModifyingComponentReplicas() = %v, isScalingUp %v", got, tt.want)
			}
		})
	}
}

func TestSetPreStopHook(t *testing.T) {
	tests := []struct {
		name                            string
		preStopHookSleepSeconds         *int64
		initialTerminationGracePeriod   *int64
		expectedPreStopHookSleepSeconds int64
		expectedTerminationGracePeriod  int64
	}{
		{
			name:                            "Default preStop sleep seconds",
			preStopHookSleepSeconds:         nil,
			initialTerminationGracePeriod:   nil,
			expectedPreStopHookSleepSeconds: config.DefaultPreStopSleepSeconds,
			expectedTerminationGracePeriod:  consts.DefaultPodTerminationGracePeriodSeconds + config.DefaultPreStopSleepSeconds,
		},
		{
			name:                            "Default termination grace period",
			preStopHookSleepSeconds:         ptr.To[int64](15),
			initialTerminationGracePeriod:   nil,
			expectedPreStopHookSleepSeconds: 15,
			expectedTerminationGracePeriod:  consts.DefaultPodTerminationGracePeriodSeconds + 15,
		},
		{
			name:                            "Custom termination grace period",
			initialTerminationGracePeriod:   ptr.To[int64](60),
			expectedPreStopHookSleepSeconds: config.DefaultPreStopSleepSeconds,
			expectedTerminationGracePeriod:  60,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			cluster := &v1.Cluster{}
			tc := &opv1alpha1.TidbCluster{
				Spec: opv1alpha1.TidbClusterSpec{
					TiDB: &opv1alpha1.TiDBSpec{
						ComponentSpec: opv1alpha1.ComponentSpec{
							TerminationGracePeriodSeconds: tt.initialTerminationGracePeriod,
						},
					},
				},
			}
			cfg := config.ClusterControllerConfig{
				TiDBPreStopHookSleepSeconds: tt.preStopHookSleepSeconds,
			}

			setPreStopHook(ctx, cluster, "group1", tc, cfg)

			if tc.Spec.TiDB.Lifecycle == nil || tc.Spec.TiDB.Lifecycle.PreStop == nil {
				t.Fatalf("PreStop hook was not set")
			}

			if tc.Spec.TiDB.Lifecycle.PreStop.Exec == nil || len(tc.Spec.TiDB.Lifecycle.PreStop.Exec.Command) != 2 {
				t.Fatalf("PreStop hook command is incorrect")
			}

			if tc.Spec.TiDB.Lifecycle.PreStop.Exec.Command[0] != "sleep" {
				t.Errorf("Expected 'sleep' command, got %s", tc.Spec.TiDB.Lifecycle.PreStop.Exec.Command[0])
			}

			if tc.Spec.TiDB.Lifecycle.PreStop.Exec.Command[1] != fmt.Sprintf("%d", tt.expectedPreStopHookSleepSeconds) {
				t.Errorf("Expected sleep duration %d, got %s", tt.expectedPreStopHookSleepSeconds, tc.Spec.TiDB.Lifecycle.PreStop.Exec.Command[1])
			}

			if *tc.Spec.TiDB.TerminationGracePeriodSeconds != tt.expectedTerminationGracePeriod {
				t.Errorf("Expected termination grace period %d, got %d", tt.expectedTerminationGracePeriod, *tc.Spec.TiDB.TerminationGracePeriodSeconds)
			}
		})
	}
}