package builder

import (
	"testing"

	"github.com/stretchr/testify/assert"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"

	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common/audit"
	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common/feature"
)

func TestBuildTiDBContainer(t *testing.T) {
	cases := []struct {
		name     string
		features feature.Gates
		res      corev1.ResourceRequirements
		want     *corev1.Container
	}{
		{
			name:     "no features",
			features: feature.NewGates(),
			res: corev1.ResourceRequirements{
				Requests: corev1.ResourceList{
					corev1.ResourceCPU:    resource.MustParse("50m"),
					corev1.ResourceMemory: resource.MustParse("4Gi"),
				},
				Limits: corev1.ResourceList{
					corev1.ResourceMemory: resource.MustParse("8Gi"),
				},
			},
			want: &corev1.Container{
				Name: "tidb",
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("50m"),
						corev1.ResourceMemory: resource.MustParse("4Gi"),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceMemory: resource.MustParse("8Gi"),
					},
				},
				Env: []corev1.EnvVar{
					{
						Name:  "TZ",
						Value: "UTC",
					},
					{
						Name: "NAMESPACE",
						ValueFrom: &corev1.EnvVarSource{
							FieldRef: &corev1.ObjectFieldSelector{
								FieldPath: "metadata.namespace",
							},
						},
					},
					{
						Name: "POD_NAME",
						ValueFrom: &corev1.EnvVarSource{
							FieldRef: &corev1.ObjectFieldSelector{
								FieldPath: "metadata.name",
							},
						},
					},
				},
			},
		},
		{
			name: "enable all related features",
			features: feature.NewGates(
				feature.IndexAccelerate,
				feature.MountPlanReplayerV1,
				feature.MountOOMPathV2,
				feature.SetGoMemLimitEnv,
			),
			res: corev1.ResourceRequirements{
				Requests: corev1.ResourceList{
					corev1.ResourceCPU:    resource.MustParse("50m"),
					corev1.ResourceMemory: resource.MustParse("4Gi"),
				},
				Limits: corev1.ResourceList{
					corev1.ResourceMemory: resource.MustParse("8Gi"),
				},
			},
			want: &corev1.Container{
				Name: "tidb",
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("50m"),
						corev1.ResourceMemory: resource.MustParse("4Gi"),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceMemory: resource.MustParse("8Gi"),
					},
				},
				Env: []corev1.EnvVar{
					{
						Name:  "TZ",
						Value: "UTC",
					},
					{
						Name: "NAMESPACE",
						ValueFrom: &corev1.EnvVarSource{
							FieldRef: &corev1.ObjectFieldSelector{
								FieldPath: "metadata.namespace",
							},
						},
					},
					{
						Name: "POD_NAME",
						ValueFrom: &corev1.EnvVarSource{
							FieldRef: &corev1.ObjectFieldSelector{
								FieldPath: "metadata.name",
							},
						},
					},
					{
						Name:  "GOMEMLIMIT",
						Value: "6871947673",
					},
				},
				VolumeMounts: []corev1.VolumeMount{
					{
						Name:      AddIndexAccelerateVolumeName,
						MountPath: AddIndexAccelerateMountPath,
					},
					{
						Name:        TiDBPlanReplayerVolumeName,
						MountPath:   TiDBPlanReplayerContainerPathV1,
						SubPathExpr: TiDBPlanReplayerMountSubPathExpr,
					},
					{
						Name:        TiDBOOMRecordVolumeName,
						MountPath:   TiDBOOMRecordContainerPathV2,
						SubPathExpr: TiDBOOMRecordMountSubPathExpr,
					},
				},
			},
		},
	}
	for i := range cases {
		c := &cases[i]
		t.Run(c.name, func(tt *testing.T) {
			tt.Parallel()
			assert.Equalf(tt, c.want, buildTiDBContainer(c.features, c.res), "buildTiDBContainer(%v)", c.res)
		})
	}
}

func TestBuildTiDBVolumes(t *testing.T) {
	cases := []struct {
		name         string
		features     feature.Gates
		tenantID     string
		projectID    string
		clusterID    string
		instanceType string
		want         []corev1.Volume
	}{
		{
			name:     "no features",
			features: feature.NewGates(),
		},
		{
			name: "enable all related features",
			features: feature.NewGates(
				feature.IndexAccelerate,
				feature.MountPlanReplayerV2,
				feature.MountOOMPathV2,
				feature.UseLegacyAuditPlugin,
				feature.SetupGCPCredentials,
			),
			tenantID:  "1",
			projectID: "2",
			clusterID: "3",
			want: []corev1.Volume{
				{
					Name: AddIndexAccelerateVolumeName,
					VolumeSource: corev1.VolumeSource{
						Ephemeral: &corev1.EphemeralVolumeSource{
							VolumeClaimTemplate: &corev1.PersistentVolumeClaimTemplate{
								ObjectMeta: metav1.ObjectMeta{},
								Spec: corev1.PersistentVolumeClaimSpec{
									AccessModes: []corev1.PersistentVolumeAccessMode{
										corev1.ReadWriteOnce,
									},
									StorageClassName: ptr.To(GetAddIndexAccelerateStorageClassName("")),
									Resources: corev1.VolumeResourceRequirements{
										Requests: corev1.ResourceList{
											corev1.ResourceStorage: resource.MustParse(AddIndexAccelerateVolumeSize),
										},
									},
								},
							},
						},
					},
				},
				{
					Name: TiDBPlanReplayerVolumeName,
					VolumeSource: corev1.VolumeSource{
						HostPath: &corev1.HostPathVolumeSource{
							Path: "/var/tidb_plan_replayer/1/2/3",
							Type: ptr.To(corev1.HostPathDirectoryOrCreate),
						},
					},
				},
				{
					Name: TiDBOOMRecordVolumeName,
					VolumeSource: corev1.VolumeSource{
						HostPath: &corev1.HostPathVolumeSource{
							Path: "/var/tidb_oom_record/1/2/3",
							Type: ptr.To(corev1.HostPathDirectoryOrCreate),
						},
					},
				},
				{
					Name: "audit-config",
					VolumeSource: corev1.VolumeSource{
						ConfigMap: &corev1.ConfigMapVolumeSource{
							LocalObjectReference: corev1.LocalObjectReference{
								Name: "3-audit",
							},
						},
					},
				},
				{
					Name: "audit-credentials",
					VolumeSource: corev1.VolumeSource{
						Secret: &corev1.SecretVolumeSource{
							SecretName: "db-audit",
						},
					},
				},
			},
		},
		{
			name: "enable all related features",
			features: feature.NewGates(
				feature.IndexAccelerate,
			),
			tenantID:     "1",
			projectID:    "2",
			clusterID:    "3",
			instanceType: "n4-highcpu-16",
			want: []corev1.Volume{
				{
					Name: AddIndexAccelerateVolumeName,
					VolumeSource: corev1.VolumeSource{
						Ephemeral: &corev1.EphemeralVolumeSource{
							VolumeClaimTemplate: &corev1.PersistentVolumeClaimTemplate{
								ObjectMeta: metav1.ObjectMeta{},
								Spec: corev1.PersistentVolumeClaimSpec{
									AccessModes: []corev1.PersistentVolumeAccessMode{
										corev1.ReadWriteOnce,
									},
									StorageClassName: ptr.To("index-acceleration-hyperdisk-balanced"),
									Resources: corev1.VolumeResourceRequirements{
										Requests: corev1.ResourceList{
											corev1.ResourceStorage: resource.MustParse(AddIndexAccelerateVolumeSize),
										},
									},
								},
							},
						},
					},
				},
			},
		},
	}
	for i := range cases {
		c := &cases[i]
		t.Run(c.name, func(tt *testing.T) {
			tt.Parallel()
			assert.Equalf(tt, c.want, buildTiDBVolumes(c.features, c.instanceType, c.tenantID, c.projectID, c.clusterID), c.name)
		})
	}
}

func Test_buildContainersForTiDB(t *testing.T) {
	slowlogVolumeName := "tidb-" + audit.SlowlogVolumeName
	mustHaveContainers := []corev1.Container{
		{Name: "tidb"},
		{Name: "statementlog"},
	}

	tests := []struct {
		name     string
		features feature.Gates
		cfg      *TiDBGroupConfig
		want     []corev1.Container
	}{
		{
			name:     "no features",
			features: feature.NewGates(),
			cfg:      &TiDBGroupConfig{},
			want:     mustHaveContainers,
		},
		{
			name: "enable all related features",
			features: feature.NewGates(
				feature.UseLegacyAuditPlugin,
				feature.SetupGCPCredentials,
			),
			cfg: &TiDBGroupConfig{
				AuditLogSidecarImage:   "auditlog",
				AuditLogSidecarCommand: []string{"test"},
			},
			want: append(mustHaveContainers, corev1.Container{
				Name:    "auditlog",
				Image:   "auditlog",
				Command: []string{"test"},
				Env: []corev1.EnvVar{
					{
						Name: "POD_NAME",
						ValueFrom: &corev1.EnvVarSource{
							FieldRef: &corev1.ObjectFieldSelector{
								FieldPath: "metadata.name",
							},
						},
					},
					{
						Name:  "GOOGLE_APPLICATION_CREDENTIALS",
						Value: audit.GCPCredentialPath,
					},
				},
				Resources: corev1.ResourceRequirements{
					Limits: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("100m"),
						corev1.ResourceMemory: resource.MustParse("256Mi"),
					},
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("0m"),
						corev1.ResourceMemory: resource.MustParse("0Mi"),
					},
				},
				VolumeMounts: []corev1.VolumeMount{
					{
						Name:      "audit-config",
						MountPath: "/etc/audit",
					},
					{
						Name:      "ti-vol-tidb-slowlog",
						MountPath: audit.MountPath,
					},
					{
						Name:      "audit-credentials",
						MountPath: "/.gcp",
					},
				},
			}),
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			got := buildContainersForTiDB(tt.features, tt.cfg, slowlogVolumeName)
			if len(got) != len(tt.want) {
				t.Errorf("buildContainersForTiDB() = %v, want %v", got, tt.want)
			}
			// Only check the auditlog container if it exists
			if len(got) > len(mustHaveContainers) {
				assert.Equalf(t, tt.want[len(tt.want)-1], got[len(got)-1], "buildInitContainersForTiDB, expect %v, but got %v", tt.want, got)
			}
		})
	}
}