package tidbmanager

import (
	"context"
	"time"

	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	"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"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

// ReplaceTidb replaces the old TiDB pod with a new one.
// It first activates a new TiDB instance then kills the old one.
func (t *tidbManager) ReplaceTidb(ctx context.Context, podName string, deleteFn func(ctx context.Context, podName, namespace string) error) error {
	podIP, err := t.getPodIP(ctx, podName)
	if err != nil {
		return err
	}
	// Get pod's token
	endpoint := ep.GenEndpointFromPod(podName, podIP)
	token := t.assignCache.findTokenByEndpoint(endpoint)
	if token == "" {
		return errors.Annotatef(utils.ErrNotFoundToken, "endpoint not found in assign cache: %s", endpoint)
	}

	// activate new pod
	tInfo, ok := t.assignCache.getTidbInfo(token)
	if !ok {
		return errors.Annotatef(utils.ErrNotFoundToken, "token not found in assign cache: %s", token)
	}

	start := time.Now()
	defer func() {
		stage := "total"
		if err != nil {
			stage = "total_fail"
		}
		metrics.Reclaimer.GetHistogramVec(metrics.ReplaceTidbHistogramName, token).WithLabelValues(stage, t.tidbPoolNamespace).
			Observe(time.Since(start).Seconds())
	}()

	param, err := t.getActivateParameter(ctx, token)
	if err != nil {
		return err
	}

	var endpoints []string
	endpoints, err = t.assignWithLock(ctx, token, len(tInfo.endpoints)+1, false, param)
	if err != nil {
		return err
	}
	t.Logger.Info("replace tidb success", "podName", podName, "endpoints", endpoints)

	// NOTE: wait for new backend to be synced.
	//nolint: gomnd // not need
	time.Sleep(time.Second * 5)

	// delete old pod
	err = deleteFn(ctx, podName, t.tidbPoolNamespace)
	if err != nil {
		return err
	}

	t.tidbScaler.TryScale()

	ticker := time.NewTicker(time.Second * 3)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			err = ctx.Err()
			return err
		case <-ticker.C:
			currentToken := t.assignCache.findTokenByEndpoint(endpoint)
			if currentToken == "" || currentToken != token {
				return err
			}
		}
	}
}

func (t *tidbManager) getActivateParameter(ctx context.Context, token string) (*tidb.ActivateParameter, error) {
	regionalInfra := infra.NewRegionalInfraWithClient(t.infraClusterCli, t.cfg.KubeClusterName)
	var meta *infra.ClusterOrBranchMeta
	meta, err := infra.GetClusterOrBranchMeta(ctx, regionalInfra, t.infraApi, utils.TrimToken(token), t.cfg.KubeClusterName)
	if err != nil {
		return nil, err
	}

	var maxIdleSeconds uint = 0 // use pool default value
	// RCU>0 means essential plan, rewrite maxIdleSeconds with the highest priority.
	if meta.IsEssential() {
		maxIdleSeconds = uint(types.TidbMaxIdleSeconds_EssentialPlan)
	}

	return &tidb.ActivateParameter{
		EnableAuditLog:   meta.EnableAuditLog,
		AuditLogRedacted: meta.AuditLogRedacted,
		MaxIdleSeconds:   maxIdleSeconds,
		KeyspaceID:       meta.KeyspaceID,
		ProjectID:        meta.ProjectID,
		ClusterID:        meta.ClusterID,
		TenantID:         meta.TenantID,
	}, nil
}

func (t *tidbManager) getPodIP(ctx context.Context, podName string) (string, error) {
	pod := &corev1.Pod{}
	err := t.client.Get(ctx, client.ObjectKey{Name: podName, Namespace: t.tidbPoolNamespace}, pod)
	if err != nil {
		return "", err
	}
	if pod.Status.PodIP == "" {
		return "", errors.New("pod IP is empty")
	}
	return pod.Status.PodIP, nil
}