package nodegroup

import (
	"fmt"
	"math"

	corev1 "k8s.io/api/core/v1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/controllers/common"
	"github.com/tidbcloud/infra-provider/pkg/gitops"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	"github.com/tidbcloud/infra-provider/pkg/util"
)

type Set struct {
	Name       string
	NodeGroups []cloudv1.NodegroupSpec
}

type CountPerZone struct {
	Count int
	Zone  string
}

type NodeTopology struct {
	Cluster    *cloudv1.Cluster
	Provider   gitops.Provider
	Components []cloudv1.ClusterComponent
	Zones      []string
	Tags       map[string]string

	OwnerName string

	// CurrentSpreads describe the current node count per zone as a reference when
	// calculating the new desired state, count is considered as 0 if the zone does not
	// appear in CurrentSpreads.
	CurrentSpreads    map[cloudv1.ComponentType][]CountPerZone
	CurrentPVCSpreads map[cloudv1.ComponentType][]CountPerZone

	ComponentCurrentActualReplicas map[cloudv1.ComponentType]int

	TiFlashDeployAZs      []string
	MixedDeployPDWithTiKV bool
	SchedulePreferPVC     bool
}

// Builder takes a NodeTopologySpec and build the nodegroup.Set to match the spec
type Builder struct {
	ConfigProvider gitops.ConfigProvider

	// This field is defined to ignore replica validation
	// Builder will not check current replicas and just build the expected set if it is set
	IgnoreReplicaValidation bool

	// This field is defined to support scale out 1 extra tidb node in a zone when rolling update
	// If it is enabled, max = min + 1
	// Now it is only enabled for tidb operator v2
	ExtraTiDBNode bool
}

func (b *Builder) Build(topo *NodeTopology, useZeroPDNodesNowForMixedDeployPDWithTiKV bool, tiproxyMixedDeployedWithPD bool) (*Set, error) {
	if topo.Components == nil {
		return &Set{Name: topo.OwnerName, NodeGroups: nil}, nil
	}
	var ngs []cloudv1.NodegroupSpec
	baseTags := common.ConcatenateTags(topo.Tags,
		map[string]string{common.UsedBy: common.CustomerTiDB},
	)
	// Cluster component node groups
	components := topo.Components
	for _, comp := range components {
		compActualReplicas := topo.ComponentCurrentActualReplicas[comp.GetComponentType()]

		ms, err := b.ConfigProvider.GetMachineSpec(topo.Provider, comp)
		if err != nil {
			return nil, err
		}
		// make spec copies, to avoid interference
		newSpec := *topo
		newSpec.Tags = common.ConcatenateComponetTags(baseTags, string(comp.GetComponentType()))
		var compNgs []cloudv1.NodegroupSpec
		if comp.GetComponentType() == cloudv1.ComponentTiFlash && len(newSpec.TiFlashDeployAZs) == 1 {
			compNgs, err = b.buildNodeGroupsForSingleAZ(&newSpec, newSpec.TiFlashDeployAZs[0], comp, compActualReplicas, ms)
		} else {
			compNgs, err = b.buildNodeGroupsPerZone(&newSpec, useZeroPDNodesNowForMixedDeployPDWithTiKV, tiproxyMixedDeployedWithPD, comp, compActualReplicas, ms)
		}
		if err != nil {
			return nil, err
		}
		ngs = append(ngs, compNgs...)
	}

	return &Set{Name: topo.OwnerName, NodeGroups: ngs}, nil
}

func (b *Builder) buildNodeGroupsForSingleAZ(spec *NodeTopology, zone string, component cloudv1.ClusterComponent, componentCurrentActualReplicas int, ms *cloudv1.MachineSpec) ([]cloudv1.NodegroupSpec, error) {
	expectedNodes := component.GetReplicas() + int(spec.Cluster.GetStandbyNodeCnt(component.GetComponentType(), zone))
	componentType := component.GetComponentType()

	// Check if the pod is really scaled-in to avoid starting the node scaling down when all scale-in pods crash and scaleOutCount becomes 0
	// pods will be scaled in tidbcluster task
	if scaleInPodCount := expectedNodes - componentCurrentActualReplicas; scaleInPodCount < 0 && !b.IgnoreReplicaValidation {
		return nil, fmt.Errorf("desired replicas %d is less than current updated replicas %d of component %s, pods must be scale-in first before nodes", expectedNodes, componentCurrentActualReplicas, string(component.GetComponentType()))
	}

	// if component pods are scaled as expected, we can scale the nodes
	currentSpread, ok := spec.CurrentSpreads[componentType]
	if !ok {
		// all pods are newly created, so just schedule them into the desired zone
		return []cloudv1.NodegroupSpec{b.buildNodeGroup(spec, component, ms, expectedNodes, zone)}, nil
	}
	totalScheduled := 0
	for idx := range currentSpread {
		totalScheduled += currentSpread[idx].Count
	}
	if expectedNodes-totalScheduled < 0 && !b.IgnoreReplicaValidation {
		return nil, fmt.Errorf("desired nodes %d is less than current requirements %d of component %s, pods must be scale-in first before nodes", expectedNodes, totalScheduled, string(component.GetComponentType()))
	}
	// some pods are already scheduled, do not touch these pods and node group
	otherAZScheduled := 0
	var ngs []cloudv1.NodegroupSpec
	for idx := range currentSpread {
		countPerZone := currentSpread[idx]
		if countPerZone.Zone != zone {
			// keep the existing node group as is
			ngs = append(ngs, b.buildNodeGroup(spec, component, ms, countPerZone.Count, countPerZone.Zone))
			otherAZScheduled += currentSpread[idx].Count
		}
	}
	thisAZCount := expectedNodes - otherAZScheduled
	ngs = append(ngs, b.buildNodeGroup(spec, component, ms, thisAZCount, zone))
	return ngs, nil
}

// buildNodeGroupsPerZone spread the node replicas to the given zone list evenly (best-effort), e.g.
// replicas: 1 zones: [a, b, c] => [{zone: a, replicas: 1}]
// replicas: 2 zones: [a, b, c] => [{zone: a, replicas: 1}, {zone: b, replicas: 1}]
// replicas: 3 zones: [a, b, c] => [{zone: a, replicas: 1}, {zone: b, replicas: 1}, {zone: c, replicas: 1}]
// replicas: 5 zones: [a, b, c] => [{zone: a, replicas: 2}, {zone: b, replicas: 2}, {zone: c, replicas: 1}]
// spec.CurrentSpreads will be taken into account, the final spreads might not be even with uneven CurrentSpreads, e.g.
// replicas: 6 CurrentSpreads [{a, 3}], final topology would be [{a, 3}, {b, 2}, {c, 1}]
func (b *Builder) buildNodeGroupsPerZone(spec *NodeTopology, useZeroPDNodesNowForMixedDeployPDWithTiKV bool, tiproxyMixedDeployedWithPD bool,
	component cloudv1.ClusterComponent, componentCurrentActualReplicas int, ms *cloudv1.MachineSpec,
) ([]cloudv1.NodegroupSpec, error) {
	expectedNodes := component.GetReplicas()
	zones := spec.Zones

	if component.GetComponentType() == cloudv1.ComponentTiProxy && tiproxyMixedDeployedWithPD {
		// if the tiproxy is not standalone(or mix deployed with PD), we should not adjust nodegroup for TiProxy nodes
		return nil, nil
	}

	// Check if the pod is really scaled-in to avoid starting the node scaling down when all scale-in pods crash and scaleOutCount becomes 0
	// pods will be scaled in tidbcluster task
	if scaleInPodCount := expectedNodes - componentCurrentActualReplicas; scaleInPodCount < 0 && !b.IgnoreReplicaValidation {
		return nil, fmt.Errorf("desired replicas %d is less than current updated replicas %d of component %s, pods must be scale-in first before nodes", expectedNodes, componentCurrentActualReplicas, string(component.GetComponentType()))
	}

	// if component pods are scaled as expected, we can scale the nodes
	currentRequirements := 0
	var countPerZones []CountPerZone
	var pvcCountPerZones []CountPerZone
	if !(useZeroPDNodesNowForMixedDeployPDWithTiKV && component.GetComponentType() == cloudv1.ComponentPD) {
		for _, zone := range zones {
			countPerZone := nodeCountInZone(spec.CurrentSpreads[component.GetComponentType()], zone)
			currentRequirements += countPerZone.Count
			countPerZones = append(countPerZones, countPerZone)

			pvcCountPerZone := nodeCountInZone(spec.CurrentPVCSpreads[component.GetComponentType()], zone)
			pvcCountPerZones = append(pvcCountPerZones, pvcCountPerZone)
		}
	}

	// the number of nodes that will be added in scale out
	// for scale-in, the count should zero since we will wait the pods to be scaled-in first
	scaleOutCount := expectedNodes - currentRequirements

	// if a cluster created in mixed-deplop-pd mode, that no need to create any pd nodes
	// if a cluster scale-in to mixed-deplop-pd mode, that no need to create new pd nodes
	if spec.MixedDeployPDWithTiKV && component.GetComponentType() == cloudv1.ComponentPD {
		scaleOutCount = 0
	}

	if scaleOutCount < 0 && !b.IgnoreReplicaValidation {
		return nil, fmt.Errorf("desired nodes %d is less than current requirements %d of component %s, pods must be scale-in first before nodes", expectedNodes, currentRequirements, string(component.GetComponentType()))
	}

	for i := 0; i < scaleOutCount; i++ {
		if spec.SchedulePreferPVC {
			scheduleOneNodePreferPVC(countPerZones, pvcCountPerZones, component.GetComponentType())
		} else {
			scheduleOneNode(countPerZones)
		}
	}

	var ngs []cloudv1.NodegroupSpec
	for _, countPerZone := range countPerZones {
		nodeCountPerZone := countPerZone.Count + int(spec.Cluster.GetStandbyNodeCnt(component.GetComponentType(), countPerZone.Zone))
		if nodeCountPerZone > 0 || (b.ExtraTiDBNode && component.GetComponentType() == cloudv1.ComponentTiDB) {
			ngs = append(ngs, b.buildNodeGroup(spec, component, ms, nodeCountPerZone, countPerZone.Zone))
		}
	}
	return ngs, nil
}

// scheduleOneNode schedule one node to a zone
func scheduleOneNode(countPerZones []CountPerZone) {
	leastZoneIndex := 0
	leastZoneCount := math.MaxInt32
	for i := range countPerZones {
		if countPerZones[i].Count < leastZoneCount {
			leastZoneCount = countPerZones[i].Count
			leastZoneIndex = i
		}
	}
	countPerZones[leastZoneIndex].Count += 1
}

// scheduleOneNodePreferPVC one node to a zone
//
//  1. find zones that have the least node count, because we want to balance the node count.
//  2. find zones that have the most PVC count if there are multiple zones with the least
//     node count, it is a workaround for https://jira.tidbcloud.com/browse/TCOC-781.
//  3. increase the node count of the fisrt zone in the above list
func scheduleOneNodePreferPVC(nodeCountPerZones, pvcCountPerZones []CountPerZone, component cloudv1.ComponentType) {
	// find candidate zones that have the least node count
	candidateForSkew := make([]int, 0)
	leaseNodeCount := math.MaxInt32
	for zoneIndex := range nodeCountPerZones {
		nodeCount := nodeCountPerZones[zoneIndex].Count

		if nodeCount == leaseNodeCount {
			candidateForSkew = append(candidateForSkew, zoneIndex)
		} else if nodeCount < leaseNodeCount {
			candidateForSkew = []int{zoneIndex}
			leaseNodeCount = nodeCount
		}
	}

	if len(candidateForSkew) == 1 || component == cloudv1.ComponentTiKV {
		nodeCountPerZones[candidateForSkew[0]].Count += 1
		return
	}

	// find candidate zones that have the most PVC count
	candidateForPreferPVC := make([]int, 0)
	mostPVCCount := math.MinInt32
	for _, zoneIndex := range candidateForSkew {
		pvcCount := pvcCountPerZones[zoneIndex].Count

		if pvcCount == mostPVCCount {
			candidateForPreferPVC = append(candidateForPreferPVC, zoneIndex)
		} else if pvcCount > mostPVCCount {
			candidateForPreferPVC = []int{zoneIndex}
			mostPVCCount = pvcCount
		}
	}

	nodeCountPerZones[candidateForPreferPVC[0]].Count += 1
}

func nodeCountInZone(spreads []CountPerZone, zone string) CountPerZone {
	for _, countPerZone := range spreads {
		if countPerZone.Zone == zone {
			return countPerZone
		}
	}
	return CountPerZone{Zone: zone, Count: 0}
}

func (r *Builder) buildNodeGroup(spec *NodeTopology, component cloudv1.ClusterComponent, ms *cloudv1.MachineSpec, count int, zone string) cloudv1.NodegroupSpec {
	worker := cloudv1.NodegroupSpec{
		Name:          util.FnvHash(namer.NodeGroupName(spec.OwnerName, string(component.GetComponentType()), zone)),
		Labels:        r.getNodeLabels(spec, component, ms, zone),
		ComponentType: component.GetComponentType(),
		MachineSpec:   ms,
		Taints: []corev1.Taint{{
			Key:    "dedicated",
			Value:  "tidb-" + string(component.GetComponentType()),
			Effect: corev1.TaintEffectNoSchedule,
		}},
		Maximum: int32(count),
		Minimum: int32(count),
		Zones:   []string{zone},
	}
	if component.GetComponentType() == cloudv1.ComponentTiDB && r.ExtraTiDBNode {
		worker.Maximum += 1
	}
	return worker
}

func (r *Builder) getNodeLabels(spec *NodeTopology, component cloudv1.ClusterComponent, ms *cloudv1.MachineSpec, zone string) map[string]string {
	nodeLabels := make(map[string]string)

	// range over Tags is always safe despite the error
	for k, v := range spec.Tags {
		nodeLabels[k] = v
	}

	// TODO(infra): Remove this block? It seems no hurt to add these labels for other providers.
	nodeLabels[consts.GardenerWorkerNodeLabelName] = namer.NodeGroupName(spec.OwnerName, string(component.GetComponentType()), zone)
	nodeLabels[consts.GardenerWorkerPrefixLabelName] = namer.NodeGroupNamePrefix(spec.OwnerName, string(component.GetComponentType()))
	nodeLabels[consts.GardenerNodeAvailableZoneInfo] = zone
	nodeLabels[consts.GardenerNodeInstanceType] = ms.MachineType

	return nodeLabels
}