package eks

import (
	"context"
	"strings"
	"testing"

	ccpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	ec2v1beta1 "github.com/upbound/provider-aws/apis/ec2/v1beta1"
	eksv1beta1 "github.com/upbound/provider-aws/apis/eks/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/config"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/nodepool/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/labels"
)

func TestNodeGroup_ZoneSubnetIDs(t *testing.T) {
	newTask := func(regionConfig config.RegionConfig, vpcSpecZones []string, vpcSpecNodeZones []string) *nodeGroup {
		return &nodeGroup{BaseTask: &task.BaseTask{Context: &task.Context{
			Config: config.RootConfig{
				ControllerConfig: config.ControllerConfig{RegionConfig: regionConfig},
			},
			NodePool: &v1alpha1.NodePool{},
			VPC: &v1alpha1.VPC{
				Spec: v1alpha1.VPCSpec{Zones: vpcSpecZones, NodeZones: vpcSpecNodeZones},
				Status: v1alpha1.VPCStatus{NodeSubnets: []v1alpha1.Subnet{
					{AvailabilityZone: "us-west-2a", SubnetID: "subnet-a"},
					{AvailabilityZone: "us-west-2b", SubnetID: "subnet-b"},
					{AvailabilityZone: "us-west-2c", SubnetID: "subnet-c"},
				}},
			},
		}}}
	}

	t.Run("prefer workload zones when configured", func(t *testing.T) {
		task := newTask(config.RegionConfig{
			Zones:     []string{"us-west-2a", "us-west-2b", "us-west-2c"},
			NodeZones: []string{"us-west-2a", "us-west-2b"},
		}, []string{"us-west-2a", "us-west-2b", "us-west-2c"}, []string{"us-west-2a", "us-west-2b"})

		zoneToSubnetID, zones, err := task.zoneSubnetIDs()
		assert.NoError(t, err)
		assert.Equal(t, []string{"us-west-2a", "us-west-2b"}, zones)
		assert.Equal(t, "subnet-a", zoneToSubnetID["us-west-2a"])
		assert.Equal(t, "subnet-b", zoneToSubnetID["us-west-2b"])
	})

	t.Run("fallback to vpc zones when workload zones empty", func(t *testing.T) {
		task := newTask(config.RegionConfig{Zones: []string{"us-west-2a", "us-west-2b", "us-west-2c"}}, []string{"us-west-2a", "us-west-2c"}, nil)

		_, zones, err := task.zoneSubnetIDs()
		assert.NoError(t, err)
		assert.Equal(t, []string{"us-west-2a", "us-west-2c"}, zones)
	})

	t.Run("error when workload zone subnet missing", func(t *testing.T) {
		task := newTask(config.RegionConfig{
			Zones:     []string{"us-west-2a", "us-west-2b", "us-west-2c"},
			NodeZones: []string{"us-west-2a", "us-west-2d"},
		}, []string{"us-west-2a", "us-west-2b", "us-west-2c"}, []string{"us-west-2a", "us-west-2d"})

		_, _, err := task.zoneSubnetIDs()
		assert.Error(t, err)
	})
}

func TestNodeGroup_PerZonePerInstanceTypeNodegroupName_ResourcePrefix(t *testing.T) {
	newTask := func(configPrefix, nodePoolPrefix string) *nodeGroup {
		return &nodeGroup{BaseTask: &task.BaseTask{Context: &task.Context{
			Config: config.RootConfig{
				ResourceNamePrefix: configPrefix,
			},
			NodePool: &v1alpha1.NodePool{
				Spec: v1alpha1.NodePoolSpec{
					ResourceNamePrefix: nodePoolPrefix,
				},
			},
			APPBox: &v1alpha1.APPBox{},
		}}}
	}

	t.Run("use nodepool spec resource prefix when provided", func(t *testing.T) {
		ngTask := newTask("tidbcloud", "custom")
		ngTask.APPBox.Namespace = "ns"
		ngTask.APPBox.Name = "appbox"
		name := ngTask.perZonePerInstanceTypeNodegroupName("us-west-2a", "c8g.xlarge", "1.34", 8, 16)
		assert.True(t, strings.Contains(name, "custom-c8g-xlarge-8c16g"))
	})

	t.Run("fallback to config resource prefix when nodepool prefix is empty", func(t *testing.T) {
		ngTask := newTask("tidbcloud", "")
		ngTask.APPBox.Namespace = "ns"
		ngTask.APPBox.Name = "appbox"
		name := ngTask.perZonePerInstanceTypeNodegroupName("us-west-2a", "c8g.xlarge", "1.34", 8, 16)
		assert.True(t, strings.Contains(name, "tidbcloud-c8g-xlarge-8c16g"))
	})
}

func TestNodeGroup_SingleZoneSubnetID(t *testing.T) {
	ngTask := &nodeGroup{BaseTask: &task.BaseTask{Context: &task.Context{
		NodePool: &v1alpha1.NodePool{},
		VPC: &v1alpha1.VPC{
			Spec: v1alpha1.VPCSpec{Zones: []string{"us-west-2a"}},
			Status: v1alpha1.VPCStatus{NodeSubnets: []v1alpha1.Subnet{
				{AvailabilityZone: "us-west-2a", SubnetID: "subnet-a"},
			}},
		},
	}}}

	zone, subnetID, err := ngTask.singleZoneSubnetID()
	assert.NoError(t, err)
	assert.Equal(t, "us-west-2a", zone)
	assert.Equal(t, "subnet-a", subnetID)
}

func TestNodeGroup_SingleZoneSubnetIDRejectsMultipleZones(t *testing.T) {
	ngTask := &nodeGroup{BaseTask: &task.BaseTask{Context: &task.Context{
		NodePool: &v1alpha1.NodePool{},
		VPC: &v1alpha1.VPC{
			Spec: v1alpha1.VPCSpec{Zones: []string{"us-west-2a", "us-west-2b"}},
			Status: v1alpha1.VPCStatus{NodeSubnets: []v1alpha1.Subnet{
				{AvailabilityZone: "us-west-2a", SubnetID: "subnet-a"},
				{AvailabilityZone: "us-west-2b", SubnetID: "subnet-b"},
			}},
		},
	}}}

	_, _, err := ngTask.singleZoneSubnetID()
	assert.Error(t, err)
}

func TestNodeGroup_PerPartitionNodegroupName(t *testing.T) {
	ngTask := &nodeGroup{BaseTask: &task.BaseTask{Context: &task.Context{
		Config:   config.RootConfig{ResourceNamePrefix: "tidbcloud"},
		NodePool: &v1alpha1.NodePool{Spec: v1alpha1.NodePoolSpec{ResourceNamePrefix: "primary"}},
		APPBox:   &v1alpha1.APPBox{},
	}}}
	ngTask.APPBox.Namespace = "ns"
	ngTask.APPBox.Name = "appbox"

	name := ngTask.perPartitionNodegroupName("us-west-2a", 2, "c8g.xlarge", "1.34", 8, 16)
	assert.True(t, strings.Contains(name, "primary-c8g-xlarge-8c16g"))
	assert.True(t, strings.Contains(name, "us-west-2a-p2-1-34"))
}

func TestNodeGroupCreateUsesOneBasedAWSPlacementPartitionNumber(t *testing.T) {
	ctx := context.Background()
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))
	require.NoError(t, eksv1beta1.SchemeBuilder.AddToScheme(scheme))
	require.NoError(t, ec2v1beta1.SchemeBuilder.AddToScheme(scheme))
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(&eksv1beta1.Cluster{
		ObjectMeta: metav1.ObjectMeta{Name: "test-cluster"},
		Status: eksv1beta1.ClusterStatus{
			AtProvider: eksv1beta1.ClusterObservation{
				VPCConfig: []eksv1beta1.VPCConfigObservation{
					{ClusterSecurityGroupID: ptr.To("sg-cluster")},
				},
			},
		},
	}).Build()

	ngTask := &nodeGroup{BaseTask: &task.BaseTask{Context: &task.Context{
		Local: cli,
		Config: config.RootConfig{
			ResourceNamePrefix: "tidbcloud",
			AdditionalCloudTags: map[string]string{
				"ManagedBy": "PingCAP",
			},
		},
		NodePool: &v1alpha1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "pool",
				Namespace: "ns",
				Labels: map[string]string{
					"tags.tidbcloud.com/shared-pool-nodepool": "nodepool-pool",
				},
			},
			Spec: v1alpha1.NodePoolSpec{
				InstanceType:      "c8g.xlarge",
				KubernetesVersion: "1.34",
				CpuCore:           8,
				MemoryGi:          16,
			},
		},
		APPBox: &v1alpha1.APPBox{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "appbox",
				Namespace: "ns",
				Labels: map[string]string{
					"tags.tidbcloud.com/shared-pool": "appbox-pool",
				},
			},
			Spec:   v1alpha1.APPBoxSpec{Region: "us-west-2"},
			Status: v1alpha1.APPBoxStatus{KubernetesClusterName: "test-cluster"},
		},
	}}}

	partitionIdx := 0
	require.NoError(t, ngTask.createNodeGroupResources(ctx, "ng-p0", "subnet-a", "sg-cluster", map[string]string{"owner": "pool"}, map[string]*string{}, nil, "node-role", "", "BOTTLEROCKET_ARM_64", &partitionIdx))

	lt := &ec2v1beta1.LaunchTemplate{}
	require.NoError(t, cli.Get(ctx, client.ObjectKey{Name: "ng-p0"}, lt))
	require.Equal(t, "PingCAP", ptr.Deref(lt.Spec.ForProvider.Tags["ManagedBy"], ""))
	require.Equal(t, "nodepool-pool", ptr.Deref(lt.Spec.ForProvider.Tags["shared-pool-nodepool"], ""))
	require.Nil(t, lt.Spec.ForProvider.Tags["shared-pool"])
	for _, spec := range lt.Spec.ForProvider.TagSpecifications {
		require.Equal(t, "PingCAP", ptr.Deref(spec.Tags["ManagedBy"], ""))
		require.Equal(t, "nodepool-pool", ptr.Deref(spec.Tags["shared-pool-nodepool"], ""))
		require.Nil(t, spec.Tags["shared-pool"])
	}
	require.Len(t, lt.Spec.ForProvider.Placement, 1)
	require.NotNil(t, lt.Spec.ForProvider.Placement[0].PartitionNumber)
	assert.Equal(t, float64(1), *lt.Spec.ForProvider.Placement[0].PartitionNumber)
	require.Len(t, lt.Spec.ForProvider.VPCSecurityGroupIds, 1)
	assert.Equal(t, "sg-cluster", ptr.Deref(lt.Spec.ForProvider.VPCSecurityGroupIds[0], ""))

	ng := &eksv1beta1.NodeGroup{}
	require.NoError(t, cli.Get(ctx, client.ObjectKey{Name: "ng-p0"}, ng))
	require.NotNil(t, ng.Spec.ForProvider.Labels["tidbcloud.com/pg-partition"])
	assert.Equal(t, "0", *ng.Spec.ForProvider.Labels["tidbcloud.com/pg-partition"])
}

func TestNodeGroupDeleteDeletesPlacementGroupPartitionResources(t *testing.T) {
	ctx := context.Background()
	nodePool := &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{Name: "pool", Namespace: "ns"},
		Spec:       v1alpha1.NodePoolSpec{},
	}
	ownerLabels := labels.NewOwnerLabels(nodePool)
	objects := make([]client.Object, 0, 6)
	for _, name := range []string{"ng-p0", "ng-p1", "ng-p2"} {
		objects = append(objects, &eksv1beta1.NodeGroup{
			ObjectMeta: metav1.ObjectMeta{Name: name, Labels: ownerLabels},
			Spec:       eksv1beta1.NodeGroupSpec{ResourceSpec: ccpv1.ResourceSpec{ManagementPolicies: ccpv1.ManagementPolicies{ccpv1.ManagementActionAll}}},
		})
		objects = append(objects, &ec2v1beta1.LaunchTemplate{
			ObjectMeta: metav1.ObjectMeta{Name: name, Labels: ownerLabels},
			Spec:       ec2v1beta1.LaunchTemplateSpec{ResourceSpec: ccpv1.ResourceSpec{ManagementPolicies: ccpv1.ManagementPolicies{ccpv1.ManagementActionAll}}},
		})
	}
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))
	require.NoError(t, eksv1beta1.SchemeBuilder.AddToScheme(scheme))
	require.NoError(t, ec2v1beta1.SchemeBuilder.AddToScheme(scheme))
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(objects...).Build()
	ngTask := &nodeGroup{BaseTask: &task.BaseTask{Context: &task.Context{
		Local:    cli,
		NodePool: nodePool,
	}}}

	result := ngTask.Delete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	nodeGroups := &eksv1beta1.NodeGroupList{}
	require.NoError(t, cli.List(ctx, nodeGroups, client.MatchingLabels(ownerLabels)))
	require.Empty(t, nodeGroups.Items)
	launchTemplates := &ec2v1beta1.LaunchTemplateList{}
	require.NoError(t, cli.List(ctx, launchTemplates, client.MatchingLabels(ownerLabels)))
	require.Empty(t, launchTemplates.Items)
}