package org.getopt.luke;

import org.apache.lucene.search.TopDocCollector;
import org.apache.lucene.search.TopDocs;

public class CountLimitedHitCollector extends LimitedHitCollector {
  private long maxSize;
  private long count;
  private int lastDoc;
  private TopDocCollector tdc;
  private TopDocs topDocs = null;
  
  public CountLimitedHitCollector(long maxSize) {
    this.maxSize = maxSize;
    count = 0;
    tdc = new TopDocCollector((int)maxSize);
  }

  @Override
  public long limitSize() {
    return maxSize;
  }

  @Override
  public int limitType() {
    return TYPE_SIZE;
  }

  /* (non-Javadoc)
   * @see org.getopt.luke.AllHitsCollector#collect(int, float)
   */
  @Override
  public void collect(int doc, float score) {
    count++;
    if (count > maxSize) {
      count--;
      throw new LimitedException(TYPE_SIZE, maxSize, count, lastDoc);
    }
    lastDoc = doc;
    tdc.collect(doc, score);
  }

  /* (non-Javadoc)
   * @see org.getopt.luke.AccessibleHitCollector#getDocId(int)
   */
  @Override
  public int getDocId(int pos) {
    if (topDocs == null) {
      topDocs = tdc.topDocs();
    }
    return topDocs.scoreDocs[pos].doc;
  }

  /* (non-Javadoc)
   * @see org.getopt.luke.AccessibleHitCollector#getScore(int)
   */
  @Override
  public float getScore(int pos) {
    if (topDocs == null) {
      topDocs = tdc.topDocs();
    }
    return topDocs.scoreDocs[pos].score;
  }

  /* (non-Javadoc)
   * @see org.getopt.luke.AccessibleHitCollector#getTotalHits()
   */
  @Override
  public int getTotalHits() {
    return (int)count;
  }
}
