package tidbmanager

import (
	"context"
	"net"
	"strconv"
	"time"

	"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"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	k8stypes "k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/util/workqueue"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

const (
	standbyCheckStatusPort = 10080
	standbyCheckTimeout    = 2 * time.Second
	standbyCheckBaseDelay  = time.Second
	standbyCheckMaxDelay   = 10 * time.Second

	standbyCheckResultRestored   = "restored"
	standbyCheckResultAssigned   = "assigned"
	standbyCheckResultPodGone    = "pod_gone"
	standbyCheckResultReplaced   = "replaced"
	standbyCheckResultActivated  = "activated"
	standbyCheckResultCanceled   = "canceled"
	standbyCheckResultIneligible = "ineligible"
)

type standbyCheck struct {
	podName string
	podIP   string
}

type standbyCheckQueue interface {
	workqueue.TypedRateLimitingInterface[standbyCheck]
}

func newStandbyCheckQueue() standbyCheckQueue {
	return workqueue.NewTypedRateLimitingQueue(
		workqueue.NewTypedItemExponentialFailureRateLimiter[standbyCheck](
			standbyCheckBaseDelay,
			standbyCheckMaxDelay,
		),
	)
}

func (t *tidbManager) enqueueStandbyCheck(podName, podIP string) {
	if t.standbyCheckQueue == nil {
		return
	}
	item := standbyCheck{podName: podName, podIP: podIP}
	t.standbyCheckMu.Lock()
	if t.standbyChecks == nil {
		t.standbyChecks = make(map[standbyCheck]struct{})
	}
	t.standbyChecks[item] = struct{}{}
	t.freeCache.quarantinePodIfIP(podName, podIP)
	t.syncStandbyQuarantinedPodsMetricLocked()
	t.standbyCheckMu.Unlock()
	t.standbyCheckQueue.Add(item)
	t.Info("[pod lifecycle] enqueue standby check", "pod name", podName, "pod ip", podIP)
}

func (t *tidbManager) processStandbyCheck(workerCount int) {
	if t.standbyCheckQueue == nil {
		return
	}
	for range workerCount {
		go t.standbyCheckWorker()
	}
	go func() {
		<-t.ctx.Done()
		t.standbyCheckQueue.ShutDown()
	}()
}

func (t *tidbManager) standbyCheckWorker() {
	for t.processNextStandbyCheck() {
	}
}

func (t *tidbManager) processNextStandbyCheck() bool {
	item, shutdown := t.standbyCheckQueue.Get()
	if shutdown {
		return false
	}
	defer t.standbyCheckQueue.Done(item)

	if t.checkStandbyAndSetFreeCache(t.ctx, item) {
		t.standbyCheckQueue.AddRateLimited(item)
	} else {
		t.standbyCheckQueue.Forget(item)
	}
	return true
}

// checkStandbyAndSetFreeCache returns true when the Pod should be checked again.
func (t *tidbManager) checkStandbyAndSetFreeCache(ctx context.Context, item standbyCheck) bool {
	logger := t.WithName("standbyCheck").WithValues("pod name", item.podName, "pod ip", item.podIP)
	if !t.retryStandbyCheck(item) {
		return false
	}

	pod := &corev1.Pod{}
	if err := t.client.Get(ctx, client.ObjectKey{Name: item.podName, Namespace: t.tidbPoolNamespace}, pod); err != nil {
		if apierrors.IsNotFound(err) {
			t.finishStandbyCheck(item, standbyCheckResultPodGone)
			return false
		}
		if ctx.Err() != nil {
			t.finishStandbyCheck(item, standbyCheckResultCanceled)
			return false
		}
		logger.Error(err, "get pod failed")
		return t.retryStandbyCheck(item)
	}

	expectedUID := pod.UID
	if pod.Status.PodIP != item.podIP {
		t.finishStandbyCheck(item, standbyCheckResultReplaced)
		return false
	}
	if pod.DeletionTimestamp != nil ||
		pod.Status.Phase == corev1.PodSucceeded ||
		pod.Status.Phase == corev1.PodFailed {
		t.finishStandbyCheck(item, standbyCheckResultPodGone)
		return false
	}

	endpoint := ep.GenEndpointFromPod(item.podName, item.podIP)
	if token := t.assignCache.findTokenByEndpoint(endpoint); token != "" {
		t.finishStandbyCheck(item, standbyCheckResultAssigned)
		logger.Info("pod has already been assigned", "token", token)
		return false
	}

	if !kubeapi.IsPodReady(pod) {
		return t.retryStandbyCheck(item)
	}

	statusCtx, cancel := context.WithTimeout(ctx, standbyCheckTimeout)
	status, err := t.tidbController.GetStatus(
		statusCtx,
		net.JoinHostPort(item.podIP, strconv.Itoa(standbyCheckStatusPort)),
	)
	cancel()
	if err != nil {
		if ctx.Err() != nil {
			t.finishStandbyCheck(item, standbyCheckResultCanceled)
			return false
		}
		logger.Error(err, "get tidb status failed")
		return t.retryStandbyCheck(item)
	}

	if status.IsActivated() {
		t.finishStandbyCheck(item, standbyCheckResultActivated)
		logger.Info("pod has already been activated", "keyspace name", status.KeyspaceName)
		return false
	}
	if !status.IsStandby() || status.KeyspaceName != "" {
		logger.Info("pod is not standby yet", "state", status.State, "keyspace name", status.KeyspaceName)
		return t.retryStandbyCheck(item)
	}

	return t.finishStandbyCheckWithCurrentPod(ctx, item, expectedUID)
}

func (t *tidbManager) retryStandbyCheck(item standbyCheck) bool {
	t.standbyCheckMu.Lock()
	defer t.standbyCheckMu.Unlock()

	return t.retryStandbyCheckLocked(item)
}

// retryStandbyCheckLocked requires standbyCheckMu to be held.
func (t *tidbManager) retryStandbyCheckLocked(item standbyCheck) bool {
	if _, quarantined := t.standbyChecks[item]; !quarantined {
		return false
	}
	t.freeCache.quarantinePodIfIP(item.podName, item.podIP)
	return true
}

func (t *tidbManager) finishStandbyCheck(item standbyCheck, result string) {
	t.standbyCheckMu.Lock()
	defer t.standbyCheckMu.Unlock()

	t.finishStandbyCheckLocked(item, result)
}

// finishStandbyCheckLocked requires standbyCheckMu to be held.
func (t *tidbManager) finishStandbyCheckLocked(item standbyCheck, result string) {
	if !t.removeStandbyCheckLocked(item) {
		return
	}
	t.freeCache.removePodIfIP(item.podName, item.podIP)
	t.recordStandbyCheckCompleted(result)
}

// removeStandbyCheckLocked requires standbyCheckMu to be held.
func (t *tidbManager) removeStandbyCheckLocked(item standbyCheck) bool {
	if _, quarantined := t.standbyChecks[item]; !quarantined {
		return false
	}
	delete(t.standbyChecks, item)
	t.syncStandbyQuarantinedPodsMetricLocked()
	return true
}

// syncStandbyQuarantinedPodsMetricLocked requires standbyCheckMu to be held.
func (t *tidbManager) syncStandbyQuarantinedPodsMetricLocked() {
	metrics.StandbyQuarantinedPodsGauge.WithLabelValues(t.tidbPoolNamespace).Set(float64(len(t.standbyChecks)))
}

func (t *tidbManager) recordStandbyCheckCompleted(result string) {
	metrics.StandbyCheckCompletedCounter.WithLabelValues(t.tidbPoolNamespace, result).Inc()
}

// finishStandbyCheckWithCurrentPod performs the final Pod read, validation, and
// free-cache mutation under the same lock used by SyncCache.
func (t *tidbManager) finishStandbyCheckWithCurrentPod(
	ctx context.Context,
	item standbyCheck,
	expectedUID k8stypes.UID,
) bool {
	logger := t.WithName("standbyCheck").WithValues("pod name", item.podName, "pod ip", item.podIP)
	t.standbyCheckMu.Lock()
	defer t.standbyCheckMu.Unlock()

	if _, quarantined := t.standbyChecks[item]; !quarantined {
		return false
	}

	currentPod := &corev1.Pod{}
	if err := t.client.Get(ctx, client.ObjectKey{Name: item.podName, Namespace: t.tidbPoolNamespace}, currentPod); err != nil {
		if apierrors.IsNotFound(err) {
			t.finishStandbyCheckLocked(item, standbyCheckResultPodGone)
			return false
		}
		if ctx.Err() != nil {
			t.finishStandbyCheckLocked(item, standbyCheckResultCanceled)
			return false
		}
		logger.Error(err, "get pod after standby check failed")
		return t.retryStandbyCheckLocked(item)
	}

	if currentPod.Status.PodIP != item.podIP {
		t.finishStandbyCheckLocked(item, standbyCheckResultReplaced)
		logger.Info("pod was replaced during standby check",
			"expected uid", expectedUID, "pod uid", currentPod.UID, "pod ip", currentPod.Status.PodIP)
		return false
	}
	if currentPod.UID != expectedUID {
		logger.Info("pod was replaced with the same ip during standby check",
			"expected uid", expectedUID, "pod uid", currentPod.UID)
		return t.retryStandbyCheckLocked(item)
	}
	if currentPod.DeletionTimestamp != nil ||
		currentPod.Status.Phase == corev1.PodSucceeded ||
		currentPod.Status.Phase == corev1.PodFailed {
		t.finishStandbyCheckLocked(item, standbyCheckResultPodGone)
		return false
	}
	if !kubeapi.IsPodReady(currentPod) {
		return t.retryStandbyCheckLocked(item)
	}

	endpoint := ep.GenEndpointFromPod(item.podName, item.podIP)
	if token := t.assignCache.findTokenByEndpoint(endpoint); token != "" {
		t.finishStandbyCheckLocked(item, standbyCheckResultAssigned)
		logger.Info("pod was assigned during standby check", "token", token)
		return false
	}

	if !t.removeStandbyCheckLocked(item) {
		return false
	}
	if t.restorePodToFreeCacheIfEligibleLocked(currentPod, item.podIP) {
		t.recordStandbyCheckCompleted(standbyCheckResultRestored)
		logger.Info("[pod lifecycle] standby pod returned to free cache")
	} else if t.assignCache.findTokenByEndpoint(endpoint) != "" {
		t.recordStandbyCheckCompleted(standbyCheckResultAssigned)
	} else {
		t.recordStandbyCheckCompleted(standbyCheckResultIneligible)
	}
	return false
}