package nodegroup

import (
	"testing"

	"github.com/stretchr/testify/assert"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

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

func buildTestCluster(name, ns, region string) *cloudv1.Cluster {
	return &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
			Labels: map[string]string{
				consts.RegionLabel: region,
			},
		},
		Spec: cloudv1.ClusterSpec{
			Kubernetes: cloudv1.Kubernetes{
				Provider: "eks",
			},
		},
	}
}

func TestGetTiDBClusterNodeGroupName(t *testing.T) {
	tests := []struct {
		name           string
		cluster        *cloudv1.Cluster
		machineType    string
		componentType  string
		k8sVersion     string
		zones          []string
		extraGroupName string
		want           string
	}{
		{
			name:          "dev: pd",
			cluster:       buildTestCluster("10442564784814776972", "nightly-ms-3200228", "us-west-2"),
			machineType:   "t4g.medium",
			componentType: "pd",
			k8sVersion:    "1.32",
			zones:         []string{"us-west-2a"},
			want:          "pd-t4g-medium-nightly-ms-3200228-10442-468d5495-us-west-2a-1-32",
		},
		{
			name:          "dev: tikv",
			cluster:       buildTestCluster("10442564784814776972", "nightly-ms-3200228", "us-west-2"),
			machineType:   "t4g.medium",
			componentType: "tikv",
			k8sVersion:    "1.32",
			zones:         []string{"us-west-2b"},
			want:          "tikv-t4g-medium-nightly-ms-3200228-104-468d5495-us-west-2b-1-32",
		},
		{
			name:          "prod: tidb",
			cluster:       buildTestCluster("10789907685234305975", "prod-1372813089409131283", "ap-southeast-1"),
			machineType:   "m6g.xlarge",
			componentType: "tidb",
			k8sVersion:    "1.31",
			zones:         []string{"ap-southeast-1a"},
			want:          "tidb-m6g-xlarge-prod-137281308940-b44c6953-ap-southeast-1a-1-31",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := GetTiDBClusterNodeGroupNameIn3Layer(tt.cluster, tt.machineType, tt.componentType, tt.k8sVersion, tt.zones, tt.extraGroupName)
			assert.Equal(t, tt.want, got)
		})
	}
}

func TestGetAMITypeByInstanceType(t *testing.T) {
	arm64InstanceType := cloudv1.InstanceType{Name: "m6g.large", Arch: consts.ARM64ArchName}
	x86InstanceType := cloudv1.InstanceType{Name: "m5.large", Arch: consts.AMD64ArchName}
	instanceTypes := []cloudv1.InstanceType{arm64InstanceType, x86InstanceType}

	tests := []struct {
		name         string
		instanceType string
		want         string
		wantFound    bool
	}{
		{
			name:         "arm64 instance type returns ARM AMI",
			instanceType: "m6g.large",
			want:         BottleRocketARMAMIType,
			wantFound:    true,
		},
		{
			name:         "x86 instance type returns x86 AMI",
			instanceType: "m5.large",
			want:         BottleRocketX86AMIType,
			wantFound:    true,
		},
		{
			name:         "unknown instance type falls back to x86",
			instanceType: "unknown.type",
			want:         BottleRocketX86AMIType,
			wantFound:    false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, found := GetAMITypeByInstanceType(instanceTypes, tt.instanceType)
			assert.Equal(t, tt.want, got)
			assert.Equal(t, tt.wantFound, found)
		})
	}
}