package action

import (
	"fmt"
	"net"
	"strings"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
)

type ApplyLoadBalanceAction struct {
	sharedPoolRepo  repo.SharedPoolRepo
	tidbClusterRepo repo.TiDBClusterRepo
	tidbGroupRepo   repo.TiDBGroupRepo
	client          dplanemgrv1.LoadBalancerServiceClient
	dataplaneClient dplanemgrv1.DataPlaneServiceClient
}

const ApplyLoadBalanceActionName = "apply-load-balance"

func (a *ApplyLoadBalanceAction) Name() string {
	return ApplyLoadBalanceActionName
}

func NewApplyLoadBalanceAction(
	sharedPoolRepo repo.SharedPoolRepo,
	tidbClusterRepo repo.TiDBClusterRepo,
	tidbGroupRepo repo.TiDBGroupRepo,
	client dplanemgrv1.LoadBalancerServiceClient,
	dataplaneClient dplanemgrv1.DataPlaneServiceClient,
) actions.Action {
	return &ApplyLoadBalanceAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbClusterRepo: tidbClusterRepo,
		tidbGroupRepo:   tidbGroupRepo,
		client:          client,
		dataplaneClient: dataplaneClient,
	}
}

func (a *ApplyLoadBalanceAction) Invoke(args actions.ActionArgs) error {
	// generate context
	ctx := utils.GenerateContext(args)
	tilog.Info(ctx, "apply load balance action")

	needUpdate, _ := args.GetString(workflow.UpdateLoadBalancerKey)
	if needUpdate != workflow.UpdateLoadBalancerValue {
		tilog.Info(ctx, "no need to update load balancer")
		return nil
	}

	sharedPoolId, err := args.GetString(workflow.SharedPoolIdKey)
	if err != nil {
		tilog.Error(ctx, "failed to get sharedPool id", zap.Error(err))
		return fmt.Errorf("failed to get sharedPool id: %w", err)
	}

	tidbGroupId, _ := args.GetString(workflow.TiDBGroupIdKey)
	if tidbGroupId == "" {
		tilog.Warn(ctx, "tidb group id is empty, skip apply load balance")
		return nil
	}

	// load tidb group
	tidbGroup, err := a.tidbGroupRepo.Load(ctx, tidbGroupId, false)
	if err != nil {
		tilog.Error(ctx, "failed to load tidb group", zap.Error(err))
		return fmt.Errorf("failed to load tidb group: %w", err)
	}

	// load sharedPool
	sharedPool, err := a.sharedPoolRepo.Load(ctx, sharedPoolId, false)
	if err != nil {
		tilog.Error(ctx, "failed to load sharedPool", zap.Error(err))
		return fmt.Errorf("failed to load sharedPool: %w", err)
	}

	// load tidb cluster
	tidbCluster, err := a.tidbClusterRepo.Load(ctx, tidbGroup.GetClusterId(), false)
	if err != nil {
		tilog.Error(ctx, "failed to load tidb cluster", zap.Error(err))
		return fmt.Errorf("failed to load tidb cluster: %w", err)
	}

	var tempTidbGroup entity.TiDBGroup
	// get endpoint config from args
	tidbGroupStr, err := args.GetString(workflow.TiDBGroupKey)
	if err != nil {
		tilog.Warn(ctx, "failed to get tidb group from args, will delete tidb group", zap.Error(err))
		tempTidbGroup = entity.TiDBGroup{
			TiDBGroup: apiv1.TiDBGroup{
				Id: tidbGroup.GetId(),
				SqlEndpoint: &entity.SQLEndpoint{
					PrivateLb: &entity.LoadBalancer{
						Enabled: false,
					},
					PublicLb: &entity.LoadBalancer{
						Enabled: false,
					},
				},
			},
		}
	} else {
		err = protojson.Unmarshal([]byte(tidbGroupStr), &tempTidbGroup)
		if err != nil {
			tilog.Error(ctx, "failed to unmarshal tidb group", zap.Error(err))
			return fmt.Errorf("failed to unmarshal tidb group: %w", err)
		}
	}

	tidbGroup.SqlEndpoint = tempTidbGroup.SqlEndpoint

	if tidbGroup.GetSqlEndpoint().GetPrivateLb().GetEnabled() {
		dataplane, err := a.dataplaneClient.GetDataPlane(ctx, &dplanemgrv1.GetDataPlaneRequest{
			DataPlaneId: tidbGroup.GetDataplaneName(),
		})
		if err != nil {
			return fmt.Errorf("failed to get dataplane %s: %w", tidbGroup.GetDataplaneName(), err)
		}
		normalizedCIDRs, err := normalizePrivateLBCIDRs(
			tidbGroup.GetSqlEndpoint().GetPrivateLb().GetIpAllowList(),
			dataplane.GetVpc().GetCidr(),
		)
		if err != nil {
			return err
		}
		tidbGroup.GetSqlEndpoint().GetPrivateLb().IpAllowList = normalizedCIDRs
	}

	lbComp := entity.FindTiProxyComponent(sharedPool, tidbGroup.GetClusterId())

	if lbComp == nil {
		tilog.Error(ctx, "lb component not found",
			zap.String("tidbGroupId", tidbGroup.GetId()),
			zap.String("sharedPoolId", sharedPool.GetId()))
		return fmt.Errorf("tiproxy component not found for tidb group %s", tidbGroup.GetId())
	}

	if tidbGroup.GetSqlEndpoint().GetPublicLb().GetEnabled() {
		// first check if the load balancer already exists
		_, err := a.client.GetLoadBalancer(ctx, &dplanemgrv1.GetLoadBalancerRequest{
			DataPlaneId: tidbGroup.GetDataplaneName(),
			Id:          entity.GetPublicSvcName(tidbGroup.GetId()),
		})
		if err != nil {
			if status.Code(err) == codes.NotFound {
				tilog.Info(ctx, "public load balancer not found, will create it",
					zap.String("tidbGroupId", tidbGroup.GetId()),
					zap.String("sharedPoolId", sharedPool.GetId()),
				)

				_, err := a.client.CreateLoadBalancer(ctx, &dplanemgrv1.CreateLoadBalancerRequest{
					DataPlaneId:  tidbGroup.GetDataplaneName(),
					Id:           entity.GetPublicSvcName(tidbGroup.GetId()),
					Loadbalancer: buildPublicLoadBalancer(tidbCluster, tidbGroup, lbComp, sharedPool),
				})
				if err != nil {
					tilog.Error(ctx, "failed to create public load balancer",
						zap.String("tidbGroupId", tidbGroup.GetId()),
						zap.Error(err),
					)
					return fmt.Errorf("failed to create public load balancer: %w", err)
				}
			} else {
				tilog.Error(ctx, "failed to get public load balancer",
					zap.String("tidbGroupId", tidbGroup.GetId()),
					zap.Error(err))
				return fmt.Errorf("failed to get public load balancer: %w", err)
			}
		} else {
			// update the existing load balancer
			_, err = a.client.UpdateLoadBalancer(ctx, &dplanemgrv1.UpdateLoadBalancerRequest{
				DataPlaneId:  tidbGroup.GetDataplaneName(),
				Id:           entity.GetPublicSvcName(tidbGroup.GetId()),
				Loadbalancer: buildPublicLoadBalancer(tidbCluster, tidbGroup, lbComp, sharedPool),
			})
			if err != nil {
				tilog.Error(ctx, "failed to update public load balancer",
					zap.String("tidbGroupId", tidbGroup.GetId()),
					zap.Error(err),
				)
				return fmt.Errorf("failed to update public load balancer: %w", err)
			}
		}

	} else {
		_, err = a.client.DeleteLoadBalancer(ctx, &dplanemgrv1.DeleteLoadBalancerRequest{
			DataPlaneId: tidbGroup.GetDataplaneName(),
			Id:          entity.GetPublicSvcName(tidbGroup.GetId()),
		})
		if err != nil {
			tilog.Error(ctx, "failed to delete public load balancer",
				zap.String("tidbGroupId", tidbGroup.GetId()),
				zap.Error(err),
			)
			return fmt.Errorf("failed to delete public load balancer: %w", err)
		}
	}

	if tidbGroup.GetSqlEndpoint().GetPrivateLb().GetEnabled() {
		// check if the load balancer already exists
		_, err := a.client.GetLoadBalancer(ctx, &dplanemgrv1.GetLoadBalancerRequest{
			DataPlaneId: tidbGroup.GetDataplaneName(),
			Id:          entity.GetPrivateSvcName(tidbGroup.GetId()),
		})

		if err != nil {
			if status.Code(err) == codes.NotFound {
				tilog.Info(ctx, "private load balancer not found, will create it",
					zap.String("tidbGroupId", tidbGroup.GetId()),
					zap.String("sharedPoolId", sharedPool.GetId()),
				)
				_, err = a.client.CreateLoadBalancer(ctx, &dplanemgrv1.CreateLoadBalancerRequest{
					DataPlaneId:  tidbGroup.GetDataplaneName(),
					Id:           entity.GetPrivateSvcName(tidbGroup.GetId()),
					Loadbalancer: buildPrivateLoadBalancer(tidbCluster, tidbGroup, lbComp, sharedPool),
				})
				if err != nil {
					tilog.Error(ctx, "failed to create private load balancer",
						zap.String("tidbGroupId", tidbGroup.GetId()),
						zap.Error(err),
					)
					return fmt.Errorf("failed to create private load balancer: %w", err)
				}
			} else {
				tilog.Error(ctx, "failed to get private load balancer",
					zap.String("tidbGroupId", tidbGroup.GetId()),
					zap.Error(err))
				return fmt.Errorf("failed to get private load balancer: %w", err)
			}
		} else {
			_, err = a.client.UpdateLoadBalancer(ctx, &dplanemgrv1.UpdateLoadBalancerRequest{
				DataPlaneId:  tidbGroup.GetDataplaneName(),
				Id:           entity.GetPrivateSvcName(tidbGroup.GetId()),
				Loadbalancer: buildPrivateLoadBalancer(tidbCluster, tidbGroup, lbComp, sharedPool),
			})
			if err != nil {
				tilog.Error(ctx, "failed to update private load balancer",
					zap.String("tidbGroupId", tidbGroup.GetId()),
					zap.Error(err),
				)
				return fmt.Errorf("failed to update private load balancer: %w", err)
			}
		}
	} else {
		_, err = a.client.DeleteLoadBalancer(ctx, &dplanemgrv1.DeleteLoadBalancerRequest{
			DataPlaneId: tidbGroup.GetDataplaneName(),
			Id:          entity.GetPrivateSvcName(tidbGroup.GetId()),
		})
		if err != nil {
			tilog.Error(ctx, "failed to delete private load balancer",
				zap.String("tidbGroupId", tidbGroup.GetId()),
				zap.Error(err),
			)
			return fmt.Errorf("failed to delete private load balancer: %w", err)
		}
	}

	tilog.Info(ctx, "successfully applied load balancer for tidb group",
		zap.String("tidbGroupId", tidbGroup.GetId()),
		zap.String("sharedPoolId", sharedPool.GetId()))

	return nil
}

func buildPrivateLoadBalancer(tidbCluster *entity.TiDBCluster, tidbGroup *entity.TiDBGroup, lbComp *entity.Component, sharedPool *entity.SharedPool) *dplanemgrv1.LoadBalancer {

	// old tidbcluster not have identity, use the tidbgroup identity instead
	subIdentity := tidbGroup.GetSqlEndpoint().GetDnsIdentity()
	identity := tidbCluster.GetGatewayIdentity()
	if identity == "" {
		identity = tidbGroup.GetSqlEndpoint().GetDnsIdentity()
		subIdentity = ""
	}

	privateLBPort := tidbGroup.GetSqlEndpoint().GetPrivateLb().GetPort()
	annotations := entity.BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool, privateLBPort)
	gracefulRestartEnabled := len(annotations) > 0
	return &dplanemgrv1.LoadBalancer{
		Id:                       entity.GetPrivateSvcName(tidbGroup.GetId()),
		Namespace:                entity.GetDataPlaneNamespace(tidbGroup.GetSharedPoolId()),
		Identity:                 identity,
		SubIdentity:              subIdentity,
		Type:                     dplanemgrv1.LoadBalancer_PRIVATE,
		Port:                     privateLBPort,
		Selectors:                buildLoadBalancerSelectors(lbComp),
		Tags:                     entity.BuildTiDBGroupTags(tidbGroup),
		IpWhiteList:              buildIpWhiteList(tidbGroup.GetSqlEndpoint().GetPrivateLb()),
		Annotations:              annotations,
		GracefulUnhealthyEnabled: gracefulRestartEnabled,
	}
}

func buildPublicLoadBalancer(tidbCluster *entity.TiDBCluster, tidbGroup *entity.TiDBGroup, lbComp *entity.Component, sharedPool *entity.SharedPool) *dplanemgrv1.LoadBalancer {

	// old tidbcluster not have identity, use the tidbgroup identity instead
	subIdentity := tidbGroup.GetSqlEndpoint().GetDnsIdentity()
	identity := tidbCluster.GetGatewayIdentity()
	if identity == "" {
		identity = tidbGroup.GetSqlEndpoint().GetDnsIdentity()
		subIdentity = ""
	}
	if tidbGroup.GetSqlEndpoint().GetDisablePublicDnsIdentity() {
		subIdentity = ""
	}

	publicLBPort := tidbGroup.GetSqlEndpoint().GetPublicLb().GetPort()
	annotations := entity.BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool, publicLBPort)
	gracefulRestartEnabled := len(annotations) > 0
	lb := &dplanemgrv1.LoadBalancer{
		Id:                       entity.GetPublicSvcName(tidbGroup.GetId()),
		Namespace:                entity.GetDataPlaneNamespace(tidbGroup.GetSharedPoolId()),
		Identity:                 identity,
		SubIdentity:              subIdentity,
		Type:                     dplanemgrv1.LoadBalancer_PUBLIC,
		Port:                     publicLBPort,
		Selectors:                buildLoadBalancerSelectors(lbComp),
		Tags:                     entity.BuildTiDBGroupTags(tidbGroup),
		IpWhiteList:              buildIpWhiteList(tidbGroup.GetSqlEndpoint().GetPublicLb()),
		Annotations:              annotations,
		GracefulUnhealthyEnabled: gracefulRestartEnabled,
	}
	if entity.HasSharedTiProxy(sharedPool) {
		lb.PreserveClientIp = new(bool)
	}
	return lb
}

func buildLoadBalancerSelectors(lbComp *entity.Component) map[string]string {
	selectors := map[string]string{
		"app.kubernetes.io/component": entity.GetComponentTypeName(lbComp),
		"app.kubernetes.io/instance":  entity.GetComponentGroupName(lbComp),
		"pingcap.com/group":           entity.GetComponentGroupName(lbComp),
	}
	return selectors
}

func buildIpWhiteList(lb *entity.LoadBalancer) []string {

	iplist := make([]string, 0, len(lb.GetIpAllowList()))
	for _, ip := range lb.GetIpAllowList() {
		cidr := ip.GetCidrBlock()
		if cidr == "" {
			continue
		}
		iplist = append(iplist, cidr)
	}
	return iplist
}

func normalizePrivateLBCIDRs(userCIDRs []*entity.IpAllowList, dataplaneCIDR string) ([]*entity.IpAllowList, error) {
	dataplaneCIDR = strings.TrimSpace(dataplaneCIDR)
	if dataplaneCIDR == "" {
		return nil, fmt.Errorf("dataplane vpc cidr is empty")
	}
	if _, _, err := net.ParseCIDR(dataplaneCIDR); err != nil {
		return nil, fmt.Errorf("dataplane vpc cidr %s is invalid: %w", dataplaneCIDR, err)
	}

	seen := make(map[string]struct{}, len(userCIDRs)+1)
	result := make([]*entity.IpAllowList, 0, len(userCIDRs)+1)

	appendCIDR := func(cidr string) {
		cidr = strings.TrimSpace(cidr)
		if cidr == "" {
			return
		}
		if _, ok := seen[cidr]; ok {
			return
		}
		seen[cidr] = struct{}{}
		result = append(result, &entity.IpAllowList{CidrBlock: cidr})
	}

	for _, item := range userCIDRs {
		appendCIDR(item.GetCidrBlock())
	}
	appendCIDR(dataplaneCIDR)

	return result, nil
}