package gitops

import (
	"testing"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/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))
		}
	}
}