package commonaws

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	eksv1alpha1 "github.com/crossplane-contrib/provider-aws/apis/eks/manualv1alpha1"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	rsv1alpha1 "github.com/tidbcloud/observability/agent-control-server/api/regionalserver/v1alpha1"
	"github.com/tidbcloud/observability/agent-control-server/pkg/agentbuild"
	"github.com/tidbcloud/observability/agent-control-server/pkg/config"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
	"github.com/tidbcloud/observability/agent-control-server/pkg/deploy"
	"github.com/tidbcloud/observability/agent-control-server/pkg/reconciler"
	"github.com/tidbcloud/observability/agent-control-server/pkg/util"
)

const (
	defaultMinSize        = 1
	defaultMaxSize        = 3
	defaultDesiredSize    = 1
	defaultDiskSize       = 20 // in GB
	defaultVersion        = "1.32"
	defaultReleaseVersion = "1.32.8-20250915"
	defaultAMIType        = "AL2023_x86_64_STANDARD"
	defaultCapacityType   = "ON_DEMAND"
	defaultSpecName       = "default"
)

// AL2AMIType is the deprecated AL2_x86_64 AMI type. EKS support ends at k8s
// 1.32; NodeGroups on this AMI must be migrated to BottlerocketAMIType when
// upgrading to 1.33+.
const AL2AMIType = "AL2_x86_64"

// BottlerocketAMIType is the AMI type that legacy AL2_x86_64 NodeGroups are
// migrated to. It doubles as the name-suffix trigger: NodeGroups on this AMI
// are named with a "-bottlerocket" suffix to enable parallel migration from
// legacy AL2 NGs (which use the unsuffixed name).
const BottlerocketAMIType = "BOTTLEROCKET_x86_64"

// BottlerocketSuffix is the NodeGroup name suffix appended to migrated
// BOTTLEROCKET NodeGroups. Exposed so HTTP-layer code can spell the same
// candidate name without re-deriving the suffix-decision logic.
const BottlerocketSuffix = "-bottlerocket"

// MaxAL2Version is the last k8s version where AL2_x86_64 is supported by
// EKS. AL2 NodeGroups are migrated to BOTTLEROCKET once the Agent moves
// past this version, so upgrades at or below 1.32 stay on AL2.
const MaxAL2Version = "1.32"

// ShouldMigrateAMI reports whether an Agent on the given (annotation) AMIType
// and k8s version should be migrated to BOTTLEROCKET_x86_64. Migration is
// gated on BOTH:
//   - AMIType is currently AL2_x86_64 (the only deprecated type), and
//   - k8s version is above MaxAL2Version (the last version EKS AL2 supports).
//
// Empty or unparseable versions return false so an Agent with unknown version
// never silently triggers a migration.
func ShouldMigrateAMI(amiType, version string) bool {
	if amiType != AL2AMIType {
		return false
	}
	if version == "" {
		return false
	}
	ok, err := util.CompareVersions(version, ">", MaxAL2Version)
	if err != nil {
		return false
	}
	return ok
}

// NameSuffix returns the NodeGroup name suffix corresponding to the given
// (annotation) AMIType, suffix-required flag, and k8s version.
//
// The suffix is added IFF the Agent has been (or is being) migrated from
// legacy AL2 to BOTTLEROCKET:
//   - AMIType is BOTTLEROCKET_x86_64 AND suffixRequired == "true": sticky
//     signal set permanently when migration happened, so the suffixed NG
//     name persists across reconciles even though AMIType now reads
//     BOTTLEROCKET; OR
//   - ShouldMigrateAMI(amiType, version): migration is desired on this
//     reconcile but the flag has not been persisted yet (first cycle).
//
// Native BOTTLEROCKET Agents (bootstrapped from a BOTTLEROCKET admin NG
// without ever being AL2) and AL2 Agents on k8s < 1.33 get no suffix.
func NameSuffix(amiType, suffixRequired, version string) string {
	if amiType == BottlerocketAMIType && suffixRequired == "true" {
		return BottlerocketSuffix
	}
	if ShouldMigrateAMI(amiType, version) {
		return BottlerocketSuffix
	}
	return ""
}

// MigrateAMIType returns the desired AMI type for the given current AMIType
// and k8s version. The deprecated AL2_x86_64 AMI is migrated to
// BOTTLEROCKET_x86_64 once the Agent moves past MaxAL2Version.
// All other AMI types — including AL2 on k8s <= 1.32, AL2_ARM_64, AL2023,
// and existing BOTTLEROCKET variants — are returned unchanged.
func MigrateAMIType(currentAMIType, version string) string {
	if ShouldMigrateAMI(currentAMIType, version) {
		return BottlerocketAMIType
	}
	return currentAMIType
}

// NodeGroupTask manages the EKS NodeGroup for the observability agent.
//
// NodeGroupTask creates, updates, and deletes the NodeGroup based on the desired specifications.
// It also handles scaling operations and ensures the NodeGroup is in the desired state.
// The task uses annotations on the agent resource to track the desired and previous specifications.
type NodeGroupTask[T client.Object] struct {
	*reconciler.BaseTask[T]
	clusterName     string
	nameFunc        func(string) string
	specs           config.CustomerAgentSpecs
	taints          []eksv1alpha1.Taint
	nodeMatchLabels map[string]string
	vpcMatchLabels  map[string]string
}

func NewNodeGroupTask[T client.Object](
	ctx *reconciler.Context[T],
	clusterName string,
	nameFunc func(string) string,
	nodegroupLabels map[string]string,
	vpcMatchLabels map[string]string,
	specs config.CustomerAgentSpecs,
	taints []eksv1alpha1.Taint,
) *NodeGroupTask[T] {
	return &NodeGroupTask[T]{
		BaseTask:        &reconciler.BaseTask[T]{Context: ctx},
		clusterName:     clusterName,
		nameFunc:        nameFunc,
		nodeMatchLabels: nodegroupLabels,
		vpcMatchLabels:  vpcMatchLabels,
		specs:           specs,
		taints:          taints,
	}
}

func (t *NodeGroupTask[T]) Name() string {
	return "NodeGroup"
}

func (t *NodeGroupTask[T]) Create(ctx context.Context) deploy.Result {
	return deploy.PanicOrDone(t.create(ctx))
}

func (t *NodeGroupTask[T]) PostCreate(ctx context.Context) deploy.Result {
	desiredSpec := util.GetAnnotationWithDefault(t.Agent, consts.AnnotationKeyO11yDesiredSpec, defaultSpecName)
	currentName := t.nameFunc(desiredSpec)
	if err := agentbuild.CheckNodeGroupReady(ctx, t.Local, currentName); err != nil {
		t.clearStalledExternalCreate(ctx, currentName)
		return deploy.Panic(err)
	}
	// Skip stale-NG cleanup during SpecScale transitions to preserve the previousSpec NG.
	if status := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus); status == "" || status == consts.SpecScaleStatusNormal {
		if err := t.cleanupStaleNodeGroups(ctx, currentName); err != nil {
			return deploy.Panic(err)
		}
	}
	switch util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus) {
	case consts.SpecScaleStatusPendingScale:
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus, consts.SpecScaleStatusNewNodeGroupSynced)
	}
	return deploy.Done()
}

// cleanupStaleNodeGroups deletes NodeGroups matching t.nodeMatchLabels except
// excludeName. Used to garbage-collect legacy AL2 NodeGroups once the
// BOTTLEROCKET replacement (named excludeName) is Active.
func (t *NodeGroupTask[T]) cleanupStaleNodeGroups(ctx context.Context, excludeName string) error {
	deleted, err := agentbuild.DeleteNodeGroupWithLabels(ctx, t.Local, t.nodeMatchLabels, excludeName)
	if err != nil {
		return err
	}
	if len(deleted) > 0 {
		t.L.Info("deleted stale nodegroups", "nodegroups", deleted)
	}
	return nil
}

// clearStalledExternalCreate best-effort clears crossplane's stalled
// external-create-pending annotation on the nodegroup so crossplane resumes
// reconciling it. See agentbuild.ClearStalledExternalCreate for why this is safe.
func (t *NodeGroupTask[T]) clearStalledExternalCreate(ctx context.Context, name string) {
	cleared, err := agentbuild.ClearStalledExternalCreate(ctx, t.Local,
		&eksv1alpha1.NodeGroup{ObjectMeta: metav1.ObjectMeta{Name: name}})
	switch {
	case err != nil:
		t.L.Error(err, "failed to clear stalled crossplane external-create-pending annotation", "name", name)
	case cleared:
		t.L.Info("cleared stalled crossplane external-create-pending annotation for crossplane to resume", "name", name)
	}
}

func (t *NodeGroupTask[T]) Delete(ctx context.Context) deploy.Result {
	// capable with legacy
	if err := agentbuild.DeleteNodeGroup(ctx, t.Local, t.nameFunc(defaultSpecName)); err != nil {
		return deploy.Panic(err)
	}
	deleted, err := agentbuild.DeleteNodeGroupWithLabels(ctx, t.Local, t.nodeMatchLabels)
	if err != nil {
		return deploy.Panic(err)
	}
	if len(deleted) > 0 {
		t.L.Info("deleted nodegroups", "nodegroups", deleted)
	}
	return deploy.Done()
}

func (t *NodeGroupTask[T]) PostDelete(ctx context.Context) deploy.Result {
	// capable with legacy
	name := t.nameFunc(defaultSpecName)
	if err := agentbuild.CheckNodeGroupDeleted(ctx, t.Local, name); err != nil {
		t.clearStalledExternalCreate(ctx, name)
		return deploy.Panic(err)
	}
	return deploy.PanicOrDone(agentbuild.CheckNodeGroupDeletedWithLabels(ctx, t.Local, t.nodeMatchLabels))
}

func (t *NodeGroupTask[T]) create(ctx context.Context) error {
	desiredSpec := util.GetAnnotationWithDefault(t.Agent, consts.AnnotationKeyO11yDesiredSpec, defaultSpecName)
	spec := t.specs.Select(desiredSpec)
	if spec == nil {
		return fmt.Errorf("desired spec not found: %s", desiredSpec)
	}
	if err := t.reconcile(ctx, t.nameFunc(desiredSpec), spec.InstanceType); err != nil {
		return err
	}

	switch util.GetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus) {
	case consts.SpecScaleStatusPendingScale:
		// If we are in a pending-scale state, remember to reconcile the original spec
		// until the new spec is ready. This is to avoid a situation where the old spec
		// cannot be further edited when the new spec is delayed.
		previousSpec := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yPreviousSpec)
		if previousSpec == "" {
			return fmt.Errorf("unexpected empty previous spec name")
		}
		spec := t.specs.Select(previousSpec)
		if spec == nil {
			return fmt.Errorf("previous spec not found: %s", previousSpec)
		}
		if err := t.reconcile(ctx, t.nameFunc(previousSpec), spec.InstanceType); err != nil {
			return err
		}

	case consts.SpecScaleStatusResourceUpdated:
		previousSpec := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yPreviousSpec)
		if previousSpec == "" {
			return fmt.Errorf("unexpected empty previous spec name")
		}
		prevName := t.nameFunc(previousSpec)
		if err := agentbuild.DeleteNodeGroup(ctx, t.Local, prevName); err != nil {
			return fmt.Errorf("failed to delete previous nodegroup: %w", err)
		}
		if err := agentbuild.CheckNodeGroupDeleted(ctx, t.Local, prevName); err != nil {
			t.clearStalledExternalCreate(ctx, prevName)
			return err
		}
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11ySpecScaleStatus, consts.SpecScaleStatusNormal)
	}
	return nil
}

func (t *NodeGroupTask[T]) updateSizeAnnotation(key string, val int) (int, error) {
	s := util.SetAnnotationIfEmpty(t.Agent, key, strconv.Itoa(val))
	i, err := strconv.Atoi(s)
	if err != nil {
		return 0, fmt.Errorf("invalid key '%s': %w", key, err)
	}
	return i, nil
}

// reconcile ensures that the NodeGroup with the given name and instance type exists
// and is configured according to the specifications derived from the agent annotations
// and the admin NodeGroup if necessary.
// It updates the NodeGroup resource in the Kubernetes cluster accordingly.
func (t *NodeGroupTask[T]) reconcile(ctx context.Context, name, instanceType string) error {
	nodeRole := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeRole)
	providerConfigReference := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeProviderConfigReference)
	amiType := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeAMIType)
	version := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeVersion)
	releaseVersion := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeReleaseVersion)
	if nodeRole == "" || providerConfigReference == "" || amiType == "" || version == "" || releaseVersion == "" {
		var nodegroups eksv1alpha1.NodeGroupList
		if err := t.Local.List(ctx, &nodegroups, client.MatchingFields{
			".spec.forProvider.clusterName": t.clusterName,
		}); err != nil {
			return fmt.Errorf("failed to list nodegroup: %w", err)
		}
		var ng *eksv1alpha1.NodeGroup
		for _, item := range nodegroups.Items {
			if strings.HasPrefix(item.Name, "tidbcloud-dedicated-") || strings.HasPrefix(item.Name, "tidbcloud-prod-") {
				ng = &item
				break
			}
		}
		if ng == nil {
			return fmt.Errorf("admin nodegroup not found")
		}

		nodeRole = ng.Spec.ForProvider.NodeRole
		providerConfigReference = ng.Spec.ProviderConfigReference.Name
		if ng.Spec.ForProvider.Version != nil {
			version = *ng.Spec.ForProvider.Version
		} else {
			version = defaultVersion
		}
		if ng.Spec.ForProvider.ReleaseVersion != nil {
			releaseVersion = *ng.Spec.ForProvider.ReleaseVersion
		} else {
			releaseVersion = defaultReleaseVersion
		}
		if ng.Spec.ForProvider.AMIType != nil {
			amiType = *ng.Spec.ForProvider.AMIType
		} else {
			amiType = defaultAMIType
		}
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeRole, nodeRole)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeProviderConfigReference, providerConfigReference)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeAMIType, amiType)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeVersion, version)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeReleaseVersion, releaseVersion)
	}

	desiredAMIType := MigrateAMIType(amiType, version)
	if desiredAMIType != amiType {
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeAMIType, desiredAMIType)
		// Persist the suffix past this cycle: once AMIType flips to
		// BOTTLEROCKET the AL2 trigger is gone, so without this flag the
		// suffixed NG would be treated as stale and cleaned up.
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeAMISuffixRequired, "true")
	}

	vpcName := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeVPCName)
	zone := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeZone)
	subnet := util.GetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeSubnet)
	if vpcName == "" || zone == "" || subnet == "" {
		var vpcs rsv1alpha1.VPCList
		if err := t.Local.List(ctx, &vpcs, client.MatchingLabels(t.vpcMatchLabels)); err != nil {
			return fmt.Errorf("failed to list vpc: %w", err)
		}
		if len(vpcs.Items) == 0 {
			return fmt.Errorf("no vpc found for %v", t.vpcMatchLabels)
		}
		vpc := vpcs.Items[0]
		if len(vpc.Status.NodeSubnets) == 0 {
			return fmt.Errorf("vpc %s has no node subnets", vpc.Name)
		}

		vpcName = vpc.Name
		zone = vpc.Status.NodeSubnets[0].AvailabilityZone
		subnet = vpc.Status.NodeSubnets[0].SubnetID
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeVPCName, vpcName)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeZone, zone)
		util.SetAnnotation(t.Agent, consts.AnnotationKeyO11yNodeSubnet, subnet)
	}

	minSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeMinSize, defaultMinSize)
	if err != nil {
		return err
	}
	maxSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeMaxSize, defaultMaxSize)
	if err != nil {
		return err
	}
	desiredSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeDesiredSize, defaultDesiredSize)
	if err != nil {
		return err
	}
	diskSize, err := t.updateSizeAnnotation(consts.AnnotationKeyO11yNodeDiskSize, defaultDiskSize)
	if err != nil {
		return err
	}

	capacityType := util.SetAnnotationIfEmpty(t.Agent, consts.AnnotationKeyO11yNodeCapacityType, defaultCapacityType)
	nodegroup := &eksv1alpha1.NodeGroup{
		ObjectMeta: metav1.ObjectMeta{Name: name},
	}
	_, err = controllerutil.CreateOrUpdate(ctx, t.Local, nodegroup, func() error {
		if nodegroup.Labels == nil {
			nodegroup.Labels = make(map[string]string)
		}
		for k, v := range t.nodeMatchLabels {
			nodegroup.Labels[k] = v
		}

		nodegroup.Spec = eksv1alpha1.NodeGroupSpec{
			ResourceSpec: xpv1.ResourceSpec{
				DeletionPolicy:     xpv1.DeletionDelete,
				ManagementPolicies: xpv1.ManagementPolicies{xpv1.ManagementActionAll},
				ProviderConfigReference: &xpv1.Reference{
					Name: providerConfigReference,
				},
			},
			ForProvider: eksv1alpha1.NodeGroupParameters{
				AMIType:       &desiredAMIType,
				CapacityType:  &capacityType,
				ClusterName:   t.clusterName,
				DiskSize:      ptr.To(int32(diskSize)),
				InstanceTypes: []string{instanceType},
				Labels:        t.nodeMatchLabels,
				NodeRole:      nodeRole,
				Region:        t.Config.O11yConfig.Region,
				ScalingConfig: &eksv1alpha1.NodeGroupScalingConfig{
					DesiredSize: ptr.To(int32(desiredSize)),
					MaxSize:     ptr.To(int32(maxSize)),
					MinSize:     ptr.To(int32(minSize)),
				},
				Subnets: []string{subnet},
				Tags: map[string]string{
					"Name":                         nodegroup.Name,
					"tidbcloud.com/nodegroup-name": nodegroup.Name,
				},
				Taints: t.taints,
				UpdateConfig: &eksv1alpha1.NodeGroupUpdateConfig{
					MaxUnavailable: ptr.To[int32](1),
				},
				Version: &version,
			},
		}
		return nil
	})
	return err
}