/**
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.aurora.scheduler.updater;

import java.util.Set;

import com.google.common.base.Optional;
import com.google.common.collect.ImmutableSet;
import com.google.inject.Inject;

import org.apache.aurora.scheduler.base.TaskGroupKey;
import org.apache.aurora.scheduler.preemptor.BiCache;
import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

/**
 * Reserves agents for instances being updated. Multiple instance keys can be registered against
 * a single agent.
 */
public interface UpdateAgentReserver {

  /**
   * Reserves the agent id for the given key. Should behave like a multi-map under the hood.
   *
   * @param agentId The agent id to reserve.
   * @param key The instance key that will use the reservation.
   */
  void reserve(String agentId, IInstanceKey key);

  /**
   * Releases the reservation on an agent id for the given key.
   *
   * @param agentId The agent id to release the reservation on.
   * @param key The instance key that should be removed.
   */
  void release(String agentId, IInstanceKey key);

  /**
   * Returns the agent id associated with the given instance key.
   *
   * @param key The instance key to look up.
   * @return An optional agent id string.
   */
  Optional<String> getAgent(IInstanceKey key);

  /**
   * Get all reservations for a given agent id. Useful for skipping over the agent between the
   * reserve/release window.
   *
   * @param agentId The agent id to look up reservations for.
   * @return A set of keys reserved for that agent.
   */
  Set<IInstanceKey> getReservations(String agentId);

  /**
   * Check if the agent reserver has any reservations for the provided key.
   *
   * @param groupKey The key to check.
   * @return True if there are reservations against any instances in that key.
   */
  boolean hasReservations(TaskGroupKey groupKey);

  /**
   * Implementation of the update reserver backed by a BiCache (the same mechanism we use for
   * preemption). This means it will expire reservations that haven't been explicitly released
   * after the configured timeout.
   */
  class UpdateAgentReserverImpl implements UpdateAgentReserver {
    private static final Logger LOG = LoggerFactory.getLogger(UpdateAgentReserverImpl.class);

    private final BiCache<IInstanceKey, String> cache;

    @Inject
    UpdateAgentReserverImpl(BiCache<IInstanceKey, String> cache) {
      this.cache = requireNonNull(cache);
    }

    public void reserve(String agentId, IInstanceKey key) {
      LOG.info("Reserving {} for {}", agentId, key);
      cache.put(key, agentId);
    }

    public void release(String agentId, IInstanceKey key) {
      LOG.info("Releasing reservation on {} for {}", agentId, key);
      cache.remove(key, agentId);
    }

    public Set<IInstanceKey> getReservations(String agentId) {
      return cache.getByValue(agentId);
    }

    @Override
    public boolean hasReservations(TaskGroupKey groupKey) {
      return cache.asMap().entrySet().stream()
          .filter(entry -> entry.getKey().getJobKey().equals(groupKey.getTask().getJob()))
          .findFirst()
          .isPresent();
    }

    @Override
    public Optional<String> getAgent(IInstanceKey key) {
      return cache.get(key);
    }
  }

  /**
   * Used to effectively disable reservations.
   */
  class NullAgentReserver implements UpdateAgentReserver {
    @Override
    public void reserve(String agentId, IInstanceKey key) {
      // noop
    }

    @Override
    public void release(String agentId, IInstanceKey key) {
      // noop
    }

    @Override
    public Optional<String> getAgent(IInstanceKey key) {
      return Optional.absent();
    }

    @Override
    public Set<IInstanceKey> getReservations(String agentId) {
      return ImmutableSet.of();
    }

    @Override
    public boolean hasReservations(TaskGroupKey groupKey) {
      return false;
    }
  }
}