package node_controller

import (
	"context"
	"math/rand"
	"sort"
	"time"

	"github.com/go-logr/logr"
	"github.com/pingcap/errors"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/mixeddeploypd"
	"github.com/tidbcloud/infra-api/controllers/utils"
	"github.com/tidbcloud/infra-api/pkg/consts"
	"github.com/tidbcloud/infra-api/pkg/gardener"
	"github.com/tidbcloud/infra-api/pkg/gitops"
	"github.com/tidbcloud/infra-api/pkg/k8s"
	metricsctrl "github.com/tidbcloud/infra-api/pkg/metrics/controller"
	"github.com/tidbcloud/infra-api/pkg/nodegroup"
)

const (
	replicaFactor = 3

	defaultRequeueInterval = 30 * time.Second
)

type ClusterNodeScaler struct {
	ConfigProvider gitops.ConfigProvider
	InfraClient    client.Client

	log         logr.Logger
	shootClient client.Client
	shootName   string
	gardener    gardener.GardenerAPI
}

func (c *ClusterNodeScaler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	start := time.Now()
	defer func() {
		metricsctrl.ReconcileDuration.WithLabelValues("node_scaler").Observe(time.Since(start).Seconds())
	}()
	c.log.Info("check scaling tidb cluster", "tc", req.NamespacedName)
	tc := &v1alpha1.TidbCluster{}
	err := c.shootClient.Get(ctx, req.NamespacedName, tc)
	if err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}
	if tc.Annotations[nodegroup.NodeScalerIgnore] == "true" {
		return ctrl.Result{}, nil
	}
	onlyCdc := tc.Spec.TiCDC != nil && tc.Spec.TiDB == nil && tc.Spec.TiFlash == nil && tc.Spec.TiKV == nil
	if onlyCdc {
		if err = c.ensureFinalizer(ctx, tc); err != nil {
			return ctrl.Result{}, errors.AddStack(err)
		}
	}

	clusterNamespace := tc.Labels[consts.ClusterNamespaceLabelKey]
	clusterName := tc.Labels[consts.ClusterNameLabelKey]
	cdcName := tc.Labels[consts.ClusterTicdcLabelKey]
	if clusterNamespace == "" {
		return ctrl.Result{}, errors.New("cluster namespace is empty in tc label keys")
	}
	var cluster *cloudv1.Cluster
	if clusterName != "" {
		cluster = &cloudv1.Cluster{}
		err = c.InfraClient.Get(ctx, types.NamespacedName{Namespace: clusterNamespace, Name: clusterName}, cluster)
		// should ignore NotFound error?
		if err != nil {
			return ctrl.Result{}, errors.AddStack(err)
		}
		// cluster is deleted, skip node group scaling
		if cluster.DeletionTimestamp != nil {
			return ctrl.Result{}, nil
		}
	}
	var ticdc *cloudv1.TiCDC
	// if cdc is not nil, add it for scale nodes
	if cdcName != "" {
		ticdc = &cloudv1.TiCDC{}
		// check nodeGroup whether ticdc is deleted or not
		err = c.InfraClient.Get(ctx, types.NamespacedName{Namespace: clusterNamespace, Name: cdcName}, ticdc)
		if err != nil {
			if !apierrors.IsNotFound(err) {
				return ctrl.Result{}, errors.Annotatef(err, "failed to get cdc resource")
			} else {
				// cloudv1.TiCDC is deleted
				ticdc = &cloudv1.TiCDC{
					ObjectMeta: metav1.ObjectMeta{
						Namespace: clusterNamespace,
						Name:      cdcName,
					},
					Spec: cloudv1.TiCDCSpec{
						NetworkName:         tc.Labels[consts.ClusterNetworkLabelKey],
						UpstreamClusterName: tc.Labels[consts.UpstreamClusterNameLabelKey],
						ClusterName:         tc.Labels[consts.ClusterNameLabelKey],
						Replicas:            0,
					},
				}
			}
		} else if !ticdc.DeletionTimestamp.IsZero() {
			ticdc.Spec.Replicas = 0
		}
	}
	network := &cloudv1.Network{}
	var networkRef types.NamespacedName
	if cluster != nil {
		networkRef = cluster.GetNetworkRef()
	} else {
		networkRef = ticdc.GetNetworkRef()
	}
	err = c.InfraClient.Get(ctx, networkRef, network)
	if err != nil {
		return ctrl.Result{}, errors.Annotate(err, "get network failed")
	}
	// only update workers when shoot name match
	status, err := cloudv1.GardenerNetworkStatusFromNetwork(network)
	if err != nil {
		return ctrl.Result{}, errors.Annotate(err, "get network status failed")
	}
	if status.ShootName != c.shootName {
		return ctrl.Result{}, nil
	}
	provider := &cloudv1.Provider{}
	err = c.InfraClient.Get(ctx, types.NamespacedName{Name: network.Spec.Provider}, provider)
	if err != nil {
		return ctrl.Result{}, err
	}
	sctx, err := c.gardener.GetShootContext(ctx, c.shootName)
	if err != nil {
		return ctrl.Result{}, err
	}
	scalingCtx := &scalingContext{
		ClusterNodeScaler:     c,
		provider:              provider,
		network:               network,
		cluster:               cluster,
		tc:                    tc,
		cdc:                   ticdc,
		ctx:                   ctx,
		onlyCdc:               onlyCdc,
		tiflashSingleAZDeploy: sctx.TiFlashSingleAZDeploy,
		PDMixedDeployConfig:   sctx.PDMixedDeployConfig,
		scheduleNodePreferPVC: sctx.ScheduleNodePreferPVC,
	}
	result, err := scalingCtx.scale()
	if err != nil || result.RequeueAfter > 0 {
		return result, err
	}

	if tc.DeletionTimestamp != nil {
		if err = c.removeFinalizer(ctx, tc); err != nil {
			return ctrl.Result{}, err
		}
	}
	return ctrl.Result{}, nil
}

func (c *ClusterNodeScaler) ensureFinalizer(ctx context.Context, tc *v1alpha1.TidbCluster) error {
	if tc.DeletionTimestamp.IsZero() && !controllerutil.ContainsFinalizer(tc, nodegroup.NodeScalerFinalizer) {
		controllerutil.AddFinalizer(tc, nodegroup.NodeScalerFinalizer)
		return c.shootClient.Update(ctx, tc)
	}
	return nil
}

func (c *ClusterNodeScaler) removeFinalizer(ctx context.Context, tc *v1alpha1.TidbCluster) error {
	if controllerutil.ContainsFinalizer(tc, nodegroup.NodeScalerFinalizer) {
		controllerutil.RemoveFinalizer(tc, nodegroup.NodeScalerFinalizer)
		if tc.Annotations == nil {
			tc.Annotations = map[string]string{}
		}
		tc.Annotations[nodegroup.NodeScalerIgnore] = "true"
		return c.shootClient.Update(ctx, tc)
	}
	return nil
}

type scalingContext struct {
	*ClusterNodeScaler

	provider *cloudv1.Provider
	network  *cloudv1.Network
	// cluster will be nil if cdc is deployed in downstream
	// and downstream has no infra-api Cluster
	cluster *cloudv1.Cluster
	cdc     *cloudv1.TiCDC
	tc      *v1alpha1.TidbCluster
	ctx     context.Context
	onlyCdc bool

	tiflashSingleAZDeploy bool
	PDMixedDeployConfig   *cloudv1.MixedDeployPDWithTiKV
	scheduleNodePreferPVC bool
}

type scaledComponent struct {
	nodeCount int

	cloudv1.ClusterComponent
}

func (c *scaledComponent) GetReplicas() int {
	return c.nodeCount
}

func scaleComponentTo(comp cloudv1.ClusterComponent, nodeCount int) cloudv1.ClusterComponent {
	return &scaledComponent{
		nodeCount:        nodeCount,
		ClusterComponent: comp,
	}
}

// scale update the cluster node groups according to the requirement of cluster and tidbcluster:
// a. the node template is specified by cluster spec
// b. the replicas of each node group is calculated from the current state of tidbcluster resource
func (c *scalingContext) scale() (ctrl.Result, error) {

	// 1. calculate deisred node count of each component
	var components []cloudv1.ClusterComponent
	var componentTemplates []cloudv1.ClusterComponent
	if !c.onlyCdc {
		componentTemplates = c.cluster.GetComponents()
		componentTemplates = append(componentTemplates, &AdminTiDB{TiDB: &c.cluster.Spec.Cluster.TiDB})
	}
	if c.cdc != nil {
		componentTemplates = append(componentTemplates, c.cdc)
	}
	currentSpreads := map[cloudv1.ComponentType][]nodegroup.CountPerZone{}
	currentPVCSpreads := map[cloudv1.ComponentType][]nodegroup.CountPerZone{}
	var mixedDeployPDWithTiKV bool
	var useZeroPDNodesNowForMixedDeployPDWithTiKV bool
	for _, template := range componentTemplates {
		componentType := template.GetComponentType()
		memberType := utils.GetMemberType(componentType)
		// If one component is suspended, skip the nodegroup building
		if c.tc.ComponentIsSuspended(memberType) {
			c.log.Info("Component is suspended, skip building nodegroup", "component", componentType)
			continue
		}
		// calculate the real node replicas we need
		var desiredNodeCount int
		var err error
		switch componentType {
		case cloudv1.ComponentTiDB:
			desiredNodeCount = c.tidbDesiredNodeCount()
		case cloudv1.ComponentTiFlash:
			desiredNodeCount = c.tiflashDesiredNodeCount()
		case cloudv1.ComponentTiKV:
			desiredNodeCount = c.tikvDesiredNodeCount()
		case cloudv1.ComponentTiProxy:
			desiredNodeCount = c.tiproxyDesiredNodeCount()
		case cloudv1.ComponentPD:
			desiredNodeCount, mixedDeployPDWithTiKV, err = c.pdDesiredNodeCount()
			if err != nil {
				return ctrl.Result{}, err
			}
			if desiredNodeCount == 0 {
				useZeroPDNodesNowForMixedDeployPDWithTiKV = true
			}
		case cloudv1.ComponentTiCDC:
			desiredNodeCount = c.ticdcDesiredNodeCount()
		default:
			// other components required by cluster are passed in as is
			desiredNodeCount = template.GetReplicas()
		}
		if desiredNodeCount < template.GetReplicas() && !mixedDeployPDWithTiKV {
			return ctrl.Result{}, errors.Errorf("%s: desired node count %d is less then component replica count %d",
				componentType, desiredNodeCount, template.GetReplicas())
		}
		// scale the component to the real desired replicas
		components = append(components, scaleComponentTo(template, desiredNodeCount))

		// calculate current pod distribution to guide the node scheduling
		countPerZone, err := c.componentCountPerZone(componentType)
		if err != nil {
			return ctrl.Result{}, errors.Annotate(err, "GetComponentCountPerZone")
		}
		if countPerZone != nil {
			currentSpreads[componentType] = countPerZone
		}

		// calculate current pvc distribution to guide the node scheduling
		if c.scheduleNodePreferPVC {
			pvcCountPerZone, err := c.componentPVCCountPerZone(componentType)
			if err != nil {
				return ctrl.Result{}, errors.Annotate(err, "GetComponentPVCCountPerZone")
			}
			if pvcCountPerZone != nil {
				currentPVCSpreads[componentType] = pvcCountPerZone
			}
			c.log.V(4).Info("Component pvc count per zone", "component", componentType, "count", pvcCountPerZone)
		}
	}

	// 2. build node group set
	builder := &nodegroup.Builder{
		ConfigProvider: c.ConfigProvider,
	}
	zones, err := c.provider.FindRegionZones(c.network.Spec.Region)
	if err != nil {
		return ctrl.Result{}, errors.AddStack(err)
	}
	if len(zones) != replicaFactor {
		return ctrl.Result{}, errors.Errorf("region %s have %d zones, expect %d zones", c.network.Spec.Region, len(zones), replicaFactor)
	}
	var (
		ownerName string
		lbs       map[string]string
		refObj    client.Object
	)
	if c.cluster != nil {
		ownerName = c.cluster.PrefixedId()
		lbs = c.cluster.Labels
		refObj = c.cluster
	} else {
		lbs = c.cdc.Labels
		refObj = c.cdc
	}
	if c.onlyCdc {
		ownerName = utils.TiCDCNodeGroupNamePrefix(c.onlyCdc, c.cdc.Spec.UpstreamClusterName, c.cluster)
	}

	tiflashAZs := zones
	if c.tiflashSingleAZDeploy {
		if c.cluster.Annotations == nil {
			c.cluster.Annotations = make(map[string]string)
		}
		zone, ok := c.cluster.Annotations[nodegroup.NodeScalerTiFlashAz]
		if !ok {
			// randomly select a zone for different clusters tiflash to avoid all customers tiflash deployed to the same zone
			randIdx := rand.Intn(len(zones))
			zone = zones[randIdx]
			c.cluster.Annotations[nodegroup.NodeScalerTiFlashAz] = zone
			// persist tiflash AZ info in cluster annotation
			if err := c.InfraClient.Update(c.ctx, c.cluster); err != nil {
				return ctrl.Result{}, err
			}
		}
		c.log.Info("Tiflash single AZ deploy", "tc", c.cluster.ClusterScopeName(), "zone", zone)
		tiflashAZs = []string{zone}
	}

	ngSet, err := builder.Build(&nodegroup.NodeTopologySpec{
		Provider:              c.provider,
		Components:            components,
		Zones:                 zones,
		Tags:                  cloudv1.TagsFromLabels(lbs),
		OwnerName:             ownerName,
		CurrentSpreads:        currentSpreads,
		CurrentPVCSpreads:     currentPVCSpreads,
		TiFlashDeployAZs:      tiflashAZs,
		MixedDeployPDWithTiKV: mixedDeployPDWithTiKV,
		SchedulePreferPVC:     c.scheduleNodePreferPVC,
	}, useZeroPDNodesNowForMixedDeployPDWithTiKV)
	if err != nil {
		return ctrl.Result{}, errors.AddStack(err)
	}
	// 3. apply node group set
	// TODO: policy could be exposed in CR or controlled by gitops
	// default to scale-only policy to achieve the same behavior with legacy logic
	if utils.HasDryrunAnnotation(refObj) {
		changes, err := c.gardener.ApplyNodeGroupSet(c.ctx, c.shootName, ngSet, gardener.NodeGroupPolicyDryrun)
		if err != nil {
			return ctrl.Result{}, errors.Annotate(err, "Dryrun ApplyNodeGroupSet")
		}
		for _, change := range changes {
			c.log.Info("[Dryrun] update node groups", "cluster", c.tc.Name, "change", change.String())
		}
	} else {
		changes, err := c.gardener.ApplyNodeGroupSet(c.ctx, c.shootName, ngSet, gardener.NodeGroupPolicyScaleOnly)
		if err != nil {
			if apierrors.IsConflict(err) {
				c.log.Info("update shoot conflicts, requeue")
				return ctrl.Result{RequeueAfter: defaultRequeueInterval}, nil
			}
			return ctrl.Result{}, errors.Annotate(err, "ApplyNodeGroupSet")
		}
		for _, change := range changes {
			c.log.Info("update node groups", "cluster", c.tc.Name, "change", change.String())
		}
	}

	return ctrl.Result{}, nil
}

func (c *scalingContext) componentCountPerZone(component cloudv1.ComponentType) ([]nodegroup.CountPerZone, error) {
	topo, err := k8s.GetComponentTopology(c.ctx, c.shootClient, types.NamespacedName{Namespace: c.tc.Namespace, Name: c.tc.Name}, component)
	if err != nil {
		return nil, errors.Annotate(err, "GetComponentTopology")
	}

	var zones []string
	for zone := range topo.Scheduled {
		zones = append(zones, zone)
	}
	sort.Strings(zones)

	var countPerZoneList []nodegroup.CountPerZone
	for _, zone := range zones {
		countPerZoneList = append(countPerZoneList, nodegroup.CountPerZone{
			Zone:  zone,
			Count: len(topo.Scheduled[zone]),
		})
	}
	return countPerZoneList, nil
}

func (c *scalingContext) componentPVCCountPerZone(component cloudv1.ComponentType) ([]nodegroup.CountPerZone, error) {
	pvcTopo, err := k8s.GetComponentPVCTopology(c.ctx, c.shootClient, types.NamespacedName{Namespace: c.tc.Namespace, Name: c.tc.Name}, component)
	if err != nil {
		return nil, errors.Annotate(err, "GetComponentPVCTopology")
	}

	c.log.V(4).Info("component topology", "component", component, "pvcTopo", pvcTopo)

	var zones []string
	for zone := range pvcTopo.Bound {
		zones = append(zones, zone)
	}
	sort.Strings(zones)

	var countPerZoneList []nodegroup.CountPerZone
	for _, zone := range zones {
		countPerZoneList = append(countPerZoneList, nodegroup.CountPerZone{
			Zone:  zone,
			Count: len(pvcTopo.Bound[zone]),
		})
	}
	return countPerZoneList, nil
}

func (c *scalingContext) tidbDesiredNodeCount() int {
	nodeCount := c.tc.TiDBStsDesiredReplicas()
	// if the actual replicas is greater than desired replicas, we should provide enough nodes to tidb-operator
	// and wait operator complete DB layer scale-in first
	if nodeCount < c.tc.TiDBStsActualReplicas() {
		nodeCount = c.tc.TiDBStsActualReplicas()
	}
	return int(nodeCount)
}

func (c *scalingContext) tiflashDesiredNodeCount() int {
	nodeCount := c.tc.TiFlashStsDesiredReplicas()
	if nodeCount < c.tc.TiFlashStsActualReplicas() {
		nodeCount = c.tc.TiFlashStsActualReplicas()
	}
	return int(nodeCount)
}

func (c *scalingContext) tiproxyDesiredNodeCount() int {
	nodeCount := c.tc.TiProxyStsDesiredReplicas()
	if nodeCount < c.tc.TiProxyStsActualReplicas() {
		nodeCount = c.tc.TiProxyStsActualReplicas()
	}
	return int(nodeCount)
}

func (c *scalingContext) tikvDesiredNodeCount() int {
	nodeCount := c.tc.TiKVStsDesiredReplicas()
	if nodeCount < c.tc.TiKVStsActualReplicas() {
		nodeCount = c.tc.TiKVStsActualReplicas()
	}
	return int(nodeCount)
}

func (c *scalingContext) pdDesiredNodeCount() (pdNodeCount int, mixedDeployPDWithTiKV bool, err error) {
	nodeCount := c.tc.PDStsDesiredReplicas()
	if nodeCount < c.tc.PDStsActualReplicas() {
		nodeCount = c.tc.PDStsActualReplicas()
	}

	zeroPDNodes, mixedDeployPDWithTiKV, err := mixeddeploypd.ShouldUseZeroPDNodes(int(nodeCount), c.tikvDesiredNodeCount(), c.tidbDesiredNodeCount(),
		c.tiflashDesiredNodeCount(), "gardener scaling", c.tc, c.provider, nil, c.shootClient, c.cluster, c.PDMixedDeployConfig)
	if err != nil {
		return int(nodeCount), false, err
	}
	if zeroPDNodes {
		return 0, mixedDeployPDWithTiKV, nil
	}
	return int(nodeCount), mixedDeployPDWithTiKV, nil
}

func (c *scalingContext) ticdcDesiredNodeCount() int {
	return c.cdc.Spec.Replicas
}