package server

import (
	"context"
	"encoding/json"
	"strconv"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

func (s *serverImpl) ensureTiProxyGracefulRestartConfigured(ctx context.Context, actionCtx *tiProxyBackendActionContext) error {
	if err := patchTiProxyGroupGracefulRestart(ctx, actionCtx.remoteCli, actionCtx.dataPlaneNamespace, actionCtx.tiproxyComponent); err != nil {
		return err
	}
	lbNames, err := s.tiProxyLoadBalancerNames(ctx, actionCtx.sharedPool, actionCtx.tiproxyComponent)
	if err != nil {
		return err
	}
	for _, lbName := range lbNames {
		lb, err := s.dataplaneClient.GetLoadBalancer(ctx, &dplanemgrv1.GetLoadBalancerRequest{
			DataPlaneId: actionCtx.sharedPool.GetDataplaneName(),
			Id:          lbName,
		})
		if err != nil {
			if status.Code(err) == codes.NotFound {
				tilog.Warn(ctx, "load balancer disappeared before tiproxy graceful restart reconcile",
					zap.String("sharedPoolId", actionCtx.sharedPool.GetId()),
					zap.String("loadBalancer", lbName))
				continue
			}
			return status.Errorf(codes.Internal, "failed to get load balancer %s: %s", lbName, err)
		}
		annotations := entity.BuildTiProxyGracefulRestartLoadBalancerAnnotations(actionCtx.sharedPool, loadBalancerServicePorts(lb)...)
		if len(annotations) == 0 {
			return status.Errorf(codes.FailedPrecondition, "provider %s does not support tiproxy graceful restart load balancer annotations", actionCtx.sharedPool.GetProvider().GetType().String())
		}
		if lb.Annotations == nil {
			lb.Annotations = map[string]string{}
		}
		for key, value := range annotations {
			lb.Annotations[key] = value
		}
		lb.GracefulUnhealthyEnabled = true

		if _, err := s.dataplaneClient.UpdateLoadBalancer(ctx, &dplanemgrv1.UpdateLoadBalancerRequest{
			DataPlaneId:  actionCtx.sharedPool.GetDataplaneName(),
			Id:           lbName,
			Loadbalancer: lb,
		}); err != nil {
			return status.Errorf(codes.Internal, "failed to update load balancer %s: %s", lbName, err)
		}
	}
	return nil
}

func loadBalancerServicePorts(lb *dplanemgrv1.LoadBalancer) []int32 {
	if lb == nil {
		return nil
	}
	if ports := lb.GetPorts(); len(ports) > 0 {
		out := make([]int32, 0, len(ports))
		for _, port := range ports {
			if port.GetPort() > 0 {
				out = append(out, port.GetPort())
			}
		}
		return out
	}
	if lb.GetPort() <= 0 {
		return nil
	}
	return []int32{lb.GetPort()}
}

func patchTiProxyGroupGracefulRestart(ctx context.Context, remoteCli client.Client, namespace, tiproxyComponentId string) error {
	var tiproxyGroup v1alpha1.TiProxyGroup
	key := types.NamespacedName{Name: tiProxyGroupName(tiproxyComponentId), Namespace: namespace}
	if err := remoteCli.Get(ctx, key, &tiproxyGroup); err != nil {
		if apierrors.IsNotFound(err) {
			return status.Errorf(codes.NotFound, "tiproxy group %s/%s not found", namespace, key.Name)
		}
		return status.Errorf(codes.Internal, "failed to get tiproxy group %s/%s: %s", namespace, key.Name, err)
	}

	patch, err := json.Marshal(map[string]any{
		"spec": map[string]any{
			"template": map[string]any{
				"metadata": map[string]any{
					"annotations": map[string]any{
						entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation: strconv.FormatInt(entity.DefaultTiProxyGracefulShutdownDeleteDelaySeconds, 10),
					},
				},
			},
			"maxSurge": entity.DefaultTiProxyGracefulRestartMaxSurge,
		},
	})
	if err != nil {
		return status.Errorf(codes.Internal, "failed to build tiproxy group %s/%s graceful restart patch: %s", namespace, key.Name, err)
	}
	if err := remoteCli.Patch(ctx, &tiproxyGroup, client.RawPatch(types.MergePatchType, patch)); err != nil {
		return status.Errorf(codes.Internal, "failed to patch tiproxy group %s/%s graceful restart config: %s", namespace, key.Name, err)
	}
	return nil
}