package branch

import (
	"net"

	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/infra"
	providertypes "github.com/tidbcloud/aws-shared-provider/pkg/types"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
)

// setTidbPoolReplicas adjusts the number of TiDB replicas for a specified pool.
// It ensures the desired number of TiDB instances are running by:
//   - Setting up the replicas of specified tidbPool(must exist if replicas > 0).
//   - Prioritizing avoiding stopping the AutoID owner. It first attempts to stop non-AutoID owner instances,
//     and stops the AutoID owner at last if replicas still need to be reduced.
//   - Freeing stopped TiDB instances using the pool's manager.
//
// If the number of replicas is increased, it updates the etcd token replicas accordingly.
// If reduced, it stops and frees extra TiDB instances, handling possible errors and
// force-stopping when necessary. When replicas is set to zero, it also removes the etcd token.
func (r *reconcileContext) setTidbPoolReplicas(tidbPool string, replicas int) error {
	tManager := r.TidbManagers[tidbPool]
	if tManager == nil {
		return errors.Errorf("tManager not found in tidbPool %s", tidbPool)
	}

	if replicas > 0 {
		if err := r.setEtcdReplicas(tidbPool, replicas); err != nil {
			return errors.Annotate(err, "failed to set tidb pool in etcd")
		}
	}

	waitSeconds := 0
	if infra.IsEssentialBranch(r.branch) {
		// EssentialPlan use 8 hours to wait
		waitSeconds = 8 * 3600
	}

	endpoints := tManager.GetAssignedTidbByToken(r.keyspaceToken)
	if endpoints == nil {
		return nil
	}
	needToStop := len(endpoints.Endpoints) - replicas
	mayAutoIDOwner := make([]string, 0)
	if needToStop > 0 {
		for _, endpoint := range endpoints.Endpoints {
			podName, podIp := ep.ParsePodFromEndpoint(endpoint)
			if err := r.TidbController.Stop(r.ctx, net.JoinHostPort(podIp, "10080"), r.keyspaceToken,
				true, waitSeconds, true, false); err != nil {
				r.Error(err, "failed to stop tidb", "pod", podName, "ip", podIp)
				mayAutoIDOwner = append(mayAutoIDOwner, endpoint)
				continue
			}
			r.Info("[pod lifecycle] stop tidb success", "pod", podName, "ip", podIp)
			if err := tManager.FreeTidb(r.ctx, podName, podIp, providertypes.CommentSkipSetToFreeCache); err != nil {
				return errors.Annotatef(err, "failed to free tidb %s(%s)", podName, podIp)
			}
			needToStop--
			if needToStop <= 0 {
				break
			}
		}
	}
	if needToStop > 0 {
		if len(mayAutoIDOwner) > 1 {
			return errors.Errorf("too many mayAutoIDOwner %v", mayAutoIDOwner)
		}
		if len(mayAutoIDOwner) == 1 {
			podName, podIp := ep.ParsePodFromEndpoint(mayAutoIDOwner[0])
			if err := r.TidbController.Stop(r.ctx, net.JoinHostPort(podIp, "10080"), r.keyspaceToken,
				true, waitSeconds, false, false); err != nil {
				return errors.Annotatef(err, "failed to force stop tidb %s(%s)", podName, podIp)
			}
			r.Info("[pod lifecycle] force stop tidb success", "pod", podName, "ip", podIp)
			if err := tManager.FreeTidb(r.ctx, podName, podIp, providertypes.CommentSkipSetToFreeCache); err != nil {
				return errors.Annotatef(err, "failed to free tidb %s(%s)", podName, podIp)
			}
		}
	}

	if replicas <= 0 {
		return r.removeEtcdToken(tidbPool)
	}

	return nil
}