package nodegroup

import (
	"fmt"
	"math"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"

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

const (
	gcpPerVolumeSize = "375Gi"

	gcpLocalVolumeCount   = "pingcap.com/gcp-local-volume-count"
	gcpProviderAnnotation = "pingcap.com/gcp-provider-annotation"
	awsLocalSSDLabel      = "pingcap.com/aws-local-ssd"

	// TODO: this package needs to be refactored.
	// Currently this package serves as part of the scheduler
	// but some scheduling logic is leaked to the controller
	// which makes it hard to maintain especially for different providers
	NodeScalerFinalizer = "node.scaler.cloud.pingcap.com/finalizer"
	NodeScalerIgnore    = "node.scaler.cloud.pingcap.com/ignore"
	NodeScalerTiFlashAz = "node.scaler.cloud.pingcap.com/tiflash-az"
)

// HACK: some legacy node groups are added manually and can neither be adopted or changed,
// use this hack method to be compatible with them
// reach out to @aylei for detailed explanations if have any concern
func processManualAddedGroups(ng *cloudv1.NodegroupSpec) {
	switch ng.Name {
	// node groups of TiKV of cluster 1311285575352848384 were replaced manually before with new name "${name}-big"
	case "2ddb6314", "2ddb6315", "2ddb6316":
		ng.Name += "-big"
	default:
		// noop
	}
}

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

type CountPerZone struct {
	Count int
	Zone  string
}

type NodeTopologySpec struct {
	Provider   *cloudv1.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

	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
}

func (b *Builder) Build(spec *NodeTopologySpec, useZeroPDNodesNowForMixedDeployPDWithTiKV bool) (*Set, error) {
	if spec.Components == nil {
		return &Set{Name: spec.OwnerName, NodeGroups: nil}, nil
	}
	var ngs []cloudv1.NodegroupSpec
	baseTags := common.ConcatenateTags(spec.Tags,
		map[string]string{common.UsedBy: common.CustomerTiDB},
	)
	// Cluster component node groups
	components := spec.Components
	for _, comp := range components {
		ms, err := b.ConfigProvider.GetMachineSpec(spec.Provider, comp)
		if err != nil {
			return nil, err
		}
		// make spec copies, to avoid interference
		newSpec := *spec
		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, ms)
		} else {
			compNgs, err = b.buildNodeGroupsPerZone(&newSpec, useZeroPDNodesNowForMixedDeployPDWithTiKV, comp, ms)
		}
		if err != nil {
			return nil, err
		}
		ngs = append(ngs, compNgs...)
	}

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

func (b *Builder) buildNodeGroupsForSingleAZ(spec *NodeTopologySpec, zone string, component cloudv1.ClusterComponent, ms *cloudv1.MachineSpec) ([]cloudv1.NodegroupSpec, error) {
	expectedNodes := component.GetReplicas()
	componentType := component.GetComponentType()
	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 {
		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 *NodeTopologySpec, useZeroPDNodesNowForMixedDeployPDWithTiKV bool,
	component cloudv1.ClusterComponent, ms *cloudv1.MachineSpec) ([]cloudv1.NodegroupSpec, error) {

	expectNodes := component.GetReplicas()
	zones := spec.Zones

	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 := expectNodes - 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 {
		return nil, fmt.Errorf("desired nodes %d is less than current requirements %d of component %s, pods must be scale-in first before nodes", expectNodes, 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 {
		if countPerZone.Count > 0 {
			ngs = append(ngs, b.buildNodeGroup(spec, component, ms, countPerZone.Count, 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 *NodeTopologySpec, component cloudv1.ClusterComponent, ms *cloudv1.MachineSpec, count int, zone string) cloudv1.NodegroupSpec {
	worker := cloudv1.NodegroupSpec{
		Name:          util.FnvHash(NodeGroupName(spec.OwnerName, string(component.GetComponentType()), zone)),
		Labels:        r.getNodeLabels(spec, component, ms, zone),
		Annotations:   r.getNodeAnnotations(spec.Provider.Spec.Type, component, ms),
		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},
	}
	processManualAddedGroups(&worker)
	return worker
}

func (r *Builder) getNodeLabels(spec *NodeTopologySpec, 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[GardenerWorkerNodeLabelName] = NodeGroupName(spec.OwnerName, string(component.GetComponentType()), zone)
	nodeLabels[GardenerWorkerPrefixLabelName] = NodeGroupNamePrefix(spec.OwnerName, string(component.GetComponentType()))
	nodeLabels[GardenerNodeAvailableZoneInfo] = zone
	nodeLabels[GardenerNodeInstanceType] = ms.MachineType

	if component.GetUseLocalStorage() {
		nodeLabels[awsLocalSSDLabel] = "true"
		nodeLabels[NodeProtection] = "true"
	}

	return nodeLabels
}

func (r *Builder) getNodeAnnotations(providerType string, component cloudv1.ClusterComponent, ms *cloudv1.MachineSpec) map[string]string {
	nodeAnnotations := make(map[string]string)

	storageSize, needStorage := component.GetStorageSize()
	if providerType == consts.ProviderTypeGCP && needStorage {
		// for TiKV and TiFlash, compute the number of local SSD disks
		if component.GetComponentType() == cloudv1.ComponentTiKV || component.GetComponentType() == cloudv1.ComponentTiFlash {
			nodeAnnotations[gcpProviderAnnotation] = "true"
			if component.GetUseLocalStorage() {
				nodeAnnotations[gcpLocalVolumeCount] = fmt.Sprintf("%d", convertGCPVolumeCounts(*storageSize))
			} else {
				nodeAnnotations[gcpLocalVolumeCount] = "0"
			}
		}
	}

	return nodeAnnotations
}

func convertGCPVolumeCounts(size resource.Quantity) int {
	volumeSize := resource.MustParse(gcpPerVolumeSize)
	volumeCount := int(math.Ceil(float64(size.Value()) / float64(volumeSize.Value())))
	// volumeCount must be an even number
	if volumeCount%2 == 1 {
		return volumeCount + 1
	}
	return volumeCount
}