package aws

import (
	"reflect"
	"testing"

	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"

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

func TestGetRoleFromAccount(t *testing.T) {
	externalID := "123456"
	tests := []struct {
		name    string
		account *cloudv1.Account
		wantAr  *AssumedRole
	}{
		{
			name:    "empty account",
			account: &cloudv1.Account{},
			wantAr:  nil,
		},
		{
			name:    "account is nil",
			account: nil,
			wantAr:  nil,
		},
		{
			name: "wrong provider type",
			account: &cloudv1.Account{
				Spec: cloudv1.AccountSpec{
					Provider: cloudv1.CloudProviderGCP,
				},
			},
			wantAr: nil,
		},
		{
			name: "without external id",
			account: &cloudv1.Account{
				Spec: cloudv1.AccountSpec{
					Provider: cloudv1.CloudProviderAWS,
					AWS: &cloudv1.AWSConfig{
						AssumeRoleARN: "arn:aws:iam::123456789012:role/assume-role",
					},
				},
			},
			wantAr: &AssumedRole{
				RoleARN: "arn:aws:iam::123456789012:role/assume-role",
			},
		},
		{
			name: "with external id",
			account: &cloudv1.Account{
				Spec: cloudv1.AccountSpec{
					Provider: cloudv1.CloudProviderAWS,
					AWS: &cloudv1.AWSConfig{
						AssumeRoleARN: "arn:aws:iam::123456789012:role/assume-role",
						ExternalID:    &externalID,
					},
				},
			},
			wantAr: &AssumedRole{
				RoleARN:    "arn:aws:iam::123456789012:role/assume-role",
				ExternalID: &externalID,
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if gotAr := GetRoleFromAccount(tt.account); !reflect.DeepEqual(gotAr, tt.wantAr) {
				t.Errorf("GetRoleFromAccount() = %v, want %v", gotAr, tt.wantAr)
			}
		})
	}
}

func TestFindSubnetIdsByZones(t *testing.T) {
	type args struct {
		subnets []ec2types.Subnet
		zones   []string
	}
	tests := []struct {
		name string
		args args
		want []string
	}{
		{
			name: "should get one subnet id",
			args: args{
				subnets: []ec2types.Subnet{
					{AvailabilityZone: &[]string{"us-west-1a"}[0], SubnetId: &[]string{"subnet-1"}[0]},
					{AvailabilityZone: &[]string{"us-west-1b"}[0], SubnetId: &[]string{"subnet-2"}[0]},
					{AvailabilityZone: &[]string{"us-west-1c"}[0], SubnetId: &[]string{"subnet-3"}[0]},
				},
				zones: []string{"us-west-1a"},
			},
			want: []string{"subnet-1"},
		},
		{
			name: "should get two subnet ids",
			args: args{
				subnets: []ec2types.Subnet{
					{AvailabilityZone: &[]string{"us-west-1a"}[0], SubnetId: &[]string{"subnet-1"}[0]},
					{AvailabilityZone: &[]string{"us-west-1b"}[0], SubnetId: &[]string{"subnet-2"}[0]},
					{AvailabilityZone: &[]string{"us-west-1c"}[0], SubnetId: &[]string{"subnet-3"}[0]},
				},
				zones: []string{"us-west-1a", "us-west-1b"},
			},
			want: []string{"subnet-1", "subnet-2"},
		},
		{
			name: "should get no subnet id",
			args: args{
				subnets: []ec2types.Subnet{
					{AvailabilityZone: &[]string{"us-west-1a"}[0], SubnetId: &[]string{"subnet-1"}[0]},
					{AvailabilityZone: &[]string{"us-west-1b"}[0], SubnetId: &[]string{"subnet-2"}[0]},
					{AvailabilityZone: &[]string{"us-west-1c"}[0], SubnetId: &[]string{"subnet-3"}[0]},
				},
				zones: []string{"us-east-1a"},
			},
			want: nil,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := FindSubnetIdsByZones(tt.args.subnets, tt.args.zones); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("FindSubnetIdsByZones() = %v, want %v", got, tt.want)
			}
		})
	}
}