package az

import (
	"context"
	"fmt"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
	autoscalingv1beta1 "github.com/crossplane-contrib/provider-aws/apis/autoscaling/v1beta1"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	eksv1alpha1 "github.com/tidbcloud/eks-provider/api/v1alpha1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	cloudaws "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/namer"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

const crossplanePausedAnnotation = "crossplane.io/paused"

var asgSuspendProcesses = []string{
	"Launch",
	"Terminate",
	"ReplaceUnhealthy",
	"HealthCheck",
}

type asgRoleConfig struct {
	assumeRoleName string
	externalID     string
}

func (c asgRoleConfig) assumedRole(accountID string) *cloudaws.AssumedRole {
	ar := &cloudaws.AssumedRole{
		RoleARN: util.GetRoleARN(accountID, c.assumeRoleName),
	}
	if c.externalID != "" {
		ar.ExternalID = &c.externalID
	}
	return ar
}

func resolveASGAssumedRole(ctx context.Context, cli client.Client, cluster *cloudv1.Cluster, envProvider dataplane.EnvProvider, cfg asgRoleConfig) (*cloudaws.AssumedRole, error) {
	var network cloudv1.Network
	if err := cli.Get(ctx, cluster.GetNetworkRef(), &network); err == nil {
		if spec, err := eksv1alpha1.EKSNetworkSpecForNetwork(&network); err == nil && spec != nil && spec.AssumeRoleArn != "" {
			return cloudaws.GetRoleFromEKSNetworkSpec(spec), nil
		}
	}
	accountID := envProvider.GetAccount()
	if accountID == "" {
		return nil, fmt.Errorf("empty AWS account id from env provider")
	}
	if cfg.assumeRoleName == "" {
		return nil, fmt.Errorf("assume role name is not configured")
	}
	return cfg.assumedRole(accountID), nil
}

func listGroupASGs(ctx context.Context, infraCli client.Client, group Group, zone string) ([]*autoscalingv1beta1.AutoScalingGroup, error) {
	var list autoscalingv1beta1.AutoScalingGroupList
	if err := infraCli.List(ctx, &list, group.Namer().ClusterSelector()); err != nil {
		return nil, fmt.Errorf("list ASGs for group %s: %w", group.ID(), err)
	}
	out := make([]*autoscalingv1beta1.AutoScalingGroup, 0)
	for idx := range list.Items {
		asg := &list.Items[idx]
		if zone != "" && asg.Labels[consts.NodeGroupZonesLabelKey] != zone {
			continue
		}
		comp := asg.Labels[namer.ComponentSelectorKey]
		if !componentLabelAllowed(comp, group) {
			continue
		}
		out = append(out, asg)
	}
	return out, nil
}

func isAWSProvider(cluster *cloudv1.Cluster) bool {
	switch cluster.GetProviderType() {
	case cloudv1.ProviderTypeEKS, cloudv1.ProviderTypeAWS:
		return true
	default:
		return false
	}
}

// pauseASGs freezes Crossplane reconciliation and AWS scaling processes for
// ASGs belonging to the given group in the fault AZ. Idempotent.
func pauseASGs(
	ctx context.Context,
	infraCli client.Client,
	awsCli cloudaws.Client,
	cluster *cloudv1.Cluster,
	envProvider dataplane.EnvProvider,
	roleCfg asgRoleConfig,
	group Group,
	zone string,
) ([]string, error) {
	logger := log.FromContext(ctx).WithValues("groupId", group.ID())
	if awsCli == nil || !isAWSProvider(cluster) {
		logger.Info("skip ASG pause: not an AWS cluster or AWS client unavailable")
		return nil, nil
	}
	asgs, err := listGroupASGs(ctx, infraCli, group, zone)
	if err != nil {
		return nil, err
	}
	role, err := resolveASGAssumedRole(ctx, infraCli, cluster, envProvider, roleCfg)
	if err != nil {
		return nil, cloudv1.NewWFNonRetriableError(err)
	}
	region := envProvider.GetRegion()
	asgAPI := awsCli.AUTOSCALE(role, cloudaws.WithRegion(region))

	paused := make([]string, 0, len(asgs))
	for _, asg := range asgs {
		util.SetAnnotation(asg, crossplanePausedAnnotation, "true")
		if err := infraCli.Update(ctx, asg); err != nil {
			return paused, fmt.Errorf("pause crossplane reconciliation for ASG %s: %w", asg.Name, err)
		}
		_, err := asgAPI.SuspendProcesses(ctx, &autoscaling.SuspendProcessesInput{
			AutoScalingGroupName: aws.String(asg.Name),
			ScalingProcesses:     asgSuspendProcesses,
		})
		if err != nil {
			return paused, fmt.Errorf("suspend processes for ASG %s: %w", asg.Name, err)
		}
		logger.Info("paused ASG", "asg", asg.Name)
		paused = append(paused, asg.Name)
	}
	return paused, nil
}

// unpauseASGs resumes AWS scaling processes and clears Crossplane pause for
// ASGs belonging to the given group in the AZ. Idempotent.
func unpauseASGs(
	ctx context.Context,
	infraCli client.Client,
	awsCli cloudaws.Client,
	cluster *cloudv1.Cluster,
	envProvider dataplane.EnvProvider,
	roleCfg asgRoleConfig,
	group Group,
	zone string,
) ([]string, error) {
	logger := log.FromContext(ctx).WithValues("groupId", group.ID())
	if awsCli == nil || !isAWSProvider(cluster) {
		logger.Info("skip ASG unpause: not an AWS cluster or AWS client unavailable")
		return nil, nil
	}
	asgs, err := listGroupASGs(ctx, infraCli, group, zone)
	if err != nil {
		return nil, err
	}
	role, err := resolveASGAssumedRole(ctx, infraCli, cluster, envProvider, roleCfg)
	if err != nil {
		return nil, cloudv1.NewWFNonRetriableError(err)
	}
	region := envProvider.GetRegion()
	asgAPI := awsCli.AUTOSCALE(role, cloudaws.WithRegion(region))

	unpaused := make([]string, 0, len(asgs))
	for _, asg := range asgs {
		_, err := asgAPI.ResumeProcesses(ctx, &autoscaling.ResumeProcessesInput{
			AutoScalingGroupName: aws.String(asg.Name),
			ScalingProcesses:     asgSuspendProcesses,
		})
		if err != nil {
			return unpaused, fmt.Errorf("resume processes for ASG %s: %w", asg.Name, err)
		}
		anns := asg.GetAnnotations()
		if anns != nil {
			if _, ok := anns[crossplanePausedAnnotation]; ok {
				delete(anns, crossplanePausedAnnotation)
				asg.SetAnnotations(anns)
				if err := infraCli.Update(ctx, asg); err != nil {
					return unpaused, fmt.Errorf("unpause crossplane reconciliation for ASG %s: %w", asg.Name, err)
				}
			}
		}
		logger.Info("unpaused ASG", "asg", asg.Name)
		unpaused = append(unpaused, asg.Name)
	}
	return unpaused, nil
}