package task

import (
	"context"
	"testing"

	"github.com/crossplane-contrib/provider-aws/apis/eks/manualv1alpha1"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	corev1 "k8s.io/api/core/v1"
	storagev1 "k8s.io/api/storage/v1"

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

// fakeConfig implements config.CloudConfigProvider with a fixed instance type list.
type fakeConfig struct {
	instanceTypes []cloudv1.InstanceType
}

func (f *fakeConfig) GetProviderType() string { return "aws" }
func (f *fakeConfig) GetInstanceTypes() ([]cloudv1.InstanceType, error) {
	return f.instanceTypes, nil
}
func (f *fakeConfig) GetStorageClasses() ([]*storagev1.StorageClass, error) { return nil, nil }
func (f *fakeConfig) FindInstanceType(name string) (cloudv1.InstanceType, error) {
	return cloudv1.InstanceType{}, nil
}
func (f *fakeConfig) FindRegionZones(region string) ([]string, error) { return nil, nil }

func TestBuildNodePool_AMIType(t *testing.T) {
	tests := []struct {
		name          string
		instanceType  string
		instanceTypes []cloudv1.InstanceType
		wantAMIType   string
	}{
		{
			name:         "x86 instance resolves to x86 AMI",
			instanceType: "c5.xlarge",
			instanceTypes: []cloudv1.InstanceType{
				{Name: "c5.xlarge", Arch: "amd64"},
			},
			wantAMIType: nodegroup.BottleRocketX86AMIType,
		},
		{
			name:         "arm64 instance resolves to ARM AMI",
			instanceType: "c6g.xlarge",
			instanceTypes: []cloudv1.InstanceType{
				{Name: "c6g.xlarge", Arch: "arm64"},
			},
			wantAMIType: nodegroup.BottleRocketARMAMIType,
		},
		{
			name:          "empty config falls back to x86 AMI",
			instanceType:  "c5.xlarge",
			instanceTypes: nil,
			wantAMIType:   nodegroup.BottleRocketX86AMIType,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cfg := &fakeConfig{instanceTypes: tt.instanceTypes}
			provider, err := NewAWSProvider(
				"s3://bucket/prefix",
				"us-west-2",
				"kube-cluster",
				"1.30",
				"nodeRole",
				"accountID",
				"assumeRoleARN",
				"externalID",
				"providerConfigName",
				"tenantID",
				"projectID",
				"clusterID",
				cloudv1.CloudAWS{},
				cfg,
			)
			require.NoError(t, err)

			np := &NodePool{
				Name:          "test-pool",
				InstanceType:  tt.instanceType,
				Region:        "us-west-2",
				KubeClusterID: "kube-cluster",
				Labels:        map[string]string{},
				Tags:          map[string]string{},
				Taints:        []corev1.Taint{},
			}
			objs, err := provider.BuildNodePool(context.Background(), "providerConfig", np)
			require.NoError(t, err)
			require.Len(t, objs, 2)

			// Second object is the NodeGroup
			ng := objs[1].(*manualv1alpha1.NodeGroup)
			require.NotNil(t, ng.Spec.ForProvider.AMIType)
			assert.Equal(t, tt.wantAMIType, *ng.Spec.ForProvider.AMIType)
		})
	}
}