class ClusterAgentAutoScaler {

    @VisibleForTesting
    Completable doAgentScaling() {
        return Completable.defer(() -> {
            if (!configuration.isAutoScalingAgentsEnabled()) {
                logger.debug("Auto scaling agents is not enabled");
                return Completable.complete();
            }

            List<Completable> actions = new ArrayList<>();

            Map<String, Job> allJobs = getAllJobs();
            Map<String, Task> allTasks = getAllTasks();
            List<AgentInstanceGroup> activeInstanceGroups = getActiveInstanceGroups();
            Map<AgentInstanceGroup, List<AgentInstance>> instancesForActiveInstanceGroups = getInstancesForInstanceGroups(activeInstanceGroups);
            Map<String, List<AgentInstance>> instancesForActiveInstanceGroupsById = instancesForActiveInstanceGroups.entrySet().stream()
                    .collect(Collectors.toMap(e -> e.getKey().getId(), Map.Entry::getValue));
            Map<String, Long> numberOfTasksOnAgents = getNumberOfTasksOnAgents(allTasks.values());
            Map<Tier, Set<String>> failedTaskIdsByTier = getFailedTaskIdsByTier(
                    schedulingService.getLastTaskPlacementFailures(), FailureKind.NEVER_TRIGGER_AUTOSCALING
            );

            long now = clock.wallTime();

            for (Tier tier : Tier.values()) {
                logger.info("Starting scaling actions for tier: {}", tier);
                TierAutoScalingConfiguration tierConfiguration = ClusterOperationFunctions.getTierConfiguration(tier, configuration);
                TierAutoScalerExecution tierAutoScalerExecution = tierTierAutoScalerExecutions.computeIfAbsent(
                        tier, k -> new TierAutoScalerExecution(tier, titusRuntime.getRegistry())
                );

                String primaryInstanceType = tierConfiguration.getPrimaryInstanceType();
                // This will throw an exception if not properly configured
                ResourceDimension tierResourceDimension = agentManagementService.getResourceLimits(primaryInstanceType);

                List<AgentInstanceGroup> activeScalableInstanceGroupsForTier = activeInstanceGroups.stream()
                        .filter(ig -> ig.getTier() == tier && ig.getInstanceType().equals(primaryInstanceType))
                        .collect(Collectors.toList());
                logger.info("{} active instance groups({}): {}", tier, activeScalableInstanceGroupsForTier.size(), activeScalableInstanceGroupsForTier);

                List<AgentInstance> idleInstancesForTier = getIdleInstancesForTier(tier, primaryInstanceType,
                        instancesForActiveInstanceGroups, numberOfTasksOnAgents, now, tierConfiguration.getIdleInstanceGracePeriodMs());
                tierAutoScalerExecution.getTotalIdleInstancesGauge().set(idleInstancesForTier.size());
                logger.info("{} idle instances({}): {}", tier, idleInstancesForTier.size(), idleInstancesForTier);

                Set<String> failedTaskIds = failedTaskIdsByTier.getOrDefault(tier, Collections.emptySet());
                tierAutoScalerExecution.getTotalFailedTasksGauge().set(failedTaskIds.size());
                logger.info("{} failed tasks({}): {}", tier, failedTaskIds.size(), failedTaskIds);

                int agentCountToScaleUp = 0;
                Set<String> potentialTaskIdsForScaleUp = new HashSet<>();

                boolean usedScaleUpCooldown = false;
                if (hasTimeElapsed(tierAutoScalerExecution.getLastScaleUp().get(), now, tierConfiguration.getScaleUpCoolDownMs())) {
                    int minIdleForTier = tierConfiguration.getMinIdle();
                    if (idleInstancesForTier.size() < minIdleForTier) {
                        int instancesNeededForMinIdle = minIdleForTier - idleInstancesForTier.size();
                        logger.info("{} needs {} instances to satisfy min idle {}", tier, instancesNeededForMinIdle, minIdleForTier);
                        agentCountToScaleUp += instancesNeededForMinIdle;
                    }

                    Set<String> placementFailureTaskIds = getFailedTaskIdsByTier(
                            schedulingService.getLastTaskPlacementFailures(), IGNORED_FAILURE_KINDS_WITH_LAUNCHGUARD
                    ).getOrDefault(tier, Collections.emptySet());
                    logger.info("{} had the placement excluding launch guard failures({}): {}", tier, placementFailureTaskIds.size(), placementFailureTaskIds);

                    Set<String> scalablePlacementFailureTaskIds = filterOutTaskIdsForScaling(placementFailureTaskIds, allJobs, allTasks, tierResourceDimension);
                    logger.info("{} had the scalable placement failures({}): {}", tier, scalablePlacementFailureTaskIds.size(), scalablePlacementFailureTaskIds);
                    potentialTaskIdsForScaleUp.addAll(scalablePlacementFailureTaskIds);

                    if (agentCountToScaleUp > 0 || !scalablePlacementFailureTaskIds.isEmpty()) {
                        usedScaleUpCooldown = true;
                    }
                }

                Set<String> tasksPastSlo = getTasksPastSlo(failedTaskIds, allTasks, now, tierConfiguration.getTaskSloMs());
                Set<String> scalableTasksPastSlo = filterOutTaskIdsForScaling(tasksPastSlo, allJobs, allTasks, tierResourceDimension);
                tierAutoScalerExecution.getTotalTasksPastSloGauge().set(scalableTasksPastSlo.size());
                logger.info("{} had tasks past slo({}): {}", tier, scalableTasksPastSlo.size(), scalableTasksPastSlo);
                potentialTaskIdsForScaleUp.addAll(scalableTasksPastSlo);

                Set<String> taskIdsForScaleUp = new HashSet<>();
                for (String taskId : potentialTaskIdsForScaleUp) {
                    boolean previouslyScaledFor = taskIdsForPreviousScaleUps.getIfPresent(taskId) != null;
                    if (!previouslyScaledFor) {
                        taskIdsForScaleUp.add(taskId);
                        taskIdsForPreviousScaleUps.put(taskId, taskId);
                    }
                }

                tierAutoScalerExecution.getTotalTasksForScaleUpGauge().set(taskIdsForScaleUp.size());
                logger.info("{} had tasks to scale up({}): {}", tier, taskIdsForScaleUp.size(), taskIdsForScaleUp);

                int agentScaleUpCountByDominantResource = calculateAgentScaleUpCountByDominantResource(taskIdsForScaleUp, allJobs, allTasks, tierResourceDimension);
                logger.info("{} needs {} instances based on dominant resource", tier, agentScaleUpCountByDominantResource);

                agentCountToScaleUp += agentScaleUpCountByDominantResource;
                logger.info("{} needs {} instances", tier, agentCountToScaleUp);
                tierAutoScalerExecution.getTotalAgentsToScaleUpGauge().set(agentCountToScaleUp);
                boolean scalingUp = false;

                if (agentCountToScaleUp > 0) {
                    long maxTokensToTake = Math.min(SCALE_UP_TOKEN_BUCKET_CAPACITY, agentCountToScaleUp);
                    Optional<Pair<Long, ImmutableTokenBucket>> takeOpt = tierAutoScalerExecution.getLastScaleUpTokenBucket().tryTake(1, maxTokensToTake);
                    if (takeOpt.isPresent()) {
                        Pair<Long, ImmutableTokenBucket> takePair = takeOpt.get();
                        tierAutoScalerExecution.setLastScaleUpTokenBucket(takePair.getRight());
                        long tokensAvailable = takePair.getLeft();
                        Pair<Integer, Completable> scaleUpPair = createScaleUpCompletable(activeScalableInstanceGroupsForTier, (int) tokensAvailable);
                        Integer agentCountBeingScaled = scaleUpPair.getLeft();
                        tierAutoScalerExecution.getTotalAgentsBeingScaledUpGauge().set(agentCountBeingScaled);
                        if (agentCountBeingScaled > 0) {
                            actions.add(scaleUpPair.getRight());
                            logger.info("Attempting to scale up {} tier by {} agent instances", tier, agentCountBeingScaled);
                            scalingUp = true;
                            if (usedScaleUpCooldown) {
                                tierAutoScalerExecution.getLastScaleUp().set(clock.wallTime());
                            }
                        }
                    }
                }

                if (!scalingUp && hasTimeElapsed(tierAutoScalerExecution.getLastScaleDown().get(), now, tierConfiguration.getScaleDownCoolDownMs())) {
                    int agentCountToScaleDown = 0;
                    int maxIdleForTier = tierConfiguration.getMaxIdle();
                    if (idleInstancesForTier.size() > maxIdleForTier) {
                        int instancesNotNeededForMaxIdle = idleInstancesForTier.size() - maxIdleForTier;
                        logger.info("{} can remove {} instances to satisfy max idle {}", tier, instancesNotNeededForMaxIdle, maxIdleForTier);
                        agentCountToScaleDown += instancesNotNeededForMaxIdle;
                    }

                    tierAutoScalerExecution.getTotalAgentsToScaleDownGauge().set(agentCountToScaleDown);

                    if (agentCountToScaleDown > 0) {
                        long maxTokensToTake = Math.min(SCALE_DOWN_TOKEN_BUCKET_CAPACITY, agentCountToScaleDown);
                        Optional<Pair<Long, ImmutableTokenBucket>> takeOpt = tierAutoScalerExecution.getLastScaleDownTokenBucket().tryTake(1, maxTokensToTake);
                        if (takeOpt.isPresent()) {
                            Pair<Long, ImmutableTokenBucket> takePair = takeOpt.get();
                            tierAutoScalerExecution.setLastScaleDownTokenBucket(takePair.getRight());
                            long tokensAvailable = takePair.getLeft();
                            Pair<Integer, Completable> scaleDownPair = createSetRemovableOverrideStatusesCompletable(idleInstancesForTier,
                                    activeScalableInstanceGroupsForTier, instancesForActiveInstanceGroupsById, (int) tokensAvailable);
                            Integer agentCountBeingScaledDown = scaleDownPair.getLeft();
                            tierAutoScalerExecution.getTotalAgentsBeingScaledDownGauge().set(agentCountBeingScaledDown);
                            if (agentCountBeingScaledDown > 0) {
                                actions.add(scaleDownPair.getRight());
                                logger.info("Attempting to scale down {} tier by {} agent instances", tier, agentCountBeingScaledDown);
                                tierAutoScalerExecution.getLastScaleDown().set(clock.wallTime());
                            }
                        }
                    }
                }
                logger.info("Finishing scaling actions for tier: {}", tier);
            }

            List<AgentInstance> removableInstancesPastElapsedTime = getRemovableInstancesPastElapsedTime(instancesForActiveInstanceGroups,
                    now, configuration.getAgentInstanceRemovableTimeoutMs());
            logger.info("Removable instances past elapsed time({}): {}", removableInstancesPastElapsedTime.size(), removableInstancesPastElapsedTime);

            if (!removableInstancesPastElapsedTime.isEmpty()) {
                actions.add(createResetOverrideStatusesCompletable(removableInstancesPastElapsedTime));
                logger.info("Resetting agent instances({}): {}", removableInstancesPastElapsedTime.size(), removableInstancesPastElapsedTime);
            }

            return actions.isEmpty() ? Completable.complete() : Completable.concat(actions);
        }).doOnCompleted(() -> logger.debug("Completed scaling agents"))
                .timeout(CLUSTER_AGENT_AUTO_SCALE_COMPLETABLE_TIMEOUT_MS, TimeUnit.MILLISECONDS);
    }

}

class ClusterAgentAutoScalerTest {

    private final TestScheduler testScheduler;
    private final TitusRuntime titusRuntime;
    private final ClusterOperationsConfiguration configuration;
    private final AgentManagementService agentManagementService;
    private final V3JobOperations v3JobOperations;
    @SuppressWarnings("unchecked")
    private final SchedulingService<StubTaskRequest> schedulingService;

    @Test
    public void testScaleUpForNonPrimaryInstanceType() {

        AgentInstanceGroup instanceGroup = createPartition("instanceGroup1", InstanceGroupLifecycleState.Active, "m4.16xlarge", 0, 0, 10);
        when(agentManagementService.getInstanceGroups()).thenReturn(singletonList(instanceGroup));

        when(agentManagementService.getAgentInstances("instanceGroup1")).thenReturn(Collections.emptyList());
        when(agentManagementService.scaleUp(eq("instanceGroup1"), anyInt())).thenReturn(Completable.complete());

        testScheduler.advanceTimeBy(6, TimeUnit.MINUTES);

        ClusterAgentAutoScaler clusterAgentAutoScaler = new ClusterAgentAutoScaler(titusRuntime, configuration,
                agentManagementService, v3JobOperations, schedulingService, testScheduler);

        clusterAgentAutoScaler.doAgentScaling().await();

        verify(agentManagementService, times(0)).scaleUp(anyString(), anyInt());
    }

}