package spec

import (
	"fmt"

	infrautil "github.com/tidbcloud/infra-provider/pkg/util"

	"github.com/tidbcloud/eks-provider/api/v1alpha1"
)

const (
	TiflowNodegroupPrefix = "tiflow-"
)

type nodegroup struct {
	name           string
	region         string
	zone           string
	instanceNumber int
}

type tiflowResourceAllocator interface {
	Execute(infraResource *TiflowClusterInfraResource) (map[string]*nodegroup, error)
}

func getTiflowResourceAllocator(strategy v1alpha1.TiflowAllocateStrategy) tiflowResourceAllocator {
	switch strategy {
	case v1alpha1.TiflowZoneAveragingStrategy:
		return tiflowZoneAveragingAllocatorSingleton
	default:
		return nil
	}
}

var tiflowZoneAveragingAllocatorSingleton = &tiflowZoneAveragingAllocator{}

type tiflowZoneAveragingAllocator struct{}

func (t *tiflowZoneAveragingAllocator) Execute(infraResource *TiflowClusterInfraResource) (map[string]*nodegroup, error) {
	region := infraResource.Network.Spec.Region
	zones, err := infraResource.Provider.FindRegionZones(region)
	if err != nil {
		return nil, err
	}

	instanceNumberOfZoneMap := t.allocateNumberByZones(infraResource.TiflowSpec.InstanceNumber, zones)
	tiflowNgMap := map[string]*nodegroup{}
	for zone, instanceNumber := range instanceNumberOfZoneMap {
		ng := &nodegroup{
			name:           TiflowNodegroupPrefix + infrautil.FnvHash(fmt.Sprintf("%s-%s-%s", infraResource.TiflowClusterName, infraResource.TiflowSpec.SpecCode, zone)),
			region:         region,
			zone:           zone,
			instanceNumber: instanceNumber,
		}
		tiflowNgMap[ng.name] = ng
	}
	return tiflowNgMap, nil
}

func (t *tiflowZoneAveragingAllocator) allocateNumberByZones(instanceNumber int, zones []string) map[string]int {
	instanceNumberOfZoneMap := map[string]int{}
	zoneTotal := len(zones)
	s := instanceNumber / zoneTotal
	y := instanceNumber % zoneTotal
	for i, zone := range zones {
		var numberOfZone int
		if y > i {
			numberOfZone = s + 1
		} else {
			numberOfZone = s
		}

		if numberOfZone > 0 {
			instanceNumberOfZoneMap[zone] = numberOfZone
		}
	}
	return instanceNumberOfZoneMap
}