package nodepool

import (
	"context"
	"fmt"

	"github.com/hashicorp/go-multierror"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/api"
	"github.com/tidbcloud/infra-provider/pkg/config"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/controllers/utils"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/pkg/nodegroup"
)

const (
	// `registry.k8s.io/pause:3.9`
	pauseImage   = "pause"
	pauseVersion = "3.9"

	zoneCount = 3
)

// CreateOverProvisioningPods creates the Kubernetes-upgrade capacity for each
// active TiDB component. useMinDomains selects the new three-zone contract.
func CreateOverProvisioningPods(
	ctx context.Context,
	remoteCli client.Client,
	cluster *cloudv1.Cluster,
	tcName, tcNS, k8sVersion string,
	provider gitops.InstanceTypeProvider,
	useMinDomains bool,
) error {
	var errs *multierror.Error

	for _, comp := range cluster.GetAllComponents() {
		replicas := comp.GetReplicas()
		if replicas < 1 || comp.NeedSuspend(cluster.Spec.Cluster.SuspendAction) {
			continue
		}
		compType := comp.GetComponentType()
		if api.GetInstanceTypeScalePhase(cluster, compType) != cloudv1.InstanceTypeScalePhaseNone {
			continue
		}

		otherCompTypes := make([]string, 0, 3)
		for _, compTypes2 := range cluster.GetAllComponentTypes() {
			if compType != compTypes2 {
				otherCompTypes = append(otherCompTypes, string(compTypes2))
			}
		}

		// Each component has its own over-provisioning deployment
		labels := map[string]string{
			consts.InstanceLabelKey:                  tcName,
			consts.ComponentLabelKey:                 consts.OverProvisioningComponentLabelValue,
			consts.OverProvisioningComponentLabelKey: string(compType),
		}
		opDeploy := &appsv1.Deployment{
			ObjectMeta: metav1.ObjectMeta{
				Name:      namer.BuildOverProvisioningName(fmt.Sprintf("%s-%s", compType, nodegroup.NormalizeVersion(k8sVersion))),
				Namespace: tcNS,
				Labels:    labels,
			},
			Spec: appsv1.DeploymentSpec{
				Replicas: pointer.Int32(zoneCount),
				Selector: &metav1.LabelSelector{
					MatchLabels: labels,
				},
				Template: corev1.PodTemplateSpec{
					ObjectMeta: metav1.ObjectMeta{
						Labels: labels,
					},
					Spec: corev1.PodSpec{
						TerminationGracePeriodSeconds: pointer.Int64(0),
						NodeSelector: map[string]string{
							consts.K8sVersionLabelKey:   k8sVersion,
							consts.InstanceTypeLabelKey: comp.GetInstanceType(),
						},
						Affinity: &corev1.Affinity{
							PodAntiAffinity: &corev1.PodAntiAffinity{
								RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
									{
										// anti-affinity with other over-provisioning pods for all components
										NamespaceSelector: &metav1.LabelSelector{}, // `An empty selector ({}) matches all namespaces`
										LabelSelector: &metav1.LabelSelector{
											MatchLabels: map[string]string{
												consts.InstanceLabelKey:  tcName,
												consts.ComponentLabelKey: consts.OverProvisioningComponentLabelValue,
											},
										},
										TopologyKey: corev1.LabelHostname,
									},
									{
										// anti-affinity with all TiDB Cluster components in other namespaces
										NamespaceSelector: &metav1.LabelSelector{
											MatchExpressions: []metav1.LabelSelectorRequirement{
												{
													Key:      corev1.LabelMetadataName, // kubernetes.io/metadata.name
													Operator: metav1.LabelSelectorOpNotIn,
													Values:   []string{tcNS},
												},
											},
										},
										LabelSelector: &metav1.LabelSelector{
											MatchExpressions: []metav1.LabelSelectorRequirement{
												{
													Key:      consts.ComponentLabelKey,
													Operator: metav1.LabelSelectorOpIn,
													Values: []string{
														string(cloudv1.ComponentPD),
														string(cloudv1.ComponentTiKV),
														string(cloudv1.ComponentTiDB),
														string(cloudv1.ComponentTiFlash),
														string(cloudv1.ComponentTiProxy),
													},
												},
											},
										},
										TopologyKey: corev1.LabelHostname,
									},
									{
										// anit-affinity with other TiDB Cluster components in the same namespace
										LabelSelector: &metav1.LabelSelector{
											MatchExpressions: []metav1.LabelSelectorRequirement{
												{
													Key:      consts.ComponentLabelKey,
													Operator: metav1.LabelSelectorOpIn,
													Values:   otherCompTypes,
												},
											},
										},
										TopologyKey: corev1.LabelHostname,
									},
								},
							},
						},
						TopologySpreadConstraints: []corev1.TopologySpreadConstraint{
							{
								// spread over-provisioning pods evenly across all zones,
								// in fact all zones should have the same number of over-provisioning pods as `replicas` is `zoneCount*N`
								MaxSkew:           1,
								TopologyKey:       corev1.LabelTopologyZone,
								WhenUnsatisfiable: corev1.DoNotSchedule,
								LabelSelector: &metav1.LabelSelector{
									MatchLabels: labels,
								},
							},
						},
						Tolerations: utils.TolerationWithComponentShared(),
						Containers: []corev1.Container{
							{
								Name:  "over-provisioning",
								Image: config.GetThirdPartyImage(pauseImage, pauseVersion),
								Resources: corev1.ResourceRequirements{
									Requests: corev1.ResourceList{
										corev1.ResourceCPU: resource.MustParse("5m"),
									},
								},
							},
						},
					},
				},
			},
		}

		if cluster.GetProviderType() == cloudv1.ProviderTypeGKE {
			if cluster.FeatureEnabled(consts.ClusterFeatureGCP3LayerNodePool) {
				opDeploy.Spec.Template.Spec.Tolerations = append(opDeploy.Spec.Template.Spec.Tolerations, corev1.Toleration{
					Key:      "3-layer-arch",
					Operator: corev1.TolerationOpEqual,
					Value:    "true",
					Effect:   corev1.TaintEffectNoSchedule,
				})
				opDeploy.Spec.Template.Spec.NodeSelector[consts.LabelKeyFor3LayerArch] = consts.LabelValTrue
				if useMinDomains {
					constraint := &opDeploy.Spec.Template.Spec.TopologySpreadConstraints[0]
					constraint.MinDomains = pointer.Int32(zoneCount)
				}
			}

			if provider != nil {
				it, err := provider.FindInstanceType(comp.GetInstanceType())
				if err != nil {
					errs = multierror.Append(errs, fmt.Errorf(
						"failed query arch of instance type(%s): %w", comp.GetInstanceType(), err))
					continue
				}
				if it.Arch == consts.ARM64ArchName {
					opDeploy.Spec.Template.Spec.Tolerations = append(opDeploy.Spec.Template.Spec.Tolerations, *utils.TolerationWithArch(consts.ARM64ArchName))
				}
			}
		}

		if !cluster.FeatureEnabled(consts.ClusterFeatureDisableNodePoolDataPlaneLabel) {
			opDeploy.Spec.Template.Spec.NodeSelector[consts.LabelKeyNodeGroupType] = consts.LabelValNodeGroupTypeDataPlane
		}

		if err := remoteCli.Create(ctx, opDeploy); err != nil && !apierrors.IsAlreadyExists(err) {
			errs = multierror.Append(errs, fmt.Errorf("failed to create over-provisioning deployment for %s: %w", compType, err))
		}
	}

	return errs.ErrorOrNil()
}

// HasUnlabeledUpgradeOverProvisioning reports whether an in-progress
// Kubernetes upgrade predates the used-for label.
func HasUnlabeledUpgradeOverProvisioning(
	ctx context.Context,
	remoteCli client.Client,
	tcName, tcNS string,
) (bool, error) {
	opDeployList := &appsv1.DeploymentList{}
	if err := remoteCli.List(ctx, opDeployList, client.InNamespace(tcNS), client.MatchingLabels{
		consts.InstanceLabelKey:  tcName,
		consts.ComponentLabelKey: consts.OverProvisioningComponentLabelValue,
	}); err != nil {
		return false, fmt.Errorf("failed to list over-provisioning deployments: %w", err)
	}
	for _, deploy := range opDeployList.Items {
		// Older Kubernetes upgrade Deployments predate the used-for label;
		// instance-type and standby Deployments always set it.
		if namer.HasOverProvisioningNamePrefix(deploy.Name) &&
			deploy.Labels[consts.OverProvisioningUsedForLabelKey] == "" {
			return true, nil
		}
	}
	return false, nil
}

func AllOverProvisioningPodsReady(ctx context.Context, remoteCli client.Client, tcName, tcNS string) bool {
	ready, total := CountOverProvisioningPodsReady(ctx, remoteCli, tcName, tcNS)
	return total > 0 && ready == total
}

// CountOverProvisioningPodsReady returns how many over-provisioning deployments are fully ready.
func CountOverProvisioningPodsReady(ctx context.Context, remoteCli client.Client, tcName, tcNS string) (ready, total int) {

	logger := log.FromContext(ctx).WithValues("tcName", tcName, "tcNS", tcNS)

	opDeployList := &appsv1.DeploymentList{}
	if err := remoteCli.List(ctx, opDeployList, client.InNamespace(tcNS), client.MatchingLabels{
		consts.InstanceLabelKey:  tcName,
		consts.ComponentLabelKey: consts.OverProvisioningComponentLabelValue,
	}); err != nil {
		logger.Error(err, "failed to list over-provisioning deployments")
		return 0, 0
	}

	for _, deploy := range opDeployList.Items {
		if !namer.HasOverProvisioningNamePrefix(deploy.Name) {
			continue
		}
		total++
		if deploy.Spec.Replicas == nil {
			continue
		}
		if deploy.Status.ReadyReplicas == *deploy.Spec.Replicas {
			ready++
		}
	}
	return ready, total
}

// DeleteOverProvisioningDeployments deletes some (or all) over provisioning deployments after all components are restarted to the new nodes.
func DeleteOverProvisioningDeployments(ctx context.Context, remoteCli client.Client, tcName, tcNS string) error {
	opDeployList := &appsv1.DeploymentList{}
	mLabels := client.MatchingLabels{
		consts.InstanceLabelKey:  tcName,
		consts.ComponentLabelKey: consts.OverProvisioningComponentLabelValue,
	}
	if err := remoteCli.List(ctx, opDeployList, client.InNamespace(tcNS), mLabels); err != nil && !apierrors.IsNotFound(err) {
		return fmt.Errorf("failed to list over-provisioning deployments: %w", err)
	}

	for _, deploy := range opDeployList.Items {
		if !namer.HasOverProvisioningNamePrefix(deploy.Name) {
			continue
		}
		if err := remoteCli.Delete(ctx, &deploy); client.IgnoreNotFound(err) != nil {
			return fmt.Errorf("failed to delete over-provisioning deployment %s: %w", deploy.Name, err)
		}
	}

	return nil
}