package server

import (
	"context"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	clusteragentv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func (s *serverImpl) prepareTiProxyGracefulRestartForTargets(ctx context.Context, sharedPoolId string, targetComponents []*entity.TargetComponent, mode entity.SetVersionConfigMode) (map[string][]string, error) {
	if mode == apiv1.SetSharedPoolComponentConfigVersionRequest_HOT_RELOAD {
		return map[string][]string{}, nil
	}

	var sharedPool *entity.SharedPool
	legacyPodsByGroup := make(map[string][]string)
	for _, comp := range targetComponents {
		if comp.GetComponentType() != apiv1.ComponentType_TIPROXY {
			continue
		}
		tiproxyGroupId := comp.GetId()
		if tiproxyGroupId == "" {
			if sharedPool == nil {
				var err error
				sharedPool, err = s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
				if err != nil {
					return nil, status.Errorf(codes.Internal, "failed to load sharedPool %s for tiproxy graceful restart preflight: %s", sharedPoolId, err)
				}
			}
			tiproxyGroupId = sharedPool.GetDefaultComponenId()
		}
		if _, ok := legacyPodsByGroup[tiproxyGroupId]; ok {
			continue
		}

		if sharedPoolId == "" {
			return nil, status.Error(codes.InvalidArgument, "sharedPool id is required")
		}
		if tiproxyGroupId == "" {
			return nil, status.Error(codes.InvalidArgument, "tiproxy group id is required")
		}

		actionCtx, err := s.prepareTiProxyBackendActionContext(ctx, sharedPoolId, tiproxyGroupId)
		if err != nil {
			return nil, err
		}
		// This preflight writes Kubernetes/Dataplane before the transaction, so
		// guard here even though DAG admission rechecks later.
		if err := guardSharedPoolMutation(actionCtx.sharedPool, sharedPoolMutationBlockedWhilePaused); err != nil {
			return nil, err
		}
		if err := s.ensureTiProxyGracefulRestartConfigured(ctx, actionCtx); err != nil {
			return nil, err
		}
		pods, err := listTiProxyPods(ctx, actionCtx.remoteCli, actionCtx.dataPlaneNamespace, actionCtx.tiproxyGroupName)
		if err != nil {
			return nil, err
		}
		if len(pods) == 0 {
			return nil, status.Errorf(codes.FailedPrecondition, "no tiproxy pod found for group %s", actionCtx.tiproxyGroupName)
		}

		podNames := make([]string, 0, len(pods))
		for i := range pods {
			podNames = append(podNames, pods[i].Name)
		}
		probeResp, err := s.clustergatewayCli.ProbeTiProxyHealthAPI(ctx, actionCtx.sharedPool, &clusteragentv1.ProbeTiProxyHealthAPIReq{
			ClusterId: actionCtx.sharedPool.GetId(),
			PodNames:  podNames,
		})
		if err != nil {
			return nil, status.Errorf(codes.Internal, "failed to probe tiproxy health API: %s", err)
		}
		legacyPodNames, err := tiProxyPodsNeedingLegacyNoNewConnections(probeResp, podNames)
		if err != nil {
			return nil, err
		}
		legacyPodsByGroup[tiproxyGroupId] = legacyPodNames
	}
	return legacyPodsByGroup, nil
}

func tiProxyPodsNeedingLegacyNoNewConnections(probeResp *clusteragentv1.ProbeTiProxyHealthAPIResp, expectedPodNames []string) ([]string, error) {
	results := probeResp.GetResults()
	if len(results) != len(expectedPodNames) {
		return nil, status.Errorf(codes.FailedPrecondition, "tiproxy health API probe returned %d results for %d pods", len(results), len(expectedPodNames))
	}

	resultByPod := make(map[string]*clusteragentv1.TiProxyHealthAPIProbeResult, len(results))
	for _, result := range results {
		if result.GetPodName() == "" {
			return nil, status.Error(codes.FailedPrecondition, "tiproxy health API probe returned empty pod name")
		}
		resultByPod[result.GetPodName()] = result
	}

	legacyPodNames := make([]string, 0, len(expectedPodNames))
	for _, podName := range expectedPodNames {
		result, ok := resultByPod[podName]
		if !ok {
			return nil, status.Errorf(codes.FailedPrecondition, "tiproxy health API probe missed pod %s", podName)
		}
		switch result.GetStatus() {
		case clusteragentv1.TiProxyHealthAPIProbeStatus_SUPPORTED:
			continue
		case clusteragentv1.TiProxyHealthAPIProbeStatus_UNSUPPORTED:
			legacyPodNames = append(legacyPodNames, podName)
		case clusteragentv1.TiProxyHealthAPIProbeStatus_UNKNOWN, clusteragentv1.TiProxyHealthAPIProbeStatus_TIPROXY_HEALTH_API_PROBE_STATUS_UNSPECIFIED:
			return nil, status.Errorf(codes.FailedPrecondition, "tiproxy health API probe for pod %s is unknown: %s", podName, result.GetMessage())
		default:
			return nil, status.Errorf(codes.FailedPrecondition, "tiproxy health API probe for pod %s returned invalid status %s", podName, result.GetStatus().String())
		}
	}
	return legacyPodNames, nil
}