package zoneredundantevicter

import (
	"context"
	"errors"

	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/kubeapi"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	infraclient "github.com/tidbcloud/tidb-management-service/api/serverless_infra/client"
	"github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	corev1 "k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	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/log"
)

func (r *Reconciler) getZRClusterStatus(ctx context.Context) ([]*infra.ClusterStatus, error) {
	logger := log.FromContext(ctx)
	var clusterStatusList []*infra.ClusterStatus

	clusterKeyList := make([]meta.ObjectKey, 0)
	err := r.InfraClusterCli.GetIndex(ctx, "yes", "status.zone_redundant_enabled-index", &clusterKeyList)
	if err != nil && !errors.Is(err, infraclient.ErrNotFound) {
		return nil, err
	}
	for _, clusterKey := range clusterKeyList {
		c := &infra.Cluster{ObjectMeta: meta.ObjectMeta{Namespace: clusterKey.Namespace, Name: clusterKey.Name}}
		if err := r.InfraClusterCli.Get(ctx, c); err != nil {
			if errors.Is(err, infraclient.ErrNotFound) {
				continue
			}
			return nil, err
		}
		clusterStatusList = append(clusterStatusList, &c.Status)
		logger.V(2).Info("append zr cluster status", "cluster", c.Name)
	}

	var bls cloudv1.BranchList
	if err := r.InfraClient.List(ctx, &bls, client.MatchingFields(map[string]string{
		"status.fromProvider.zoneRedundantEnabled": "zoneRedundantEnabled",
	})); err != nil {
		return nil, err
	}
	for i := range bls.Items {
		c := &bls.Items[i]
		clusterStatus, err := c.AwsFreeTierClusterStatusFromBranch()
		if err != nil {
			return nil, err
		}
		clusterStatusList = append(clusterStatusList, &infra.ClusterStatus{
			TidbPool:      clusterStatus.TidbPool,
			KeyspaceToken: clusterStatus.KeyspaceToken,
		})
		logger.V(2).Info("append zr branch status", "cluster", c.Name)
	}

	return clusterStatusList, nil
}

func (r *Reconciler) getPodsOnNodes(ctx context.Context, nodes []*corev1.Node) (map[string]*corev1.Node, error) {
	logger := log.FromContext(ctx)
	podNodeMap := make(map[string]*corev1.Node)
	for _, node := range nodes {
		podList := &corev1.PodList{}
		if err := r.Client.List(ctx, podList, client.MatchingFields{"spec.nodeName": node.Name}); err != nil {
			return nil, err
		}
		for i := range podList.Items {
			if pod := &podList.Items[i]; pod.DeletionTimestamp == nil {
				podNodeMap[pod.Name] = node
				logger.V(2).Info("insert podNodeMap", "podName", pod.Name, "nodeName", node.Name)
			}
		}
	}
	return podNodeMap, nil
}

func (r *Reconciler) getEvictTidbPods(
	podNodeMap map[string]*corev1.Node,
	clusterStatusList []*infra.ClusterStatus,
) map[string][]string {
	podMap := make(map[string][]string)
	for _, clusterStatus := range clusterStatusList {
		if tManager := r.TidbManagers[clusterStatus.TidbPool]; tManager != nil {
			if endpoints := tManager.GetAssignedTidbByToken(clusterStatus.KeyspaceToken); endpoints != nil {
				for _, endpoint := range endpoints.Endpoints {
					podname, _ := ep.ParsePodFromEndpoint(endpoint)
					if podNodeMap[podname] != nil {
						podMap[clusterStatus.TidbPool] = append(podMap[clusterStatus.TidbPool], podname)
					}
				}
			}
		}
		for rgName, rgStatus := range clusterStatus.ResourceGroups {
			if tManager := r.TidbManagers[rgStatus.Pool]; tManager != nil {
				if endpoints := tManager.GetAssignedTidbByToken(rgName); endpoints != nil {
					for _, endpoint := range endpoints.Endpoints {
						podname, _ := ep.ParsePodFromEndpoint(endpoint)
						if podNodeMap[podname] != nil {
							podMap[rgStatus.Pool] = append(podMap[rgStatus.Pool], podname)
						}
					}
				}
			}
		}
	}
	return podMap
}

func (r *Reconciler) evictFreeCacheTidbPods(ctx context.Context, tidbPool, tidbNode string) error {
	for _, podname := range r.TidbManagers[tidbPool].GetTidbFromFreeCache(tidbNode, true) {
		if err := client.IgnoreNotFound(r.Client.Delete(ctx,
			&corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: podname, Namespace: tidbPool}},
			&client.DeleteOptions{GracePeriodSeconds: ptr.To[int64](0)}),
		); err != nil {
			return err
		}
		log.FromContext(ctx).Info("[pod lifecycle] delete free cache tidb pod",
			"podname", podname, "tidbPool", tidbPool, "tidbNode", tidbNode)
		metrics.ZREvictPodCounter.WithLabelValues(tidbPool, tidbNode, "free-cache-tidb", "delete").Inc()
	}
	return nil
}

func (r *Reconciler) evictTidbPod(ctx context.Context, tidbPod *corev1.Pod, tidbNode *corev1.Node) error {
	logger := log.FromContext(ctx)

	if kubeapi.IsNodeUnreachable(tidbNode) {
		if err := client.IgnoreNotFound(r.Client.Delete(ctx, tidbPod, &client.DeleteOptions{
			GracePeriodSeconds: ptr.To[int64](0),
		})); err != nil {
			return err
		}
		logger.Info("[pod lifecycle] delete assigned tidb pod", "pod", podName(tidbPod))
		metrics.ZREvictPodCounter.WithLabelValues(tidbPod.Namespace, tidbNode.Name, "assigned-tidb", "delete").Inc()
	} else {
		if err := client.IgnoreNotFound(r.Client.SubResource("eviction").Create(
			ctx, tidbPod, &policyv1.Eviction{}),
		); err != nil {
			return err
		}
		logger.Info("evict assigned tidb pod", "pod", podName(tidbPod))
		metrics.ZREvictPodCounter.WithLabelValues(tidbPod.Namespace, tidbNode.Name, "assigned-tidb", "evict").Inc()
	}
	return nil
}

func (r *Reconciler) evictZoneRedundantTidbPods(ctx context.Context, tidbNodes []*corev1.Node) (retry bool, err error) {
	if len(tidbNodes) == 0 {
		return
	}

	logger := log.FromContext(ctx).WithName("evictZoneRedundantTidbPods")

	defer func() {
		if err != nil {
			logger.Error(err, "failed to evictZoneRedundantTidbPods", "retry", retry)
			if retry {
				err = nil
			}
		}
	}()

	podNodeMap, err := r.getPodsOnNodes(ctx, tidbNodes)
	if err != nil {
		return
	}
	clusterStatusList, err := r.getZRClusterStatus(ctx)
	if err != nil {
		return
	}
	evictTidbPodMap := r.getEvictTidbPods(podNodeMap, clusterStatusList)

	for _, tidbNode := range tidbNodes {
		var updated bool
		if updated, err = r.trySetEvicterTaint(ctx, tidbNode); err != nil {
			return
		} else if updated {
			logger.Info("updated tidb node taint", "node", tidbNode.Name)
		}
		for tidbPool := range evictTidbPodMap {
			if err = r.evictFreeCacheTidbPods(ctx, tidbPool, tidbNode.Name); err != nil {
				return
			}
			r.Recorder.Eventf(tidbNode, corev1.EventTypeNormal, "DeleteZoneRedundantPods",
				"delete all free cache tidb pods in %s", tidbPool)
		}
	}
	for tidbPool, evictAssignedTidbPods := range evictTidbPodMap {
		tidbNodeMap := make(map[*corev1.Node]struct{})
		for _, evictAssignedTidbPod := range evictAssignedTidbPods {
			tidbNode := podNodeMap[evictAssignedTidbPod]
			tidbNodeMap[tidbNode] = struct{}{}
			pod := &corev1.Pod{ObjectMeta: metav1.ObjectMeta{Name: evictAssignedTidbPod, Namespace: tidbPool}}
			if err = r.evictTidbPod(ctx, pod, tidbNode); err != nil {
				if apierrors.IsTooManyRequests(err) {
					retry = true
					logger.Info(err.Error(), "tidbPod", evictAssignedTidbPod, "tidbPool", tidbPool)
					metrics.ZREvictPodRetryCounter.WithLabelValues(tidbPool, tidbNode.Name, "assigned-tidb").Inc()
					err = nil
					continue
				}
				return
			}
		}
		for tidbNode := range tidbNodeMap {
			reason := "EvictZoneRedundantPods"
			if kubeapi.IsNodeUnreachable(tidbNode) {
				reason = "DeleteZoneRedundantPods"
			}
			r.Recorder.Eventf(tidbNode, corev1.EventTypeNormal, reason,
				"delete assigned tidb pod in %s: %v", tidbPool, evictAssignedTidbPods)
		}
	}
	return
}