/* 
 * Copyright 1999,2004 The Apache Software Foundation.
 * 
 * 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.
 */

/**
 * 
 * Simple Token Bucket implementation to limit traffic flow. The
 * <code>TokenBucket</code> adds a token to the bucket at the rate of
 * <code>fillAmount</code> tokens every <code>fillInterval</code> seconds up
 * to <code>maxTokens</code>.
 * 
 * 
 * @author Chad LaVigne
 * 
 */
public class TokenBucket {
  /**
   * Number of tokens to add to the bucket, <code>fillAmount</code> tokens are
   * added to the bucket every <code>fillInterval</code> seconds.
   */
  private long fillAmount;

  /**
   * Interval at which to add tokens to the bucket, <code>fillAmount</code>
   * tokens are added to the bucket every <code>fillInterval</code> seconds.
   */
  private long fillInterval;

  /**
   * The maximum number of tokens allowed in the bucket. This becomes the peak
   * traffic burst allowed by the token bucket in <code>fillInterval</code>
   * seconds.
   */
  private long maxTokens;

  private long currentNumberOfTokens;

  private long lastTokenRemovedTime;

  public TokenBucket(long fillAmount, long fillInterval, long maxTokens) {
    this.fillAmount = fillAmount;
    this.fillInterval = fillInterval;
    this.maxTokens = maxTokens;
    this.currentNumberOfTokens = maxTokens;
    lastTokenRemovedTime = System.currentTimeMillis();
  }

  /**
   * Method to get a token from the bucket. If the bucket is not empty a token
   * is removed.
   * 
   * @return
   */
  public synchronized boolean getToken() {
    replaceTokens();
    boolean isEmpty = currentNumberOfTokens <= 0;

    if (!isEmpty) {
      currentNumberOfTokens--;
      lastTokenRemovedTime = System.currentTimeMillis();
    }

    return isEmpty;
  }

  private void replaceTokens() {
    long currentTime = System.currentTimeMillis();
    long secondsSinceLastFill = (long) ((currentTime / 1000) - (lastTokenRemovedTime / 1000));

    if (secondsSinceLastFill >= fillInterval) {
      long numberOfTokensToAdd = (secondsSinceLastFill / fillInterval) * fillAmount;
      currentNumberOfTokens = currentNumberOfTokens + numberOfTokensToAdd > maxTokens ? maxTokens : currentNumberOfTokens + numberOfTokensToAdd;
    }
  }
}