package gitops

import (
	"testing"

	corev1 "k8s.io/api/core/v1"

	"k8s.io/utils/ptr"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"

	. "github.com/onsi/gomega"
	"k8s.io/apimachinery/pkg/api/resource"
)

func TestComputeMemory(t *testing.T) {
	cases := map[string]resource.Quantity{
		"256Mi": resource.MustParse("0"),
		"512Mi": resource.MustParse("0"),
		"1Gi":   resource.MustParse("0"),
		"2Gi":   resource.MustParse("0"),
		"4Gi":   resource.MustParse("1536Mi"),
		"8Gi":   resource.MustParse("5046376857"),
		"16Gi":  resource.MustParse("12777527705"),
		"32Gi":  resource.MustParse("28925933649"),
		"64Gi":  resource.MustParse("61224087715"),
		"128Gi": resource.MustParse("125820395847"), // 117 Gi
		"256Gi": resource.MustParse("260511492997"), // 242 Gi
	}
	for total, expect := range cases {
		t.Run(total, func(t *testing.T) {
			g := NewGomegaWithT(t)
			computed := computeMemory(resource.MustParse(total))
			g.Expect(computed.Value()).To(Equal(expect.Value()))
		})
	}
}

func TestComputeMemoryLimit(t *testing.T) {
	cases := map[string]resource.Quantity{
		"16Gi":  resource.MustParse("14388140441"),
		"32Gi":  resource.MustParse("30536546385"),
		"64Gi":  resource.MustParse("62834700451"),
		"128Gi": resource.MustParse("127431008583"),
	}
	for total, expect := range cases {
		t.Run(total, func(t *testing.T) {
			g := NewGomegaWithT(t)
			computed := computeMemoryLimit(resource.MustParse(total))
			g.Expect(computed.Value()).To(Equal(expect.Value()))
		})
	}
}

func TestMachineImageVersion(t *testing.T) {
	cases := []struct {
		name              string
		c                 cloudv1.ClusterComponent
		expectErrorMsg    string
		expectMachineSpec cloudv1.MachineSpec
	}{
		{
			name: consts.AMD64ArchName,
			c: &cloudv1.TiKV{
				InstanceType: "c5.2xlarge",
			},
			expectErrorMsg: "",
			expectMachineSpec: cloudv1.MachineSpec{
				MachineType:         "c5.2xlarge",
				MachineImage:        storageNodeImage,
				MachineImageVersion: X86ImageVersion,
			},
		},
		{
			name: consts.ARM64ArchName,
			c: &cloudv1.TiKV{
				InstanceType: "c6g.2xlarge",
			},
			expectErrorMsg: "",
			expectMachineSpec: cloudv1.MachineSpec{
				MachineType:         "c6g.2xlarge",
				MachineImage:        storageNodeImage,
				MachineImageVersion: Aarch64ImageVersion,
			},
		},
		{
			name: "unknown_arch",
			c: &cloudv1.TiKV{
				InstanceType: "unknown_arch",
			},
			expectErrorMsg: "",
			expectMachineSpec: cloudv1.MachineSpec{
				MachineType:         "unknown_arch",
				MachineImage:        storageNodeImage,
				MachineImageVersion: X86ImageVersion,
			},
		},
		{
			name: "unknown_instance_type",
			c: &cloudv1.TiKV{
				InstanceType: "unknown_instance_type",
			},
			expectErrorMsg:    "unknown instance type",
			expectMachineSpec: cloudv1.MachineSpec{},
		},
		{
			name:           "empty_instance_type",
			c:              &cloudv1.TiCDC{},
			expectErrorMsg: "",
			expectMachineSpec: cloudv1.MachineSpec{
				MachineType:         "",
				MachineImage:        NodeImage,
				MachineImageVersion: X86ImageVersion,
			},
		},
	}
	g := NewGomegaWithT(t)
	configProvider := NewHardcode()
	var (
		awsProvider = &cloudv1.Provider{
			Spec: cloudv1.ProviderSpec{
				Type: consts.ProviderTypeAWS,
				InstanceTypes: []cloudv1.InstanceType{
					{
						Name: "c5.2xlarge",
						Arch: consts.AMD64ArchName,
					},
					{
						Name: "c6g.2xlarge",
						Arch: consts.ARM64ArchName,
					},
					{
						Name: "unknown_arch",
						Arch: "",
					},
				},
			},
		}
		gcpProvider = &cloudv1.Provider{
			Spec: cloudv1.ProviderSpec{
				Type: consts.ProviderTypeGCP,
			},
		}
	)
	for _, tc := range cases {
		t.Logf("test case: %s", tc.name)
		machineSpec, err := configProvider.GetMachineSpec(awsProvider, tc.c)
		if tc.expectErrorMsg == "" {
			g.Expect(err).Should(Succeed())
			g.Expect(machineSpec.MachineType).To(Equal(tc.expectMachineSpec.MachineType))
			g.Expect(machineSpec.MachineImage).To(Equal(tc.expectMachineSpec.MachineImage))
			g.Expect(machineSpec.MachineImageVersion).To(Equal(tc.expectMachineSpec.MachineImageVersion))
			// Test the non-AWS provider.
			machineSpec, err = configProvider.GetMachineSpec(gcpProvider, tc.c)
			g.Expect(err).Should(Succeed())
			g.Expect(machineSpec.MachineType).To(Equal(tc.expectMachineSpec.MachineType))
			g.Expect(machineSpec.MachineImage).To(Equal(tc.expectMachineSpec.MachineImage))
			// Should always be x86 image version.
			g.Expect(machineSpec.MachineImageVersion).To(Equal(X86ImageVersion))
		} else {
			g.Expect(err).Should(HaveOccurred())
			g.Expect(err.Error()).Should(ContainSubstring(tc.expectErrorMsg))
		}
	}
}

func TestGetResourceSpec(t *testing.T) {
	cases := map[string]struct {
		provider    Provider
		component   cloudv1.ClusterComponent
		expectError bool
		expectSpec  *cloudv1.ResourceSpec
	}{
		"PD component with default storage class on AWS": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					Type: consts.ProviderTypeAWS,
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.2xlarge", CPU: resource.MustParse("8"), Memory: resource.MustParse("16Gi")},
					},
				},
			},
			component: &cloudv1.PD{
				InstanceType: "c5.2xlarge",
			},
			expectError: false,
			expectSpec: &cloudv1.ResourceSpec{
				StorageClass: consts.PDGp3StorageClassName,
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("7"),
						corev1.ResourceMemory: computeMemory(resource.MustParse("16Gi")),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceMemory: computeMemoryLimit(resource.MustParse("16Gi")),
					},
				},
			},
		},
		"TiKV component with local storage on EKS": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					Type: consts.ProviderTypeEKS,
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.4xlarge", CPU: resource.MustParse("16"), Memory: resource.MustParse("32Gi")},
					},
				},
			},
			component: &cloudv1.TiKV{
				InstanceType:    "c5.4xlarge",
				UseLocalStorage: ptr.To(true),
			},
			expectError: false,
			expectSpec: &cloudv1.ResourceSpec{
				StorageClass: localStorageClassName,
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("15"),
						corev1.ResourceMemory: computeMemory(resource.MustParse("32Gi")),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceMemory: computeMemoryLimit(resource.MustParse("32Gi")),
					},
				},
			},
		},
		"TiFlash component with storage size on GKE": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					Type: consts.ProviderTypeGKE,
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "n1-standard-8", CPU: resource.MustParse("8"), Memory: resource.MustParse("30Gi")},
					},
				},
			},
			component: &cloudv1.TiFlash{
				InstanceType: "n1-standard-8",
				StorageSize:  resource.MustParse("500Gi"),
			},
			expectError: false,
			expectSpec: &cloudv1.ResourceSpec{
				StorageClass: consts.TiFlashGp3StorageClassName,
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:     resource.MustParse("7"),
						corev1.ResourceMemory:  computeMemory(resource.MustParse("30Gi")),
						corev1.ResourceStorage: resource.MustParse("500Gi"),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceMemory: computeMemoryLimit(resource.MustParse("30Gi")),
					},
				},
			},
		},
		"PD component on AWS": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					Type: consts.ProviderTypeAWS,
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.xlarge", CPU: resource.MustParse("4"), Memory: resource.MustParse("8Gi")},
					},
				},
			},
			component: &cloudv1.PD{
				InstanceType: "c5.xlarge",
			},
			expectError: false,
			expectSpec: &cloudv1.ResourceSpec{
				StorageClass: consts.PDGp3StorageClassName,
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("3"),
						corev1.ResourceMemory: computeMemory(resource.MustParse("8Gi")),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceMemory: computeMemoryLimit(resource.MustParse("8Gi")),
					},
				},
			},
		},
		"TiDB component on GCP": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					Type: consts.ProviderTypeGCP,
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "n1-standard-4", CPU: resource.MustParse("4"), Memory: resource.MustParse("15Gi")},
					},
				},
			},
			component: &cloudv1.TiDB{
				InstanceType: "n1-standard-4",
			},
			expectError: false,
			expectSpec: &cloudv1.ResourceSpec{
				StorageClass: consts.DefaultStorageClassName,
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("3"),
						corev1.ResourceMemory: computeMemory(resource.MustParse("15Gi")),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceMemory: computeMemoryLimit(resource.MustParse("15Gi")),
					},
				},
			},
		},
		"unknown instance type": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					Type: consts.ProviderTypeAWS,
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.2xlarge", CPU: resource.MustParse("8"), Memory: resource.MustParse("16Gi")},
					},
				},
			},
			component: &cloudv1.TiKV{
				InstanceType: "unknown.type",
			},
			expectError: true,
		},
	}

	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			configProvider := NewHardcode()

			resourceSpec, err := configProvider.GetResourceSpec(tc.provider, tc.component)

			if tc.expectError {
				g.Expect(err).Should(HaveOccurred())
			} else {
				g.Expect(err).Should(Succeed())
				g.Expect(resourceSpec.StorageClass).To(Equal(tc.expectSpec.StorageClass))

				// Validate CPU request
				cpuReq := resourceSpec.Resources.Requests[corev1.ResourceCPU]
				expectedCPU := tc.expectSpec.Resources.Requests[corev1.ResourceCPU]
				g.Expect(cpuReq.Value()).To(
					Equal(expectedCPU.Value()),
					"CPU request should match expected value",
				)

				// Validate memory request
				memReq := resourceSpec.Resources.Requests[corev1.ResourceMemory]
				expectedMemReq := tc.expectSpec.Resources.Requests[corev1.ResourceMemory]
				g.Expect(memReq.Value()).To(
					Equal(expectedMemReq.Value()),
					"Memory request should match expected value",
				)

				// Validate memory limit
				memLimit := resourceSpec.Resources.Limits[corev1.ResourceMemory]
				expectedMemLimit := tc.expectSpec.Resources.Limits[corev1.ResourceMemory]
				g.Expect(memLimit.Value()).To(
					Equal(expectedMemLimit.Value()),
					"Memory limit should match expected value",
				)

				// Validate storage if present
				if storageSize, ok := tc.component.GetStorageSize(); ok {
					g.Expect(resourceSpec.Resources.Requests[corev1.ResourceStorage]).To(Equal(*storageSize))
				}
			}
		})
	}
}

func TestGetCPUMemResourceRequirements(t *testing.T) {
	cases := map[string]struct {
		provider         InstanceTypeProvider
		instanceType     string
		expectError      bool
		expectCPU        resource.Quantity
		expectMemRequest resource.Quantity
		expectMemLimit   resource.Quantity
	}{
		"small instance 4 CPU 8Gi memory": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.xlarge", CPU: resource.MustParse("4"), Memory: resource.MustParse("8Gi")},
					},
				},
			},
			instanceType:     "c5.xlarge",
			expectError:      false,
			expectCPU:        resource.MustParse("3"),
			expectMemRequest: computeMemory(resource.MustParse("8Gi")),
			expectMemLimit:   computeMemoryLimit(resource.MustParse("8Gi")),
		},
		"medium instance 8 CPU 16Gi memory": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.2xlarge", CPU: resource.MustParse("8"), Memory: resource.MustParse("16Gi")},
					},
				},
			},
			instanceType:     "c5.2xlarge",
			expectError:      false,
			expectCPU:        resource.MustParse("7"),
			expectMemRequest: computeMemory(resource.MustParse("16Gi")),
			expectMemLimit:   computeMemoryLimit(resource.MustParse("16Gi")),
		},
		"large instance 32 CPU 128Gi memory": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.9xlarge", CPU: resource.MustParse("32"), Memory: resource.MustParse("128Gi")},
					},
				},
			},
			instanceType:     "c5.9xlarge",
			expectError:      false,
			expectCPU:        resource.MustParse("31"),
			expectMemRequest: computeMemory(resource.MustParse("128Gi")),
			expectMemLimit:   computeMemoryLimit(resource.MustParse("128Gi")),
		},
		"instance with unavailable memory": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					InstanceTypes: []cloudv1.InstanceType{
						{
							Name:              "custom",
							CPU:               resource.MustParse("16"),
							Memory:            resource.MustParse("64Gi"),
							UnavailableMemory: resource.MustParse("2Gi"),
						},
					},
				},
			},
			instanceType:     "custom",
			expectError:      false,
			expectCPU:        resource.MustParse("15"),
			expectMemRequest: computeMemory(resource.MustParse("62Gi")),
			expectMemLimit:   computeMemoryLimit(resource.MustParse("62Gi")),
		},
		"single CPU instance": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "t3.small", CPU: resource.MustParse("1"), Memory: resource.MustParse("2Gi")},
					},
				},
			},
			instanceType:     "t3.small",
			expectError:      false,
			expectCPU:        resource.MustParse("0"),
			expectMemRequest: computeMemory(resource.MustParse("2Gi")),
			expectMemLimit:   computeMemoryLimit(resource.MustParse("2Gi")),
		},
		"unknown instance type": {
			provider: &cloudv1.Provider{
				Spec: cloudv1.ProviderSpec{
					InstanceTypes: []cloudv1.InstanceType{
						{Name: "c5.2xlarge", CPU: resource.MustParse("8"), Memory: resource.MustParse("16Gi")},
					},
				},
			},
			instanceType: "unknown.type",
			expectError:  true,
		},
	}

	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			configProvider := NewHardcode()

			requirements, err := configProvider.GetCPUMemResourceRequirements(tc.provider, tc.instanceType)

			if tc.expectError {
				g.Expect(err).Should(HaveOccurred())
			} else {
				g.Expect(err).Should(Succeed())
				cpuReq := requirements.Requests[corev1.ResourceCPU]
				expectedCPU := tc.expectCPU
				g.Expect(cpuReq.Value()).To(Equal(expectedCPU.Value()))
				memReq := requirements.Requests[corev1.ResourceMemory]
				expectedMemReq := tc.expectMemRequest
				g.Expect(memReq.Value()).To(Equal(expectedMemReq.Value()))
				memLimit := requirements.Limits[corev1.ResourceMemory]
				expectedMemLimit := tc.expectMemLimit
				g.Expect(memLimit.Value()).To(Equal(expectedMemLimit.Value()))
			}
		})
	}
}