/**
 * Copyright 2019 Airbnb. Licensed under Apache-2.0. See LICENSE in the project root for license
 * information.
 */
package com.airbnb.dynein.scheduler.partition;

import com.airbnb.dynein.scheduler.Constants;
import io.fabric8.kubernetes.api.model.Pod;
import io.fabric8.kubernetes.api.model.PodList;
import io.fabric8.kubernetes.api.model.apps.ReplicaSet;
import io.fabric8.kubernetes.client.AppsAPIGroupClient;
import io.fabric8.kubernetes.client.DefaultKubernetesClient;
import io.fabric8.kubernetes.client.NamespacedKubernetesClient;
import io.fabric8.kubernetes.client.Watch;
import io.fabric8.kubernetes.client.Watcher;
import io.fabric8.kubernetes.client.dsl.FilterWatchListDeletable;
import java.text.Collator;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import javax.inject.Inject;
import javax.inject.Named;
import lombok.AllArgsConstructor;

@AllArgsConstructor(onConstructor = @__(@Inject))
public class K8SConsecutiveAllocationPolicy implements PartitionPolicy {
  private static final String K8S_CONDITION_READY = "Ready";
  private static final String K8S_CONDITION_TRUE = "True";
  private final DefaultKubernetesClient kubernetesClient;
  private final AppsAPIGroupClient appsClient;

  @Named(Constants.MAX_PARTITIONS)
  private final Integer numPartitions;

  public List<Integer> getPartitions() {
    String podName = System.getenv("K8S_POD_NAME");
    String namespace = System.getenv("K8S_NAMESPACE");
    NamespacedKubernetesClient namespacedClient = kubernetesClient.inNamespace(namespace);
    ReplicaSet replicaSet;

    Pod pod = namespacedClient.pods().withName(podName).get();
    String replicaSetName = pod.getMetadata().getOwnerReferences().get(0).getName();
    replicaSet = appsClient.replicaSets().inNamespace(namespace).withName(replicaSetName).get();
    FilterWatchListDeletable<Pod, PodList, Boolean, Watch, Watcher<Pod>> deploymentPods =
        namespacedClient.pods().withLabelSelector(replicaSet.getSpec().getSelector());

    List<String> activePods =
        deploymentPods
            .list()
            .getItems()
            .stream()
            .filter(
                isActive ->
                    isActive
                        .getStatus()
                        .getConditions()
                        .stream()
                        .anyMatch(
                            condition ->
                                condition.getType().equals(K8S_CONDITION_READY)
                                    && condition.getStatus().equals(K8S_CONDITION_TRUE)))
            .map(it -> it.getMetadata().getName())
            .sorted(Collator.getInstance())
            .collect(Collectors.toList());

    int podIndex = activePods.indexOf(podName);
    int numPods = activePods.size();

    if (numPods == 0 || podIndex == -1) {
      return new ArrayList<>();
    }
    List<Integer> partitions = new ArrayList<>();
    int split = numPartitions / numPods;
    int start = podIndex * split;
    int end = (podIndex == numPods - 1) ? numPartitions - 1 : ((podIndex + 1) * split) - 1;
    for (int i = start; i <= end; i++) {
      partitions.add(i);
    }
    return partitions;
  }
}