package tidbpodrescheduler

import (
	"context"
	"fmt"
	"time"

	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	apitypes "k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func (r *Reconciler) getRescheduledPendingPods(ctx context.Context, tidbPool string) []corev1.Pod {
	logger := log.FromContext(ctx, "tidbPool", tidbPool).WithName("getRescheduledPendingPods")
	pods := &corev1.PodList{}
	if err := r.Client.List(ctx, pods, client.InNamespace(tidbPool),
		client.MatchingFields{"tidb.rescheduler.pending": "true"}); err != nil {
		logger.Error(err, "failed to list rescheduled pending pods")
		return nil
	}
	return pods.Items
}

func (r *Reconciler) deleteRescheduledPendingPods(ctx context.Context, tidbPool string) (map[string]*corev1.Node, error) {
	logger := log.FromContext(ctx, "tidbPool", tidbPool).WithName("deleteRescheduledPendingPods")
	patchedNodes := make(map[string]*corev1.Node)
	for _, pod := range r.getRescheduledPendingPods(ctx, tidbPool) {
		if pod.Spec.NodeName != "" {
			continue
		}

		if nodeName := getNodeNameFromPodAffinity(&pod); nodeName != "" && patchedNodes[nodeName] == nil {
			rescheduleFailTimestamp := time.Now().Format(time.RFC3339)
			data := fmt.Sprintf(`{"metadata":{"annotations":{"%s":"%s"}}}`,
				types.AnnoNodeRescheduleFailTimestamp, rescheduleFailTimestamp)
			node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: nodeName}}
			if err := r.Client.Patch(ctx, node, client.RawPatch(apitypes.MergePatchType, []byte(data))); err != nil {
				return nil, err
			}
			patchedNodes[nodeName] = node
		}

		if err := r.Client.Delete(ctx, &pod); client.IgnoreNotFound(err) != nil {
			metrics.DeletePendingPodsCounter.WithLabelValues(pod.Namespace, "fail").Inc()
			logger.Error(err, "failed to delete rescheduled pod", "pod", pod.Namespace+"/"+pod.Name, "nodeName", pod.Spec.NodeName)
		} else {
			metrics.DeletePendingPodsCounter.WithLabelValues(pod.Namespace, "success").Inc()
			logger.Info("[pod lifecycle] delete rescheduled pod", "pod", pod.Namespace+"/"+pod.Name, "nodeName", pod.Spec.NodeName)
		}
	}
	return patchedNodes, nil
}

func ensureNodeNameAffinity(pod *corev1.Pod, nodeName string) {
	if pod.Spec.Affinity == nil {
		pod.Spec.Affinity = &corev1.Affinity{}
	}
	if pod.Spec.Affinity.NodeAffinity == nil {
		pod.Spec.Affinity.NodeAffinity = &corev1.NodeAffinity{}
	}

	if pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution == nil {
		pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution = &corev1.NodeSelector{}
	}
	req := pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution

	if len(req.NodeSelectorTerms) == 0 {
		req.NodeSelectorTerms = []corev1.NodeSelectorTerm{{}}
	}

	req.NodeSelectorTerms[0].MatchFields = append(req.NodeSelectorTerms[0].MatchFields, corev1.NodeSelectorRequirement{
		Key:      "metadata.name",
		Operator: corev1.NodeSelectorOpIn,
		Values:   []string{nodeName},
	})
}

func getNodeNameFromPodAffinity(pod *corev1.Pod) string {
	if pod.Spec.Affinity != nil && pod.Spec.Affinity.NodeAffinity != nil &&
		pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution != nil &&
		len(pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms) > 0 {
		for _, term := range pod.Spec.Affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms {
			for _, req := range term.MatchFields {
				if req.Key == "metadata.name" && req.Operator == corev1.NodeSelectorOpIn {
					return req.Values[0]
				}
			}
		}
	}
	return ""
}

func incomingPodCount(ctx context.Context, nodeName string, abnormalList []corev1.Pod) (count int) {
	logger := log.FromContext(ctx, "nodeName", nodeName).WithName("incomingPodCount")
	for i := range abnormalList {
		if pod := &abnormalList[i]; pod.Spec.NodeName == nodeName {
			count++
			logger.V(1).Info("incoming pod", "pod", pod.Namespace+"/"+pod.Name)
		}
	}
	return
}