package capacitycontroller

import (
	"math"
	"sort"
)

type PoolResources struct {
	PoolID                   string
	Demand                   int
	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
	Pools                       []PoolResources
}

type PoolPlan struct {
	PoolID          string  `json:"poolId"`
	Demand          int     `json:"demand"`
	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"`
	Pools        []PoolPlan `json:"pools"`
}

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 _, pool := range input.Pools {
		for _, node := range pool.Nodes {
			if node.MemoryPressure {
				pool.MemoryPressure = true
			}
		}
		current := max(pool.CurrentReplicas, 0)
		minReplicas := max(pool.MinReplicas, 0)
		configMax := max(pool.MaxReplicas, minReplicas)
		resourceMax, memoryPercent, cpuPercent := poolResourceMaximum(
			pool,
			input.MemoryTargetPercent,
			input.CPUTargetPercent,
		)
		resourceMax = min(resourceMax, configMax)
		if resourceMax < minReplicas && plan.FrozenReason == "" {
			plan.FrozenReason = "pool_resource_budget"
		}
		resourceMax = max(resourceMax, minReplicas)
		plan.Pools = append(plan.Pools, PoolPlan{
			PoolID: pool.PoolID, CurrentReplicas: current, DesiredReplicas: minReplicas,
			Demand:      pool.Demand,
			ResourceMax: resourceMax, MemoryPercent: memoryPercent, CPUPercent: cpuPercent,
		})
		plan.CurrentTotal += current
		minTotal += minReplicas
		resourceMaxTotal += resourceMax
		if pool.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,
		)
		capPoolResourceMaxima(plan.Pools, workerReplicaBudget)
		resourceMaxTotal = 0
		for _, pool := range plan.Pools {
			resourceMaxTotal += pool.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.Pools, desired)
	return plan
}

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

func poolResourceMaximum(pool PoolResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) {
	if len(pool.Nodes) > 0 {
		if pool.WorkerRequestMemoryBytes <= 0 || pool.WorkerRequestMilliCPU <= 0 {
			return 0, 0, 0
		}
		memoryMax := 0
		cpuMax := 0
		memoryPercent := float64(0)
		cpuPercent := float64(0)
		for _, node := range pool.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) / pool.WorkerRequestMemoryBytes)
			nonWorkerCPU := max(
				node.UsedMilliCPU-node.WorkerUsedMilliCPU,
				0,
			)
			cpuBudget := node.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorkerCPU
			cpuMax += int(max(cpuBudget, 0) / pool.WorkerRequestMilliCPU)
		}
		return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
	}
	memoryPercent := usagePercent(pool.UsedMemoryBytes, pool.AllocatableMemoryBytes)
	cpuPercent := usagePercent(pool.UsedMilliCPU, pool.AllocatableMilliCPU)
	memoryMax := pool.MaxReplicas
	if pool.AllocatableMemoryBytes > 0 && pool.WorkerRequestMemoryBytes > 0 {
		nonWorker := max(pool.UsedMemoryBytes-int64(pool.CurrentReplicas)*pool.WorkerRequestMemoryBytes, 0)
		budget := pool.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker
		memoryMax = int(max(budget, 0) / pool.WorkerRequestMemoryBytes)
	}
	cpuMax := pool.MaxReplicas
	if pool.AllocatableMilliCPU > 0 && pool.WorkerRequestMilliCPU > 0 {
		nonWorker := max(pool.UsedMilliCPU-int64(pool.CurrentReplicas)*pool.WorkerRequestMilliCPU, 0)
		budget := pool.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker
		cpuMax = int(max(budget, 0) / pool.WorkerRequestMilliCPU)
	}
	return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent
}

func distributeDesiredReplicas(pools []PoolPlan, desired int) {
	if len(pools) == 0 {
		return
	}
	assigned := 0
	for index := range pools {
		assigned += pools[index].DesiredReplicas
	}
	for assigned < desired {
		sort.SliceStable(pools, func(left, right int) bool {
			leftDemandPerReplica := float64(max(pools[left].Demand, 0)) / float64(max(pools[left].DesiredReplicas, 1))
			rightDemandPerReplica := float64(max(pools[right].Demand, 0)) / float64(max(pools[right].DesiredReplicas, 1))
			if leftDemandPerReplica != rightDemandPerReplica {
				return leftDemandPerReplica > rightDemandPerReplica
			}
			leftRoom := pools[left].ResourceMax - pools[left].DesiredReplicas
			rightRoom := pools[right].ResourceMax - pools[right].DesiredReplicas
			if leftRoom != rightRoom {
				return leftRoom > rightRoom
			}
			if pools[left].DesiredReplicas != pools[right].DesiredReplicas {
				return pools[left].DesiredReplicas < pools[right].DesiredReplicas
			}
			return pools[left].PoolID < pools[right].PoolID
		})
		if pools[0].DesiredReplicas >= pools[0].ResourceMax {
			break
		}
		pools[0].DesiredReplicas++
		assigned++
	}
	sort.Slice(pools, func(left, right int) bool { return pools[left].PoolID < pools[right].PoolID })
}

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