package tidbupgrader

import (
	"context"
	"errors"
	"fmt"
	"net"
	"time"

	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	maxTryTickCount    = 8
	podDeletionPoll    = 10 * time.Second
	podDeletionTimeout = 3600 * time.Second
)

// handleAssignCache manages the upgrade process for TiDB pods in the assign cache.
// It controls the batch deletion of old pods and creation of new pods while maintaining
// stability by monitoring changes in the assign cache size and handling abnormal pods.
func (tu *TidbUpgrader) handleAssignCache(batchSize int) error { //nolint:lll // by design
	allOldPodMap, oldPodInAssignCacheSize := tu.tManager.GetOldTidbFromAssignCache(tu.tidbPool.Revision, -1)
	batchOldPods := tu.filterFirstPodOfEachToken(allOldPodMap, batchSize)
	_, assignCacheSize := tu.tManager.AssignCacheSize()

	metrics.OldPodAssignCacheSizeGauge.WithLabelValues(tu.tidbPool.Revision, tu.tidbPoolNamespace).
		Set(float64(oldPodInAssignCacheSize))
	metrics.AssignCacheSizeGauge.WithLabelValues(tu.tidbPool.Revision, tu.tidbPoolNamespace).
		Set(float64(assignCacheSize))

	if tu.record.preOldPodInAssignCacheSize != oldPodInAssignCacheSize {
		tu.V(1).Info("oldPodInAssignCacheSize change",
			"pre", tu.record.preOldPodInAssignCacheSize, "now", oldPodInAssignCacheSize)
		tu.record.preOldPodInAssignCacheSize = oldPodInAssignCacheSize
		tu.record.tryTickCount = 0
		return nil
	}

	requiredNewRevisionSize := tu.tidbPool.MinSize
	if requiredNewRevisionSize == 0 && oldPodInAssignCacheSize > 0 {
		requiredNewRevisionSize = 1
	}
	newRevisionSize := len(tu.tManager.GetTidbFromFreeCacheByRevision(tu.tidbPool.Revision, true))
	if newRevisionSize < requiredNewRevisionSize {
		tu.Info("wait for enough new pods in free cache",
			"newRevisionSize", newRevisionSize,
			"requiredNewRevisionSize", requiredNewRevisionSize,
			"poolMinSize", tu.tidbPool.MinSize)
		if tu.tidbPool.MinSize == 0 {
			if err := tu.tManager.TryCreatePod(tu.ctx); err != nil {
				tu.Error(err, "create tidb pod")
			}
		}
		return errors.New("not enough new pods in free cache")
	}
	// oldPodInAssignCacheSize not change, we can delete old tidb pods in free cache
	for _, oldPodName := range tu.oldPodInFreeCache {
		if err := tu.deleteOldPodInFreeCache(oldPodName); err != nil {
			tu.Error(err, "delete tidb pod in free cache", "pod name", oldPodName)
		}
	}

	if tu.record.tryTickCount < maxTryTickCount {
		tu.Info("oldPodInAssignCacheSize not change, try again",
			"oldPodInAssignCacheSize", oldPodInAssignCacheSize, "tryTickCount", tu.record.tryTickCount)
		tu.record.tryTickCount++
		return nil
	}

	if len(batchOldPods) == 0 {
		return nil
	}
	// oldPodInAssignCacheSize keep stable, we can delete batchOldPods in assign cache every tick
	tu.Info("oldPodInAssignCacheSize keep stable, delete batchOldPods in assign cache",
		"batchOldPods size", len(batchOldPods))

	for _, oldPodName := range batchOldPods {
		if err := tu.tManager.TryCreatePod(tu.ctx); err != nil {
			tu.Error(err, "create tidb pod")
		}

		// make sure we have enough available free tidb pods to assign
		if len(tu.tManager.GetTidbFromFreeCacheByRevision(tu.tidbPool.Revision, true)) > 0 {
			if err := tu.tManager.ReplaceTidb(tu.ctx, oldPodName, tu.deleteOldPodInAssignCache); err != nil {
				tu.Error(err, "failed to replace tidb pod", "pod_name", oldPodName)
				continue
			}

			tu.record.preOldPodInAssignCacheSize--
		}
	}

	return nil
}

// filterFirstPodOfEachToken filters the first pod of each token. It ensures only one pod is deleted in one cycle of each token.
func (tu *TidbUpgrader) filterFirstPodOfEachToken(podMap map[string][]string, limit int) []string {
	batchPods := make([]string, 0, limit)
	for _, podNames := range podMap {
		if len(podNames) == 0 {
			continue
		} else if len(podNames) == 1 {
			batchPods = append(batchPods, podNames[0])
		} else {
			batchPods = append(batchPods, tu.skipAutoIDOwner(podNames))
		}

		if len(batchPods) >= limit {
			break
		}
	}

	return batchPods
}

func (tu *TidbUpgrader) skipAutoIDOwner(podNames []string) string {
	for _, podName := range podNames {
		pod := &corev1.Pod{}
		err := tu.client.Get(tu.ctx, client.ObjectKey{Name: podName, Namespace: tu.tidbPoolNamespace}, pod)
		if err != nil {
			tu.Error(err, "Get pod", "pod name", podName)
			continue
		}
		isOwner, err := tu.tidbController.IsAutoIdOwner(tu.ctx, net.JoinHostPort(pod.Status.PodIP, "10080"))
		if err != nil {
			tu.Error(err, "check auto id owner", "pod name", podName, "pod ip", pod.Status.PodIP)
			continue
		}
		if isOwner {
			tu.Info("skip delete pod which is auto id owner", "pod name", podName, "pod ip", pod.Status.PodIP)
			continue
		}
		return podName
	}

	return podNames[0]
}

func (tu *TidbUpgrader) processRegularPool() error {
	batchSize := int(float64(tu.freeCacheSize) * (1 - minOldRevPercent))
	if batchSize == 0 {
		batchSize = 1
	}
	return tu.handleAssignCache(batchSize)
}

func (tu *TidbUpgrader) deleteOldPodInAssignCache(ctx context.Context, podName, namespace string) error {
	removePod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      podName,
			Namespace: namespace,
		},
	}

	tu.Info("[pod lifecycle] delete tidb pod in assign cache", "pod name", podName)
	err := tu.client.Delete(ctx, removePod)
	if err != nil {
		return client.IgnoreNotFound(err)
	}

	// Wait for pod to be fully deleted
	return tu.waitForPodDeleted(ctx, podName, namespace)
}

// waitForPodDeleted waits for a pod to be fully deleted from the cluster
func (tu *TidbUpgrader) waitForPodDeleted(ctx context.Context, podName, namespace string) error {
	tu.Info("[pod lifecycle] waiting for pod to be fully deleted", "pod name", podName)

	pod := &corev1.Pod{}
	start := time.Now()
	return wait.PollUntilContextTimeout(ctx, podDeletionPoll, podDeletionTimeout, true, func(ctx context.Context) (bool, error) {
		err := tu.client.Get(ctx, client.ObjectKey{Name: podName, Namespace: namespace}, pod)
		if apierrors.IsNotFound(err) {
			metrics.PodDeletionDurationGauge.DeletePartialMatch(map[string]string{"namespace": namespace, "pod_name": podName})
			tu.Info("[pod lifecycle] pod fully deleted", "pod name", podName)
			return true, nil
		}
		if err != nil {
			return false, fmt.Errorf("failed to get pod %s/%s: %w", namespace, podName, err)
		}

		elapsed := time.Since(start)
		metrics.PodDeletionDurationGauge.WithLabelValues(namespace, podName).Set(float64(elapsed.Seconds()))
		tu.Info("pod still exists, waiting for deletion", "pod name", podName, "deletionTimestamp", pod.DeletionTimestamp, "elapsed", elapsed)

		return false, nil
	})
}