package eks

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	ec2v1beta1 "github.com/upbound/provider-aws/apis/ec2/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"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/consts"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/nodepool/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/labels"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/namer"
	mgrscheme "github.com/tidbcloud/infra-provider/regional-server/pkg/scheme"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
)

func TestPlacementGroupNameUsesSeparatedAppBoxIdentity(t *testing.T) {
	pgTask := newPlacementGroupTestTask(t)
	require.Equal(t, namer.PlacementGroupName("m6i.2xlarge", "ns.appbox"), pgTask.placementGroupName())
}

func TestPlacementGroupCreateAddsCloudTags(t *testing.T) {
	ctx := context.Background()
	pgTask := newPlacementGroupTestTask(t)

	result := pgTask.Create(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())

	pg := &ec2v1beta1.PlacementGroup{}
	require.NoError(t, pgTask.Local.Get(ctx, types.NamespacedName{Name: pgTask.placementGroupName()}, pg))
	tags := pg.Spec.ForProvider.Tags
	require.Nil(t, tags["shared-pool"])
	require.Equal(t, ptrString("nodepool-pool"), tags["shared-pool-nodepool"])
	require.Equal(t, ptrString("prod"), tags["env"])
	require.Equal(t, ptrString(pgTask.placementGroupName()), tags["Name"])
}

func TestPlacementGroupPostDeleteSkipsWaitWhenOtherLiveUserRemains(t *testing.T) {
	ctx := context.Background()
	pgTask := newPlacementGroupTestTask(t)
	pgTask.NodePool.Spec.PlacementGroupEnabled = true
	ownerLabels := labels.NewOwnerLabels(pgTask.APPBox)
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")

	require.NoError(t, pgTask.Local.Create(ctx, &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{Name: "other", Namespace: "ns", Labels: ownerLabels},
		Spec: v1alpha1.NodePoolSpec{
			InstanceType:          "m6i.2xlarge",
			PlacementGroupEnabled: true,
		},
	}))
	require.NoError(t, pgTask.Local.Create(ctx, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}}))

	result := pgTask.PostDelete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())
}

func TestPlacementGroupDeleteSkipsDeletionWhenOtherLiveUserRemains(t *testing.T) {
	ctx := context.Background()
	pgTask := newPlacementGroupTestTask(t)
	pgTask.NodePool.Spec.PlacementGroupEnabled = true
	ownerLabels := labels.NewOwnerLabels(pgTask.APPBox)
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")

	require.NoError(t, pgTask.Local.Create(ctx, &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{Name: "other", Namespace: "ns", Labels: ownerLabels},
		Spec: v1alpha1.NodePoolSpec{
			InstanceType:          "m6i.2xlarge",
			PlacementGroupEnabled: true,
		},
	}))
	require.NoError(t, pgTask.Local.Create(ctx, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}}))

	result := pgTask.Delete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())
	require.NoError(t, pgTask.Local.Get(ctx, types.NamespacedName{Name: pgName}, &ec2v1beta1.PlacementGroup{}))
}

func TestPlacementGroupDeleteUsesNodePoolLabelsWhenAPPBoxMissing(t *testing.T) {
	ctx := context.Background()
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")
	pgTask := newPlacementGroupTestTaskWithContext(t, &task.Context{
		NodePool: &v1alpha1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "current",
				Namespace: "ns",
				Labels: map[string]string{
					v1alpha1.APPBoxNameKey:  "appbox",
					v1alpha1.DataPlaneIDKey: "dp-1",
				},
				UID: "current-uid",
			},
			Spec: v1alpha1.NodePoolSpec{
				InstanceType:          "m6i.2xlarge",
				PlacementGroupEnabled: true,
			},
		},
		APPBox: nil,
		VPC:    nil,
		Config: config.RootConfig{
			AdditionalCloudTags: map[string]string{"env": "prod"},
		},
	}, &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "other",
			Namespace: "ns",
			Labels: map[string]string{
				v1alpha1.APPBoxNameKey:  "appbox",
				v1alpha1.DataPlaneIDKey: "dp-1",
			},
		},
		Spec: v1alpha1.NodePoolSpec{
			InstanceType:          "m6i.2xlarge",
			PlacementGroupEnabled: true,
		},
	}, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}})

	result := pgTask.Delete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())
	require.NoError(t, pgTask.Local.Get(ctx, types.NamespacedName{Name: pgName}, &ec2v1beta1.PlacementGroup{}))
}

func TestPlacementGroupDeleteWithoutSiblingUsesLabelDerivedNameWhenAPPBoxMissing(t *testing.T) {
	ctx := context.Background()
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")
	pgTask := newPlacementGroupTestTaskWithContext(t, &task.Context{
		NodePool: &v1alpha1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "current",
				Namespace: "ns",
				Labels: map[string]string{
					v1alpha1.APPBoxNameKey:  "appbox",
					v1alpha1.DataPlaneIDKey: "dp-1",
				},
				UID: "current-uid",
			},
			Spec: v1alpha1.NodePoolSpec{
				InstanceType:          "m6i.2xlarge",
				PlacementGroupEnabled: true,
			},
		},
		APPBox: nil,
		VPC:    nil,
		Config: config.RootConfig{
			AdditionalCloudTags: map[string]string{"env": "prod"},
		},
	}, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}})

	result := pgTask.Delete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())
	require.Error(t, pgTask.Local.Get(ctx, types.NamespacedName{Name: pgName}, &ec2v1beta1.PlacementGroup{}))
}

func TestPlacementGroupDeleteDeletesResourceWhenNoOtherLiveUserRemains(t *testing.T) {
	ctx := context.Background()
	pgTask := newPlacementGroupTestTask(t)
	pgTask.NodePool.Spec.PlacementGroupEnabled = true
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")

	require.NoError(t, pgTask.Local.Create(ctx, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}}))

	result := pgTask.Delete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())
	require.Error(t, pgTask.Local.Get(ctx, types.NamespacedName{Name: pgName}, &ec2v1beta1.PlacementGroup{}))
}

func TestPlacementGroupPostDeleteWaitsWhenNoOtherLiveUserRemains(t *testing.T) {
	ctx := context.Background()
	pgTask := newPlacementGroupTestTask(t)
	pgTask.NodePool.Spec.PlacementGroupEnabled = true
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")
	require.NoError(t, pgTask.Local.Create(ctx, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}}))

	result := pgTask.PostDelete(ctx)
	require.Error(t, result.Error())
}

// TestPlacementGroupDeleteRunsWithNilVPC is the regression test for the
// deletion-time PG orphan bug. PG eligibility is read from NodePool spec,
// not from dataplane or VPC, so Delete must still drive the AWS cleanup
// even when the reconciler has skipped loadDependencies on a Delete
// reconcile (VPC is nil).
func TestPlacementGroupDeleteRunsWithNilVPC(t *testing.T) {
	ctx := context.Background()
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")
	pgTask := newPlacementGroupTestTaskWithContext(t, &task.Context{
		NodePool: &v1alpha1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "current",
				Namespace: "ns",
				Labels: map[string]string{
					v1alpha1.APPBoxNameKey:  "appbox",
					v1alpha1.DataPlaneIDKey: "dp-1",
				},
				UID: "current-uid",
			},
			Spec: v1alpha1.NodePoolSpec{
				InstanceType:          "m6i.2xlarge",
				PlacementGroupEnabled: true,
			},
		},
		APPBox: nil,
		VPC:    nil,
		Config: config.RootConfig{
			AdditionalCloudTags: map[string]string{"env": "prod"},
		},
	}, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}})

	require.True(t, placementGroupEnabled(pgTask.NodePool))
	result := pgTask.Delete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())
	require.Error(t, pgTask.Local.Get(ctx, types.NamespacedName{Name: pgName}, &ec2v1beta1.PlacementGroup{}))
}

// TestPlacementGroupSkippedForLegacyPGDisabledPool is the regression test
// for the silent-migration invariant: PG eligibility is the durable record
// on NodePool spec. A legacy single-AZ EKS pool whose spec carries the
// historical PlacementGroupEnabled=false must stay out of the partition
// path on every reconcile, regardless of the current dataplane topology.
// Recreating the NodePool is the only way to flip this value.
func TestPlacementGroupSkippedForLegacyPGDisabledPool(t *testing.T) {
	ctx := context.Background()
	pgName := namer.PlacementGroupName("m6i.2xlarge", "ns.appbox")
	pgTask := newPlacementGroupTestTaskWithContext(t, &task.Context{
		NodePool: &v1alpha1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "current",
				Namespace: "ns",
				Labels: map[string]string{
					v1alpha1.APPBoxNameKey:  "appbox",
					v1alpha1.DataPlaneIDKey: "dp-1",
				},
				UID: "current-uid",
			},
			Spec: v1alpha1.NodePoolSpec{
				InstanceType:          "m6i.2xlarge",
				PlacementGroupEnabled: false,
			},
		},
		DataPlane: &v1alpha1.DataPlane{
			ObjectMeta: metav1.ObjectMeta{Name: "dp-1", Namespace: "ns"},
			// A multi-AZ dataplane must NOT migrate the legacy single-AZ,
			// PG-disabled NodePool into the partition path. Spec value is
			// the durable record; topology changes after creation do not
			// flip it.
			Spec: v1alpha1.DataPlaneSpec{NodeZones: []string{
				"us-west-2a", "us-west-2b", "us-west-2c",
			}},
		},
		APPBox: nil,
		VPC:    nil,
		Config: config.RootConfig{
			AdditionalCloudTags: map[string]string{"env": "prod"},
		},
	}, &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: pgName}})

	result := pgTask.Delete(ctx)
	require.NoError(t, result.Error())
	require.True(t, result.IsDoneState())
	// A pre-existing PlacementGroup must NOT be deleted for a
	// PG-disabled legacy pool: the controller skips Delete entirely.
	require.NoError(t, pgTask.Local.Get(ctx, types.NamespacedName{Name: pgName}, &ec2v1beta1.PlacementGroup{}))
}

func TestPlacementGroupUsersExcludeCurrentAndDeletingNodePools(t *testing.T) {
	ctx := context.Background()
	ownerLabels := map[string]string{
		consts.LabelKeyOwnedBy:   "appbox",
		consts.LabelKeyNamespace: "ns",
		consts.LabelKeyName:      "appbox",
		consts.LabelKeyUID:       "appbox-uid",
	}
	now := metav1.NewTime(time.Now())
	pgTask := newPlacementGroupTestTask(t,
		&v1alpha1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:              "deleting",
				Namespace:         "ns",
				Labels:            ownerLabels,
				DeletionTimestamp: &now,
				Finalizers:        []string{"test"},
			},
			Spec: v1alpha1.NodePoolSpec{
				InstanceType:          "m6i.2xlarge",
				PlacementGroupEnabled: true,
			},
		},
		&v1alpha1.NodePool{
			ObjectMeta: metav1.ObjectMeta{
				Name:      "other",
				Namespace: "ns",
				Labels:    ownerLabels,
			},
			Spec: v1alpha1.NodePoolSpec{
				InstanceType:          "m6i.2xlarge",
				PlacementGroupEnabled: true,
			},
		},
	)

	users, err := pgTask.otherLivePlacementGroupUsers(ctx)
	require.NoError(t, err)
	require.Equal(t, 1, users)
}

func TestPlacementGroupUsersExcludeDifferentOwner(t *testing.T) {
	ctx := context.Background()
	pgTask := newPlacementGroupTestTask(t)
	pgTask.NodePool.Spec.PlacementGroupEnabled = true

	// NodePool with different owner labels (e.g. different APPBox)
	require.NoError(t, pgTask.Local.Create(ctx, &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "different-owner",
			Namespace: "ns",
			Labels: map[string]string{
				consts.LabelKeyOwnedBy:   "appbox",
				consts.LabelKeyNamespace: "ns",
				consts.LabelKeyName:      "other-appbox",
				consts.LabelKeyUID:       "other-appbox-uid",
			},
		},
		Spec: v1alpha1.NodePoolSpec{
			InstanceType: "m6i.2xlarge",
		},
	}))

	// NodePool with same APPBox owner and instance type shares the same placement group
	require.NoError(t, pgTask.Local.Create(ctx, &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "same-appbox-different-nodepool",
			Namespace: "ns",
			Labels: map[string]string{
				consts.LabelKeyOwnedBy:   "appbox",
				consts.LabelKeyNamespace: "ns",
				consts.LabelKeyName:      "appbox",
				consts.LabelKeyUID:       "appbox-uid",
			},
		},
		Spec: v1alpha1.NodePoolSpec{
			InstanceType:          "m6i.2xlarge",
			PlacementGroupEnabled: true,
		},
	}))

	users, err := pgTask.otherLivePlacementGroupUsers(ctx)
	require.NoError(t, err)
	require.Equal(t, 1, users)
}

func newPlacementGroupTestTask(t *testing.T, objects ...client.Object) *placementGroup {
	t.Helper()
	require.NoError(t, ec2v1beta1.SchemeBuilder.AddToScheme(mgrscheme.Scheme))
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))
	require.NoError(t, ec2v1beta1.SchemeBuilder.AddToScheme(scheme))

	appBox := &v1alpha1.APPBox{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "appbox",
			Namespace: "ns",
			UID:       "appbox-uid",
			Labels: map[string]string{
				consts.SharedPoolLabelKey: "pool-1",
			},
		},
		Spec: v1alpha1.APPBoxSpec{Region: "us-west-2"},
	}
	nodePool := &v1alpha1.NodePool{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "current",
			Namespace: "ns",
			Labels: util.MergeMaps(
				labels.NewOwnerLabels(appBox),
				map[string]string{
					"tags.tidbcloud.com/shared-pool-nodepool": "nodepool-pool",
				},
			),
			UID: "current-uid",
		},
		Spec: v1alpha1.NodePoolSpec{
			InstanceType:          "m6i.2xlarge",
			// PG-enabled by default so the helper exercises the partition
			// path. Tests that need the non-PG path set this to false on
			// the returned context.
			PlacementGroupEnabled: true,
		},
	}

	return newPlacementGroupTestTaskWithContext(t, &task.Context{
		NodePool: nodePool,
		APPBox:   appBox,
		Config: config.RootConfig{
			AdditionalCloudTags: map[string]string{"env": "prod"},
		},
	}, objects...)
}

func newPlacementGroupTestTaskWithContext(t *testing.T, baseCtx *task.Context, objects ...client.Object) *placementGroup {
	t.Helper()
	require.NoError(t, ec2v1beta1.SchemeBuilder.AddToScheme(mgrscheme.Scheme))
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))
	require.NoError(t, ec2v1beta1.SchemeBuilder.AddToScheme(scheme))

	allObjects := append([]client.Object{baseCtx.NodePool}, objects...)
	cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(allObjects...).Build()
	baseCtx.Local = cli
	ctx, err := baseCtx.Default()
	require.NoError(t, err)
	return &placementGroup{BaseTask: &task.BaseTask{Context: ctx}}
}

func ptrString(value string) *string {
	return &value
}