package eks

import (
	"context"
	"fmt"
	"maps"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	ec2v1beta1 "github.com/upbound/provider-aws/apis/ec2/v1beta1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/consts"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/nodepool/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/deploy"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/k8s"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/labels"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/namer"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
)

const placementGroupPartitionCount = 3

// placementGroupEnabled returns whether the NodePool is eligible for an EC2
// partition placement group. The decision is durably stored on the NodePool
// spec as Spec.PlacementGroupEnabled: the dataplane gRPC CreateNodePool path
// runs "is AWS + single-AZ?" once at NodePool creation and writes the
// result, and this helper reads it back on every reconcile.
//
// Reading from spec (rather than re-deriving from the live dataplane / VPC
// topology) makes the value the source of truth itself: it survives across
// deletion reconciles (the NodePool CR persists while the parent dataplane
// and VPC may not), preserves the legacy per-pool intent for pools that
// were created before this code existed, and never silently flips
// eligibility on a topology change after creation. See
// docs/design/single-az-ha-placement-groups.md for the full rationale.
func placementGroupEnabled(np *v1alpha1.NodePool) bool {
	if np == nil {
		return false
	}
	return np.Spec.PlacementGroupEnabled
}

type placementGroup struct {
	*task.BaseTask
}

func NewPlacementGroup(ctx *task.Context) deploy.Task {
	return &placementGroup{BaseTask: &task.BaseTask{Context: ctx}}
}

func (t *placementGroup) Name() string {
	return "PlacementGroup"
}

func (t *placementGroup) Create(ctx context.Context) deploy.Result {
	if !placementGroupEnabled(t.NodePool) {
		return deploy.Done()
	}

	name := t.placementGroupName()
	existing := &ec2v1beta1.PlacementGroup{}
	if err := t.Local.Get(ctx, types.NamespacedName{Name: name}, existing); err == nil {
		return deploy.Done()
	} else if !apierrors.IsNotFound(err) {
		return deploy.Error(err)
	}

	metaLabels := util.MergeMaps(t.Namer.Labels(), t.APPBox.Labels)
	metaLabels = util.MergeMaps(metaLabels, labels.NewOwnerLabels(t.NodePool))

	commonTags := maps.Clone(t.NodePool.Labels)
	commonTags["tidbcloud.com/nodegroup-name"] = name
	tags := controllers.ConvertToCloudTags(commonTags, consts.ProviderTypeEKS)
	tags = util.MergeMaps(tags, t.Config.AdditionalCloudTags)
	tags = util.BuildAWSTagWithName(tags, name)

	pg := &ec2v1beta1.PlacementGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: metaLabels,
		},
		Spec: ec2v1beta1.PlacementGroupSpec{
			ForProvider: ec2v1beta1.PlacementGroupParameters{
				Region:         ptr.To(t.APPBox.Spec.Region),
				Strategy:       ptr.To("partition"),
				PartitionCount: ptr.To(float64(placementGroupPartitionCount)),
				Tags:           util.ToMapPointer(tags),
			},
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{Name: t.ProviderConfigName()},
				ManagementPolicies:      controllers.GetResourceManagementPolicies(),
			},
		},
	}
	return deploy.ErrorOrDone(k8s.Apply(ctx, t.Local, pg))
}

func (t *placementGroup) PostCreate(ctx context.Context) deploy.Result {
	if !placementGroupEnabled(t.NodePool) {
		return deploy.Done()
	}

	pg := &ec2v1beta1.PlacementGroup{}
	if err := t.Local.Get(ctx, types.NamespacedName{Name: t.placementGroupName()}, pg); err != nil {
		return deploy.Error(err).WithRequeuAfter(deploy.DefaultRequeueAfter)
	}
	condition := pg.Status.GetCondition(xpv1.TypeReady)
	if condition.Status != corev1.ConditionTrue {
		t.NodePool.Status.Conditions = []metav1.Condition{task.ToMetaCondition(condition)}
		return deploy.Error(fmt.Errorf("waiting for placement group %s to be ready", pg.GetName())).WithRequeuAfter(deploy.DefaultRequeueAfter)
	}
	t.NodePool.Status.Conditions = []metav1.Condition{task.ToMetaCondition(condition)}
	return deploy.Done()
}

func (t *placementGroup) Delete(ctx context.Context) deploy.Result {
	if !placementGroupEnabled(t.NodePool) {
		return deploy.Done()
	}

	inUse, err := t.otherLivePlacementGroupUsers(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	if inUse > 0 {
		return deploy.Done()
	}

	pg := &ec2v1beta1.PlacementGroup{ObjectMeta: metav1.ObjectMeta{Name: t.placementGroupName()}}
	return deploy.ErrorOrDone(controllers.DeleteUpboundResource(ctx, t.Local, pg))
}

func (t *placementGroup) PostDelete(ctx context.Context) deploy.Result {
	if !placementGroupEnabled(t.NodePool) {
		return deploy.Done()
	}
	inUse, err := t.otherLivePlacementGroupUsers(ctx)
	if err != nil {
		return deploy.Error(err)
	}
	if inUse > 0 {
		return deploy.Done()
	}
	return deploy.ErrorOrDone(k8s.GetDeletedObject(ctx, t.Local, types.NamespacedName{Name: t.placementGroupName()}, &ec2v1beta1.PlacementGroup{}))
}

func (t *placementGroup) otherLivePlacementGroupUsers(ctx context.Context) (int, error) {
	nodePools := &v1alpha1.NodePoolList{}
	if err := t.Local.List(ctx, nodePools, client.InNamespace(t.NodePool.Namespace)); err != nil {
		return 0, err
	}

	matchesOwner := t.nodePoolOwnerMatcher()
	count := 0
	for _, np := range nodePools.Items {
		if np.Name == t.NodePool.Name {
			continue
		}
		if placementGroupEnabled(&np) && np.Spec.InstanceType == t.NodePool.Spec.InstanceType && np.DeletionTimestamp.IsZero() {
			// Only count nodepools that resolve to the same placement group owner identity.
			if matchesOwner(np) {
				count++
			}
		}
	}
	return count, nil
}

func (t *placementGroup) nodePoolOwnerMatcher() func(v1alpha1.NodePool) bool {
	if ownerLabels, ok := ownerLabelsFromNodePool(t.NodePool); ok {
		return func(np v1alpha1.NodePool) bool {
			return np.Labels[consts.LabelKeyOwnedBy] == ownerLabels[consts.LabelKeyOwnedBy] &&
				np.Labels[consts.LabelKeyNamespace] == ownerLabels[consts.LabelKeyNamespace] &&
				np.Labels[consts.LabelKeyName] == ownerLabels[consts.LabelKeyName] &&
				np.Labels[consts.LabelKeyUID] == ownerLabels[consts.LabelKeyUID]
		}
	}

	if t.APPBox != nil {
		ownerLabels := labels.NewOwnerLabels(t.APPBox)
		return func(np v1alpha1.NodePool) bool {
			return np.Labels[consts.LabelKeyOwnedBy] == ownerLabels[consts.LabelKeyOwnedBy] &&
				np.Labels[consts.LabelKeyNamespace] == ownerLabels[consts.LabelKeyNamespace] &&
				np.Labels[consts.LabelKeyName] == ownerLabels[consts.LabelKeyName] &&
				np.Labels[consts.LabelKeyUID] == ownerLabels[consts.LabelKeyUID]
		}
	}

	appBoxName := ""
	if t.NodePool.Labels != nil {
		appBoxName = t.NodePool.Labels[v1alpha1.APPBoxNameKey]
	}
	if appBoxName != "" {
		dataPlaneID := t.NodePool.Labels[v1alpha1.DataPlaneIDKey]
		return func(np v1alpha1.NodePool) bool {
			if np.Labels[v1alpha1.APPBoxNameKey] != appBoxName {
				return false
			}
			if dataPlaneID == "" {
				return true
			}
			return np.Labels[v1alpha1.DataPlaneIDKey] == dataPlaneID
		}
	}

	dataPlaneID := ""
	if t.NodePool.Labels != nil {
		dataPlaneID = t.NodePool.Labels[v1alpha1.DataPlaneIDKey]
	}
	if dataPlaneID != "" {
		return func(np v1alpha1.NodePool) bool {
			return np.Labels[v1alpha1.DataPlaneIDKey] == dataPlaneID
		}
	}

	return func(v1alpha1.NodePool) bool {
		return true
	}
}

func ownerLabelsFromNodePool(nodePool *v1alpha1.NodePool) (map[string]string, bool) {
	if nodePool == nil || nodePool.Labels == nil {
		return nil, false
	}

	keys := []string{
		consts.LabelKeyOwnedBy,
		consts.LabelKeyNamespace,
		consts.LabelKeyName,
		consts.LabelKeyUID,
	}
	ownerLabels := make(map[string]string, len(keys))
	for _, key := range keys {
		value := nodePool.Labels[key]
		if value == "" {
			return nil, false
		}
		ownerLabels[key] = value
	}
	return ownerLabels, true
}

func (t *placementGroup) placementGroupName() string {
	return namer.PlacementGroupName(t.NodePool.Spec.InstanceType, t.placementGroupIdentity())
}

func (t *placementGroup) placementGroupIdentity() string {
	if t.APPBox != nil {
		return appBoxPlacementGroupIdentity(t.APPBox)
	}
	if t.NodePool != nil && t.NodePool.Labels != nil {
		if appBoxName := t.NodePool.Labels[v1alpha1.APPBoxNameKey]; appBoxName != "" {
			return t.NodePool.Namespace + "." + appBoxName
		}
		if dataPlaneID := t.NodePool.Labels[v1alpha1.DataPlaneIDKey]; dataPlaneID != "" {
			return t.NodePool.Namespace + "." + dataPlaneID
		}
	}
	return t.NodePool.Namespace + "." + t.NodePool.Name
}

func appBoxPlacementGroupIdentity(appBox *v1alpha1.APPBox) string {
	return appBox.Namespace + "." + appBox.Name
}