package org.apache.hadoop.yarn.server.nodemanager.containermanager.localizer;

import java.util.Comparator;
import java.util.Iterator;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

import org.apache.hadoop.yarn.server.nodemanager.DeletionService;

public class ResourceRetentionSet {

  private long delSize;
  private long currentSize;
  private final long targetSize;
  private final DeletionService delService;
  private final SortedMap<LocalizedResource,LocalResourcesTracker> retain;

  ResourceRetentionSet(DeletionService delService, long targetSize) {
    this(delService, targetSize, new LRUComparator());
  }

  ResourceRetentionSet(DeletionService delService, long targetSize,
      Comparator<? super LocalizedResource> cmp) {
    this(delService, targetSize,
        new TreeMap<LocalizedResource,LocalResourcesTracker>(cmp));
  }

  ResourceRetentionSet(DeletionService delService, long targetSize,
      SortedMap<LocalizedResource,LocalResourcesTracker> retain) {
    this.retain = retain;
    this.delService = delService;
    this.targetSize = targetSize;
  }

  public void addResources(LocalResourcesTracker newTracker) {
    for (LocalizedResource resource : newTracker) {
      currentSize += resource.getSize();
      if (resource.getRefCount() > 0) {
        // always retain resources in use
        continue;
      }
      retain.put(resource, newTracker);
    }
    for (Iterator<Map.Entry<LocalizedResource,LocalResourcesTracker>> i =
           retain.entrySet().iterator();
         currentSize - delSize > targetSize && i.hasNext();) {
      Map.Entry<LocalizedResource,LocalResourcesTracker> rsrc = i.next();
      LocalizedResource resource = rsrc.getKey();
      LocalResourcesTracker tracker = rsrc.getValue();
      if (tracker.remove(resource, delService)) {
        delSize += resource.getSize();
        i.remove();
      }
    }
  }

  @Override
  public String toString() {
    StringBuilder sb = new StringBuilder();
    sb.append("Cache: ").append(currentSize).append(", ");
    sb.append("Deleted: ").append(delSize);
    return sb.toString();
  }

  static class LRUComparator implements Comparator<LocalizedResource> {
    public int compare(LocalizedResource r1, LocalizedResource r2) {
      long ret = r1.getTimestamp() - r2.getTimestamp();
      if (0 == ret) {
        return System.identityHashCode(r1) - System.identityHashCode(r2);
      }
      return ret > 0 ? 1 : -1;
    }
    public boolean equals(Object other) {
      return this == other;
    }
  }
}