package common

import (
	"testing"

	"github.com/stretchr/testify/assert"

	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/infra/repoimpl/schema"
)

func TestGetInstanceCpuGeneration(t *testing.T) {
	t.Run("AwsArm64InstanceWith6gReturnsGraviton2", func(t *testing.T) {
		t.Run("m6g.xlarge", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "aws"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "m6g.xlarge",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON2, cpuGen)
		})

		t.Run("c6g.xlarge", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "aws"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "c6g.xlarge",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON2, cpuGen)
		})
	})

	t.Run("AwsArm64InstanceWith8gReturnsGraviton4", func(t *testing.T) {
		t.Run("m8g.large", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "aws"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "m8g.large",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON4, cpuGen)
		})

		t.Run("r8g.2xlarge", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "aws"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "r8g.2xlarge",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GRAVITON4, cpuGen)
		})
	})

	t.Run("NonAwsProviderReturnsUnspecified", func(t *testing.T) {
		t.Run("n2-standard-8", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "n2-standard-8",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})

		t.Run("Standard_D8pls_v5", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "azure"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "Standard_D8pls_v5",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})
	})

	t.Run("NonArm64CpuArchReturnsUnspecified", func(t *testing.T) {
		t.Run("c5.2xlarge", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "aws"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "c5.2xlarge",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})

		t.Run("n2-highmem-16", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "n2-highmem-16",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})

		t.Run("Standard_D2s_v5", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "azure"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "Standard_D2s_v5",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})
	})

	t.Run("UnknownAwsArm64InstanceTypeReturnsError", func(t *testing.T) {
		t.Run("unknown.large", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "aws"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "unknown.large",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.NotNil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})
	})

	t.Run("GcpN4Instance", func(t *testing.T) {
		t.Run("n4-standard-4", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "n4-standard-4",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON, cpuGen)
		})

		t.Run("n4-highcpu-8", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "n4-highcpu-8",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_INTEL_XEON, cpuGen)
		})
	})

	t.Run("GcpC4AInstance", func(t *testing.T) {
		t.Run("c4a-standard-4", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "c4a-standard-4",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_AXION, cpuGen)
		})

		t.Run("c4a-highcpu-8", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "arm64",
					InstanceType: "c4a-highcpu-8",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_GENERATION4_AXION, cpuGen)
		})
	})

	t.Run("Gcp Old Instance", func(t *testing.T) {
		t.Run("n1-standard-2", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "n1-standard-2",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})

		t.Run("t2d-standard-16", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "t2d-standard-16",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})
	})

	t.Run("GcpN2Instance", func(t *testing.T) {
		t.Run("n2-standard-8", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "n2-standard-8",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})

		t.Run("n2-highmem-16", func(t *testing.T) {
			componentTierUnit := &schema.ComponentTierUnit{
				Provider: schema.Provider{Name: "gcp"},
				InstanceType: schema.InstanceType{
					CpuArch:      "amd64",
					InstanceType: "n2-highmem-16",
				},
			}
			cpuGen, err := GetInstanceCpuGeneration(componentTierUnit)
			assert.Nil(t, err)
			assert.Equal(t, globaldedicatedv1.InstanceMu_CPU_GENERATION_UNSPECIFIED, cpuGen)
		})
	})
}