package tidbmanager

import (
	"context"
	"errors"
	"fmt"
	"math"
	"net"
	"os"
	"strconv"
	"time"

	"github.com/go-logr/logr"
	perrors "github.com/pingcap/errors"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	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"
)

const (
	AntiAffinityScorePerReplica = 40
	OtherErrScoreDeduct         = 20
)

func (t *tidbManager) patchCriticalTidbPod(ctx context.Context, log logr.Logger, name string) {
	logger := log.WithName("patchCriticalTidbPod").WithValues("pod name", name, "pool", t.tidbPoolNamespace)
	serverlessPoolConfigCm := &corev1.ConfigMap{}
	if err := t.client.Get(ctx, client.ObjectKey{
		Name:      types.ServerlessPoolConfig,
		Namespace: t.tidbPoolNamespace,
	}, serverlessPoolConfigCm); err != nil {
		logger.Error(err, "get cm serverlessPoolConfig")
		return
	}
	if serverlessPoolConfigCm.Data[types.TidbCriticalKey] == "true" {
		logger.Info("patching critical tidb pod")
		data := fmt.Sprintf(`{"metadata":{"annotations":{"%s":"%s"}}}`, t.safeToEvict, "false")
		if err := t.client.Patch(ctx, &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: t.tidbPoolNamespace,
		}}, client.RawPatch(apitypes.MergePatchType, []byte(data))); err != nil {
			logger.Error(err, "patch critical tidb pod")
			return
		}
	}
}

// func (t *tidbManager) removeRunSlowAnno(ctx context.Context, log logr.Logger, name string) {
// 	logger := log.WithName("removeRunSlowAnno").WithValues("pod name", name, "pool", t.tidbPoolNamespace)
// 	tidbPod := &corev1.Pod{}
// 	if err := t.client.Get(ctx, client.ObjectKey{Name: name, Namespace: t.tidbPoolNamespace}, tidbPod); err != nil {
// 		logger.Error(err, "get tidb pod")
// 		return
// 	}
// 	if tidbPod.Annotations[types.AnnoTidbStatusRunSlow] != "" {
// 		logger.Info("removing tidb pod AnnoTidbStatusRunSlow")
// 		data := fmt.Sprintf(`{"metadata":{"annotations":{"%s":null}}}`, types.AnnoTidbStatusRunSlow)
// 		if err := t.client.Patch(ctx, &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
// 			Name:      name,
// 			Namespace: t.tidbPoolNamespace,
// 		}}, client.RawPatch(apitypes.StrategicMergePatchType, []byte(data))); err != nil {
// 			logger.Error(err, "remove tidb pod AnnoTidbStatusRunSlow")
// 			return
// 		}
// 	}
// }

func (t *tidbManager) AssignTidb(
	ctx context.Context,
	token string,
	param *tidb.ActivateParameter,
) (endpoints []string, err error) {
	if t.migrationLocker != nil && t.migrationLocker.IsLocked(utils.TrimToken(token)) {
		return nil, perrors.Annotatef(utils.ErrKeyspaceMetaServiceGroupMigrating, "AssignTidb for token %s", token)
	}
	if tInfo, ok := t.assignCache.getTidbInfo(token); ok {
		start := time.Now()
		defer func() {
			stage := "total"
			if err != nil {
				stage = "total_fail"
			}
			metrics.Reclaimer.GetHistogramVec(metrics.AssignTidbHistogramName, token).WithLabelValues(stage, t.tidbPoolNamespace).
				Observe(time.Since(start).Seconds())
		}()

		switch {
		case tInfo.replicas <= 0:
			return nil, utils.ErrZeroReplica
		case len(tInfo.endpoints) < tInfo.replicas:
			endpoints, err = t.assignWithLock(ctx, token, tInfo.replicas, true, param)
			if err != nil {
				if !errors.Is(err, utils.ErrNotFoundToken) && len(tInfo.endpoints) > 0 {
					// if token not deleted, we can return assignCache's result and log error
					t.Error(err, "assignWithLock error, but still return old endpoints")
					return tInfo.endpoints, nil
				}
				return nil, err
			}
			return endpoints, nil
		default:
			return tInfo.endpoints, nil
		}
	}
	return nil, perrors.Annotatef(utils.ErrNotFoundToken, "failed to AssignTidb for token %s", token)
}

func (t *tidbManager) assignWithLock(
	ctx context.Context,
	token string,
	replicas int,
	verifyReplicas bool,
	param *tidb.ActivateParameter,
) (endpoints []string, err error) {
	if t.isAssignCacheSizeLimitExceeded(ctx) {
		return nil, utils.ErrAssignCacheSizeLimitExceeded
	}

	tokenLog := t.WithValues("token", token)
	var name, ip, targetNode string
	defer func() {
		if err == nil && name != "" {
			t.patchCriticalTidbPod(ctx, tokenLog, name)
			// t.removeRunSlowAnno(ctx, tokenLog, name)
		}
	}()

	// stage 1: Lock
	start := time.Now()
	tLocker, closeSession, err := t.disLockerFunc(token)
	if err != nil {
		return nil, err
	}
	// Why: closeSession revokes the etcd lease that backs the distributed mutex.
	// Without it, NewSession's keepalive goroutines + lease live forever, leaking ~2 goroutines
	// and one lease per call. Deferred LIFO so Unlock runs before Close.
	defer closeSession()
	// tLocker := t.locLocker.AssignLocker(token)
	// defer t.locLocker.ReleaseLocker(token)

	tLocker.Lock()
	defer tLocker.Unlock()
	metrics.Reclaimer.GetHistogramVec(metrics.AssignTidbHistogramName, token).WithLabelValues("DisLock", t.tidbPoolNamespace).
		Observe(time.Since(start).Seconds())

	// stage 2: GetFromEtcd
	start = time.Now()
	resp, err := t.etcdKV.Get(ctx, t.etcdTokenPrefix+token+"/"+utils.KeyEndpoints)
	if err != nil {
		return nil, err
	}
	metrics.Reclaimer.GetHistogramVec(metrics.AssignTidbHistogramName, token).WithLabelValues("GetFromEtcd", t.tidbPoolNamespace).
		Observe(time.Since(start).Seconds())
	if len(resp.Kvs) == 0 {
		return nil, perrors.Annotatef(utils.ErrNotFoundToken, "failed to GetFromEtcd for token %s", token)
	}
	kv := resp.Kvs[0]
	tokenLog.V(3).Info("get real-time in etcd", "endpoints", string(kv.Value))
	tInfo := newTidbInfo(string(kv.Value), replicas, kv.ModRevision)
	if len(tInfo.endpoints) >= replicas {
		return tInfo.endpoints, nil
	}

	// stage 3: Assign
	start = time.Now()
	antiAffinityNodes := t.getNodesFromEndpoints(ctx, tInfo.endpoints)
	var otherErrNodes []string
	var otherErrRetry = 1
	for {
		if err = t.rateLimiter.wait(ctx); err != nil {
			return nil, err
		}

		name, ip, targetNode, err = t.assignOne(ctx, antiAffinityNodes, otherErrNodes)
		if err != nil {
			if errors.Is(err, utils.ErrNotEnoughPod) {
				if pool := t.tidbPool.GetPool(ctx); pool != nil && pool.MinSize == 0 {
					t.tryCreateNewPod(ctx, token+strconv.Itoa(replicas), pool.TidbPod)
				}
			}
			return nil, err
		}

		tokenLog.Info("[pod lifecycle] activating pod", "pod_name", name, "node", targetNode)
		if err = t.activator.do(ctx, net.JoinHostPort(ip, "10080"), token, name, param); err == nil {
			break
		}
		if errors.Is(err, context.Canceled) {
			t.freeCache.tryResetActivating(targetNode, name, ip)
			return nil, err
		}
		if errors.Is(err, tidb.ErrActivateTerminating) {
			metrics.ActivateTerminatingCounter.WithLabelValues(t.tidbPoolNamespace).Inc()
			tokenLog.Info("[pod lifecycle] quarantine terminating pod", "pod_name", name, "node", targetNode)
			t.enqueueStandbyCheck(name, ip)
			continue
		}
		if errors.Is(err, tidb.ErrActivateByOtherToken) {
			t.freeCache.removePodIfIP(name, ip)
		} else {
			t.freeCache.tryResetActivating(targetNode, name, ip)
			if otherErrRetry <= 0 {
				return nil, err
			}
			otherErrRetry--
			otherErrNodes = append(otherErrNodes, targetNode)
			time.Sleep(1 * time.Second)
		}
	}
	metrics.Reclaimer.GetHistogramVec(metrics.AssignTidbHistogramName, token).WithLabelValues("Assign", t.tidbPoolNamespace).
		Observe(time.Since(start).Seconds())

	tokenLog.Info("[pod lifecycle] assign pod to cluster", "pod_name", name)

	// stage 4: SyncEtcd
	start = time.Now()
	if err = t.modifyTidbInfo(ctx, token, ep.GenEndpointFromPod(name, ip), &tInfo); err != nil {
		tokenLog.Error(err, "modifyTidbInfo", "pod name", name)
		t.removePod(removePod{name: name, force: true})
		return nil, err
	}
	metrics.Reclaimer.GetHistogramVec(metrics.AssignTidbHistogramName, token).WithLabelValues("SyncEtcd", t.tidbPoolNamespace).
		Observe(time.Since(start).Seconds())

	if verifyReplicas {
		// if replicas scale down in activator, this assigned pod is an extra replica.
		if newtInfo, ok := t.assignCache.getTidbInfo(token); ok {
			if len(tInfo.endpoints) > newtInfo.replicas {
				tokenLog.Info("found extra replica", "endpoints", tInfo.endpoints, "new replicas", newtInfo.replicas)
				t.removePod(removePod{name: name, force: true})
				return nil, utils.ErrExtraReplica
			}
		} else {
			return nil, perrors.Annotatef(utils.ErrNotFoundToken, "failed to SyncEtcd for token %s", token)
		}
	}

	return tInfo.endpoints, nil
}

func (t *tidbManager) isAssignCacheSizeLimitExceeded(ctx context.Context) bool {
	serverlessPoolConfigCm := &corev1.ConfigMap{}
	if err := t.client.Get(ctx, client.ObjectKey{
		Name:      types.ServerlessPoolConfig,
		Namespace: t.tidbPoolNamespace,
	}, serverlessPoolConfigCm); err != nil {
		return false
	}

	limit, err := strconv.Atoi(serverlessPoolConfigCm.Data[types.AssignCacheSizeLimitKey])
	if err != nil || limit <= 0 {
		return false
	}

	_, endpointSize := t.assignCache.getSize()
	return endpointSize >= limit
}

func (t *tidbManager) tryCreateNewPod(ctx context.Context, reqID string, tidbPod *corev1.Pod) {
	abnormalList, err := t.tidbPool.GetAbnormalPodList(ctx)
	if err != nil {
		t.Error(err, "get abnormal tidb pod list", "reqID", reqID)
		return
	}
	for i := range abnormalList {
		if abnormalList[i].Annotations[types.AnnoTidbCreatedByTManager] == reqID {
			t.Info("the previous reqID tidb pod is not ready, skip create by tManager", "reqID", reqID)
			return
		}
	}

	pod := tidbPod.DeepCopy()
	pod.Annotations = utils.UpdateMap(pod.Annotations, types.AnnoTidbCreatedByTManager, reqID)
	if err = t.client.Create(ctx, pod); err != nil {
		t.Error(err, "failed to create new pod", "pod_name", pod.Name, "reqID", reqID)
	} else {
		t.Info("[pod lifecycle] create new pod", "pod_name", pod.Name, "reqID", reqID)
	}
}

func (t *tidbManager) getNodesFromEndpoints(ctx context.Context, endpoints []string) (nodes []string) {
	for _, endpoint := range endpoints {
		podName, _ := ep.ParsePodFromEndpoint(endpoint)
		pod := &corev1.Pod{}
		if err := t.client.Get(ctx, client.ObjectKey{Name: podName, Namespace: t.tidbPoolNamespace}, pod); err != nil {
			continue
		}

		if pod.Spec.NodeName != "" {
			nodes = append(nodes, pod.Spec.NodeName)
		}
	}
	return
}

func (t *tidbManager) assignOne(
	ctx context.Context, antiAffinityNodes, otherErrNodes []string,
) (podName, podIP, targetNode string, err error) {
	nodeScoreObj := &serverlessv1alpha1.NodeScore{}
	if err := t.client.Get(ctx, client.ObjectKey{Name: types.NodeScoreMap, Namespace: os.Getenv("NAMESPACE")},
		nodeScoreObj); err != nil {
		t.Error(err, "failed to get node score, skip sort")
	}

	var maxScore = math.MinInt
	t.freeCache.walkNodeByLock(func(nodeName string, podInfo map[string]freePodInfo) bool {
		availableNode := false
		for _, info := range podInfo {
			if info.isAvailable() {
				availableNode = true
				break
			}
		}
		if availableNode {
			nodeScore := nodeScoreObj.Status.SchedulerNodeScore[nodeName] // if node not exist, its score is 0
			nodeScore -= findDuplicateCount(nodeName, antiAffinityNodes) * AntiAffinityScorePerReplica
			nodeScore -= findDuplicateCount(nodeName, otherErrNodes) * OtherErrScoreDeduct
			if nodeScore > maxScore {
				maxScore = nodeScore
				targetNode = nodeName
			}
		}
		return true
	}, func() {
		podName, podIP = t.freeCache.assignOneOnNode(targetNode)
	})

	if podName == "" || podIP == "" {
		err = utils.ErrNotEnoughPod
		return
	}

	t.V(2).Info("assignOne success", "maxScore", maxScore, "targetNode", targetNode,
		"pod name", podName, "pod ip", podIP)
	metrics.AssignTidbScoreHistogram.WithLabelValues(t.tidbPoolNamespace).Observe(float64(maxScore))
	return
}

func findDuplicateCount(name string, list []string) (count int) {
	for _, v := range list {
		if name == v {
			count++
		}
	}
	return
}