package capacitycontroller

import (
	"math"
	"sort"
)

type SiteResources struct {
	Site                     string
	CurrentReplicas          int
	MinReplicas              int
	MaxReplicas              int
	AllocatableMemoryBytes   int64
	UsedMemoryBytes          int64
	WorkerRequestMemoryBytes int64
	AllocatableMilliCPU      int64
	UsedMilliCPU             int64
	WorkerRequestMilliCPU    int64
	MemoryPressure           bool
	Nodes                    []NodeResources
}

type NodeResources struct {
	NodeName               string
	CurrentReplicas        int
	AllocatableMemoryBytes int64
	UsedMemoryBytes        int64
	WorkerUsedMemoryBytes  int64
	AllocatableMilliCPU    int64
	UsedMilliCPU           int64
	WorkerUsedMilliCPU     int64
	MemoryPressure         bool
}

type PlanInput struct {
	Queued                      int
	Running                     int
	InstanceSlots               int
	TargetOutstandingPerReplica int
	MemoryTargetPercent         int
	MemoryHardPercent           int
	CPUTargetPercent            int
	DatabaseConnections         int
	DatabaseConnectionBudget    int
	NonWorkerConnectionBudget   int
	WorkerDatabasePoolMax       int
	SynchronousDatabasePeers    int
	ScaleUpEligible             bool
	ScaleDownEligible           bool
	Sites                       []SiteResources
}

type SitePlan struct {
	Site            string  `json:"site"`
	CurrentReplicas int     `json:"currentReplicas"`
	DesiredReplicas int     `json:"desiredReplicas"`
	ResourceMax     int     `json:"resourceMax"`
	MemoryPercent   float64 `json:"memoryPercent"`
	CPUPercent      float64 `json:"cpuPercent"`
}

type Plan struct {
	RawDesired   int        `json:"rawDesired"`
	DesiredTotal int        `json:"desiredTotal"`
	CurrentTotal int        `json:"currentTotal"`
	FrozenReason string     `json:"frozenReason,omitempty"`
	Sites        []SitePlan `json:"sites"`
}

func CalculatePlan(input PlanInput) Plan {
	if input.InstanceSlots < 1 {
		input.InstanceSlots = 1
	}
	target := input.TargetOutstandingPerReplica
	if target < 1 {
		target = 2 * input.InstanceSlots
	}
	if input.MemoryTargetPercent < 1 {
		input.MemoryTargetPercent = 75
	}
	if input.MemoryHardPercent < 1 {
		input.MemoryHardPercent = 85
	}
	if input.CPUTargetPercent < 1 {
		input.CPUTargetPercent = 70
	}
	rawDesired := int(math.Ceil(float64(max(input.Queued+input.Running, 0)) / float64(target)))
	plan := Plan{RawDesired: rawDesired}
	minTotal := 0
	resourceMaxTotal := 0
	for _, site := range input.Sites {
		for _, node := range site.Nodes {
			if node.MemoryPressure {
				site.MemoryPressure = true
			}
		}
		current := max(site.CurrentReplicas, 0)
		minReplicas := max(site.MinReplicas, 0)
		configMax := max(site.MaxReplicas, minReplicas)
		resourceMax, memoryPercent, cpuPercent := siteResourceMaximum(
			site,
			input.MemoryTargetPercent,
			input.CPUTargetPercent,
		)
		resourceMax = min(resourceMax, configMax)
		if resourceMax < minReplicas && plan.FrozenReason == "" {
			plan.FrozenReason = "site_resource_budget"
		}
		resourceMax = max(resourceMax, minReplicas)
		plan.Sites = append(plan.Sites, SitePlan{
			Site: site.Site, CurrentReplicas: current, DesiredReplicas: minReplicas,
			ResourceMax: resourceMax, MemoryPercent: memoryPercent, CPUPercent: cpuPercent,
		})
		plan.CurrentTotal += current
		minTotal += minReplicas
		resourceMaxTotal += resourceMax
		if site.MemoryPressure || memoryPercent >= float64(input.MemoryHardPercent) {
			plan.FrozenReason = "node_memory_pressure"
		}
	}
	if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
		workerReplicaBudget := max(
			(input.DatabaseConnectionBudget-input.NonWorkerConnectionBudget)/input.WorkerDatabasePoolMax,
			0,
		)
		capSiteResourceMaxima(plan.Sites, workerReplicaBudget)
		resourceMaxTotal = 0
		for _, site := range plan.Sites {
			resourceMaxTotal += site.ResourceMax
		}
		if workerReplicaBudget < minTotal && plan.FrozenReason == "" {
			plan.FrozenReason = "database_connection_budget"
		}
	}
	desired := max(rawDesired, minTotal)
	desired = min(desired, resourceMaxTotal)

	if input.SynchronousDatabasePeers < 1 && plan.FrozenReason == "" {
		plan.FrozenReason = "database_not_synchronous"
	}
	if desired > plan.CurrentTotal {
		if input.DatabaseConnectionBudget > 0 && input.WorkerDatabasePoolMax > 0 {
			additional := max(
				(input.DatabaseConnectionBudget-input.DatabaseConnections)/input.WorkerDatabasePoolMax,
				0,
			)
			desired = min(desired, plan.CurrentTotal+additional)
		}
		if !input.ScaleUpEligible || plan.FrozenReason != "" ||
			(input.DatabaseConnectionBudget > 0 && input.DatabaseConnections >= input.DatabaseConnectionBudget) {
			desired = plan.CurrentTotal
			if plan.FrozenReason == "" {
				if !input.ScaleUpEligible {
					plan.FrozenReason = "scale_up_window"
				} else {
					plan.FrozenReason = "database_connection_budget"
				}
			}
		} else {
			desired = min(desired, max(plan.CurrentTotal*2, plan.CurrentTotal+2))
		}
	} else if desired < plan.CurrentTotal && !input.ScaleDownEligible {
		desired = plan.CurrentTotal
		if plan.FrozenReason == "" {
			plan.FrozenReason = "scale_down_stabilization"
		}
	}
	desired = max(desired, minTotal)
	plan.DesiredTotal = desired
	distributeDesiredReplicas(plan.Sites, desired)
	return plan
}

func capSiteResourceMaxima(sites []SitePlan, total int) {
	minimum := 0
	original := make(map[string]int, len(sites))
	for index := range sites {
		original[sites[index].Site] = sites[index].ResourceMax
		sites[index].ResourceMax = sites[index].DesiredReplicas
		minimum += sites[index].ResourceMax
	}
	target := max(total, minimum)
	assigned := minimum
	for assigned < target {
		sort.SliceStable(sites, func(left, right int) bool {
			leftRoom := original[sites[left].Site] - sites[left].ResourceMax
			rightRoom := original[sites[right].Site] - sites[right].ResourceMax
			if leftRoom != rightRoom {
				return leftRoom > rightRoom
			}
			return sites[left].Site < sites[right].Site
		})
		if original[sites[0].Site] <= sites[0].ResourceMax {
			break
		}
		sites[0].ResourceMax++
		assigned++
	}
}

func siteResourceMaximum(site SiteResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) {
	if len(site.Nodes) > 0 {
		if site.WorkerRequestMemoryBytes <= 0 || site.WorkerRequestMilliCPU <= 0 {
			return 0, 0, 0
		}
		memoryMax := 0
		cpuMax := 0
		memoryPercent := float64(0)
		cpuPercent := float64(0)
		for _, node := range site.Nodes {
			nodeMemoryPercent := usagePercent(node.UsedMemoryBytes, node.AllocatableMemoryBytes)
			nodeCPUPercent := usagePercent(node.UsedMilliCPU, node.AllocatableMilliCPU)
			memoryPercent = max(memoryPercent, nodeMemoryPercent)
			cpuPercent = max(cpuPercent, nodeCPUPercent)
			nonWorkerMemory := max(
				node.UsedMemoryBytes-node.WorkerUsedMemoryBytes,
				0,
			)
			memoryBudget := node.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorkerMemory
			memoryMax += int(max(memoryBudget, 0) / site.WorkerRequestMemoryBytes)
			nonWorkerCPU := max(
				node.UsedMilliCPU-node.WorkerUsedMilliCPU,
				0,
			)
			cpuBudget := node.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorkerCPU
			cpuMax += int(max(cpuBudget, 0) / site.WorkerRequestMilliCPU)
		}
		return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
	}
	memoryPercent := usagePercent(site.UsedMemoryBytes, site.AllocatableMemoryBytes)
	cpuPercent := usagePercent(site.UsedMilliCPU, site.AllocatableMilliCPU)
	memoryMax := site.MaxReplicas
	if site.AllocatableMemoryBytes > 0 && site.WorkerRequestMemoryBytes > 0 {
		nonWorker := max(site.UsedMemoryBytes-int64(site.CurrentReplicas)*site.WorkerRequestMemoryBytes, 0)
		budget := site.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker
		memoryMax = int(max(budget, 0) / site.WorkerRequestMemoryBytes)
	}
	cpuMax := site.MaxReplicas
	if site.AllocatableMilliCPU > 0 && site.WorkerRequestMilliCPU > 0 {
		nonWorker := max(site.UsedMilliCPU-int64(site.CurrentReplicas)*site.WorkerRequestMilliCPU, 0)
		budget := site.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker
		cpuMax = int(max(budget, 0) / site.WorkerRequestMilliCPU)
	}
	return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
}

func distributeDesiredReplicas(sites []SitePlan, desired int) {
	if len(sites) == 0 {
		return
	}
	assigned := 0
	for index := range sites {
		assigned += sites[index].DesiredReplicas
	}
	for assigned < desired {
		sort.SliceStable(sites, func(left, right int) bool {
			leftRoom := sites[left].ResourceMax - sites[left].DesiredReplicas
			rightRoom := sites[right].ResourceMax - sites[right].DesiredReplicas
			if leftRoom != rightRoom {
				return leftRoom > rightRoom
			}
			if sites[left].DesiredReplicas != sites[right].DesiredReplicas {
				return sites[left].DesiredReplicas < sites[right].DesiredReplicas
			}
			return sites[left].Site < sites[right].Site
		})
		if sites[0].DesiredReplicas >= sites[0].ResourceMax {
			break
		}
		sites[0].DesiredReplicas++
		assigned++
	}
	sort.Slice(sites, func(left, right int) bool { return sites[left].Site < sites[right].Site })
}

func usagePercent(used int64, allocatable int64) float64 {
	if allocatable <= 0 {
		return 0
	}
	return float64(used) * 100 / float64(allocatable)
}