package common

import (
	"context"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
	"github.com/aws/smithy-go"
	"github.com/go-logr/logr"
	corev1 "k8s.io/api/core/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	awsclient "github.com/tidbcloud/infra-provider/pkg/cloud/aws"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/dataplane"
	"github.com/tidbcloud/infra-provider/pkg/metrics"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

const (
	// run every 10 minutes by default
	scaleProtectionInterval = 60 * time.Minute * 4

	scaleInProtectionAnnotationKey = "tidbcloud.com/scale-in-protection"
	scaleInProtectionEnabledValue  = "enabled"

	managedNodegroupKey = "eks.amazonaws.com/nodegroup"
	tidbclusterName     = "tidb-cluster"
)

// ScaleProtectionChecker is a background runnable that periodically checks
// whether data-plane nodes in AWS EKS self-managed nodegroups have
// scale-in protection enabled.
//
// It is intentionally decoupled from the cluster reconcile loop.
type ScaleProtectionChecker struct {
	InfraClient        client.Client
	EnvProviderManager dataplane.EnvProviderManager
	Logger             logr.Logger
	Interval           time.Duration

	AWSCli               awsclient.Client
	AssumeRoleName       string
	AssumeRoleExternalID string
}

// Start implements manager.Runnable. It runs until ctx is cancelled.
func (c *ScaleProtectionChecker) Start(ctx context.Context) error {
	interval := c.Interval
	if interval <= 0 {
		interval = scaleProtectionInterval
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	// run immediately once
	c.safeRunOnce(ctx)

	for {
		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
			c.safeRunOnce(ctx)
		}
	}
}

// safeRunOnce wraps runOnce with panic recovery so that a bug in the checker
// does not crash the whole controller process.
func (c *ScaleProtectionChecker) safeRunOnce(ctx context.Context) {
	log := c.Logger
	defer func() {
		if r := recover(); r != nil {
			// record a controller panic metric for observability
			metrics.ControllerPanic.WithLabelValues("scale_protection_checker", "").Inc()
			// log the panic so we can inspect the stack trace in logs
			log.Error(fmt.Errorf("panic in scale protection checker: %v", r), "recovered from panic in scale protection checker")
		}
	}()
	c.runOnce(ctx)
}

func (c *ScaleProtectionChecker) isSelfManagedEKS(cluster *cloudv1.Cluster) bool {
	if cluster == nil {
		return false
	}
	if cluster.Spec.Kubernetes.Provider != cloudv1.ProviderTypeEKS {
		return false
	}
	if cluster.Annotations == nil {
		return false
	}
	return cluster.Annotations[consts.UseSelfManagedNodegroupAnnotation] == consts.AnnoValTrue
}

func (c *ScaleProtectionChecker) runOnce(ctx context.Context) {
	log := c.Logger

	var cl cloudv1.ClusterList
	if err := c.InfraClient.List(ctx, &cl); err != nil {
		log.Error(err, "failed to list clusters for scale protection check")
		return
	}

	for i := range cl.Items {
		cluster := &cl.Items[i]
		if !cluster.DeletionTimestamp.IsZero() || cluster.Suspending() || cluster.Suspended() {
			continue
		}
		if !c.isSelfManagedEKS(cluster) {
			continue
		}

		c.checkClusterScaleProtection(ctx, log, cluster)
	}
}

type nodeInfo struct {
	node          *corev1.Node
	componentType cloudv1.ComponentType
	instanceID    string
}

func (c *ScaleProtectionChecker) checkClusterScaleProtection(ctx context.Context, log logr.Logger, cluster *cloudv1.Cluster) {
	asgCli, kcli, region, err := c.buildClusterClients(ctx, log, cluster)
	if err != nil {
		return
	}

	nodeList, podsByNode, err := listClusterNodesAndPods(ctx, log, kcli, cluster)
	if err != nil {
		return
	}

	candidates, instanceIDs := collectScaleProtectionCandidates(log, cluster, nodeList, podsByNode)

	// if no candidates, there is nothing to check this round
	labels := []string{
		cluster.GetProviderType(),
		cluster.GetRegion(),
		cluster.Name,
		cluster.Namespace,
		cluster.GetTenantID(),
		cluster.GetProjectID(),
		string(cluster.Status.Phase),
	}

	if len(candidates) == 0 {
		return
	}

	// query AWS for ProtectedFromScaleIn state
	protectedFromScaleIn, ok := describeProtectedFromScaleInWithRetry(
		ctx, log, asgCli, cluster.Name, region, instanceIDs,
	)
	if !ok {
		// aws describe failed permanently this round; treat as skipped rather than
		// assuming all nodes are unprotected
		recordScaleProtectionCheckSkipped(cluster, "aws_describe_failed", "")
		return
	}

	for _, info := range candidates {
		awsProtected := false
		if protectedFromScaleIn != nil {
			if v, ok := protectedFromScaleIn[info.instanceID]; ok {
				awsProtected = v
			}
		}

		annoVal := ""
		if info.node.Annotations != nil {
			annoVal = info.node.Annotations[scaleInProtectionAnnotationKey]
		}
		annoProtected := annoVal == scaleInProtectionEnabledValue

		// node is missing protection if either annotation or AWS side is not enabled
		isMissing := !annoProtected || !awsProtected
		if isMissing {
			log.Info("node missing scale-in protection according to annotation and/or aws state",
				"cluster", cluster.Name,
				"node", info.node.Name,
				"component", string(info.componentType),
				"instanceID", info.instanceID,
				"annotationProtected", annoProtected,
				"annotationValue", annoVal,
			)
		}

		recordScaleProtectionNodeMetric(labels, info, annoProtected, awsProtected, isMissing)
	}
}

func recordScaleProtectionNodeMetric(
	labels []string,
	info nodeInfo,
	annoProtected, awsProtected, isMissing bool,
) {
	value := float64(0)
	if isMissing {
		value = 1
	}

	nodeLabels := append(append([]string{}, labels...),
		info.node.Name,
		info.instanceID,
		string(info.componentType),
		strconv.FormatBool(!annoProtected),
		strconv.FormatBool(!awsProtected),
	)

	metrics.ScaleProtectionMissingNode.WithLabelValues(nodeLabels...).Set(value)
}

func recordScaleProtectionCheckSkipped(cluster *cloudv1.Cluster, reason, node string) {
	labels := []string{
		cluster.GetProviderType(),
		cluster.GetRegion(),
		cluster.Name,
		cluster.Namespace,
		cluster.GetTenantID(),
		cluster.GetProjectID(),
		string(cluster.Status.Phase),
		reason,
		node,
	}
	metrics.ScaleProtectionCheckSkipped.WithLabelValues(labels...).Inc()
}

func (c *ScaleProtectionChecker) buildClusterClients(
	ctx context.Context,
	log logr.Logger,
	cluster *cloudv1.Cluster,
) (*autoscaling.Client, client.Client, string, error) {
	ep, err := c.EnvProviderManager.GetEnvProvider(ctx, cluster)
	if err != nil {
		log.Error(err, "failed to get env provider for cluster", "cluster", cluster.Name)
		return nil, nil, "", err
	}

	region := ep.GetRegion()
	account := ep.GetAccount()

	var (
		asgCli *autoscaling.Client
	)
	if c.AWSCli != nil && c.AssumeRoleName != "" {
		roleARN := util.GetRoleARN(account, c.AssumeRoleName)
		assumedRole := &awsclient.AssumedRole{
			RoleARN:    roleARN,
			ExternalID: nil,
		}
		if c.AssumeRoleExternalID != "" {
			assumedRole.ExternalID = &c.AssumeRoleExternalID
		}
		asgCli = c.AWSCli.AUTOSCALE(assumedRole, awsclient.WithRegion(region))
	}

	kcli := ep.GetClient()
	if kcli == nil {
		log.Info("env provider client is nil, skip scale protection check for cluster", "cluster", cluster.Name)
		return nil, nil, region, errors.New("env provider client is nil")
	}

	return asgCli, kcli, region, nil
}

func listClusterNodesAndPods(
	ctx context.Context,
	log logr.Logger,
	kcli client.Client,
	cluster *cloudv1.Cluster,
) (*corev1.NodeList, map[string][]*corev1.Pod, error) {
	// step 1: list all TiDB component pods in the cluster namespace, then group by node
	podsByNode := make(map[string][]*corev1.Pod)
	podList := &corev1.PodList{}
	if err := kcli.List(ctx, podList, client.InNamespace(cluster.PrefixedId())); err != nil {
		log.Error(err, "failed to list pods for scale protection check", "cluster", cluster.Name)
		return nil, nil, err
	}

	if len(podList.Items) == 0 {
		log.Info("no pods found for scale protection check", "cluster", cluster.Name, "namespace", cluster.PrefixedId())
		return &corev1.NodeList{}, podsByNode, nil
	}

	for i := range podList.Items {
		pod := &podList.Items[i]
		// only consider pods that belong to the tidb cluster and have a component label
		if _, ok := getPodComponentType(pod); !ok {
			continue
		}
		if pod.Spec.NodeName == "" {
			continue
		}
		podsByNode[pod.Spec.NodeName] = append(podsByNode[pod.Spec.NodeName], pod)
	}

	// step 2: list all nodes and keep only those that have at least one TiDB pod.
	// We intentionally list full Node objects here (instead of only relying on pod.Spec.NodeName)
	// so that later we can read labels/annotations/providerID and avoid per-node Get calls.
	// Filtering over a single List also ensures we only process nodes that currently exist.
	allNodes := &corev1.NodeList{}
	if err := kcli.List(ctx, allNodes); err != nil {
		log.Error(err, "failed to list nodes for scale protection check", "cluster", cluster.Name)
		return nil, nil, err
	}

	nodeList := &corev1.NodeList{}
	for i := range allNodes.Items {
		n := &allNodes.Items[i]
		if _, ok := podsByNode[n.Name]; ok {
			nodeList.Items = append(nodeList.Items, *n)
		}
	}
	return nodeList, podsByNode, nil
}

func collectScaleProtectionCandidates(
	log logr.Logger,
	cluster *cloudv1.Cluster,
	nodeList *corev1.NodeList,
	podsByNode map[string][]*corev1.Pod,
) ([]nodeInfo, []string) {
	var (
		candidates []nodeInfo
	)
	instanceIDs := make([]string, 0, len(nodeList.Items))

	for _, node := range nodeList.Items {
		labels := node.Labels
		if labels == nil {
			continue
		}

		// find pods scheduled on this node and check whether any of them
		// is a TiDB component pod
		podsOnNode := podsByNode[node.Name]
		if len(podsOnNode) == 0 {
			continue
		}

		// podsByNode is already built only from TiDB component pods, so we can
		// safely take the component type from any pod on this node. Keep a small
		// defensive check in case the construction logic changes in the future.
		ct, ok := getPodComponentType(podsOnNode[0])
		if !ok {
			continue
		}
		comp := *ct

		// skip managed nodegroups; self-managed nodes do not have this label
		if _, ok := labels[managedNodegroupKey]; ok {
			continue
		}

		// parse instance ID from providerID if possible
		providerID := node.Spec.ProviderID
		if providerID == "" {
			log.Info("skip node without providerID when collecting aws state",
				"cluster", cluster.Name,
				"node", node.Name,
				"component", string(comp))
			recordScaleProtectionCheckSkipped(cluster, "missing_provider_id", node.Name)
			continue
		}
		idx := strings.LastIndex(providerID, "/")
		if idx == -1 || idx+1 >= len(providerID) {
			log.Info("unexpected providerID format, skip aws check for node",
				"cluster", cluster.Name,
				"node", node.Name,
				"providerID", providerID)
			recordScaleProtectionCheckSkipped(cluster, "invalid_provider_id", node.Name)
			continue
		}
		instanceID := providerID[idx+1:]

		info := nodeInfo{
			node:          node.DeepCopy(),
			componentType: comp,
			instanceID:    instanceID,
		}
		candidates = append(candidates, info)
		instanceIDs = append(instanceIDs, instanceID)

		val := ""
		if node.Annotations != nil {
			val = node.Annotations[scaleInProtectionAnnotationKey]
		}
		if val != scaleInProtectionEnabledValue {
			// this log can be noisy on large clusters; keep at V(1)
			log.Info("node annotation indicates missing scale-in protection in AWS self-managed nodegroup for TiDB component node",
				"cluster", cluster.Name,
				"node", node.Name,
				"component", string(comp),
				"annotationValue", val,
			)
		}
	}

	return candidates, instanceIDs
}

func getPodComponentType(pod *corev1.Pod) (*cloudv1.ComponentType, bool) {
	var component cloudv1.ComponentType

	name := pod.Labels[consts.NameLabelKey]
	if name != tidbclusterName {
		return nil, false
	}

	componentName, ok := pod.Labels[consts.ComponentLabelKey]
	if !ok {
		return nil, false
	}

	component = cloudv1.ComponentType(componentName)
	return &component, true
}

// describeProtectedFromScaleInWithRetry calls DescribeAutoScalingInstances with
// small batches and an exponential backoff retry on throttling errors.
// It returns (result, true) on success, and (nil, false) on persistent failures
// or context cancellation.
func describeProtectedFromScaleInWithRetry(
	ctx context.Context,
	log logr.Logger,
	asgCli *autoscaling.Client,
	clusterName, region string,
	instanceIDs []string,
) (map[string]bool, bool) {
	protectedFromScaleIn := make(map[string]bool)

	if asgCli == nil {
		log.Info("aws autoscaling client is nil, skip aws state check for cluster", "cluster", clusterName)
		return nil, false
	}

	const (
		batchSize    = 20
		maxAttempts  = 10
		initialDelay = 200 * time.Millisecond
	)

	for i := 0; i < len(instanceIDs); i += batchSize {
		end := i + batchSize
		if end > len(instanceIDs) {
			end = len(instanceIDs)
		}
		batch := instanceIDs[i:end]

		var (
			out *autoscaling.DescribeAutoScalingInstancesOutput
			err error
		)

		delay := initialDelay
		for attempt := 1; attempt <= maxAttempts; attempt++ {
			out, err = asgCli.DescribeAutoScalingInstances(ctx, &autoscaling.DescribeAutoScalingInstancesInput{
				InstanceIds: batch,
			})
			if err == nil {
				break
			}

			// if it's not throttling, or we've exhausted attempts, fail fast
			if !isAutoScalingThrottlingError(err) || attempt == maxAttempts {
				log.Error(err, "failed to describe auto scaling instances from aws",
					"cluster", clusterName,
					"region", region,
					"attempt", attempt,
					"maxAttempts", maxAttempts)
				return nil, false
			}

			log.Info("throttled by aws autoscaling, retrying with backoff",
				"cluster", clusterName,
				"region", region,
				"attempt", attempt,
				"maxAttempts", maxAttempts)

			select {
			case <-ctx.Done():
				return nil, false
			case <-time.After(delay):
			}
			delay *= 2
		}

		for _, inst := range out.AutoScalingInstances {
			if inst.InstanceId == nil {
				continue
			}
			id := *inst.InstanceId
			if inst.ProtectedFromScaleIn != nil && *inst.ProtectedFromScaleIn {
				protectedFromScaleIn[id] = true
			} else {
				protectedFromScaleIn[id] = false
			}
		}
	}

	return protectedFromScaleIn, true
}

func isAutoScalingThrottlingError(err error) bool {
	var apiErr smithy.APIError
	if !errors.As(err, &apiErr) {
		return false
	}

	switch apiErr.ErrorCode() {
	case "Throttling", "ThrottlingException", "RequestLimitExceeded":
		return true
	default:
		return false
	}
}