package com.uas.search.grouping;

import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.SimpleCollector;

public class DistinctGroupCollector extends SimpleCollector {

	private final String groupField;

	private final Set<String> fieldsToLoad;

	private LeafReader reader;

	private Set<Long> values = new HashSet<Long>();

	public DistinctGroupCollector(String groupField) {
		this.groupField = groupField;
		this.fieldsToLoad = new HashSet<String>();
		this.fieldsToLoad.add(groupField);
	}

	@Override
	public boolean needsScores() {
		return false;
	}

	@Override
	protected void doSetNextReader(LeafReaderContext context) throws IOException {
		super.doSetNextReader(context);
		reader = context.reader();
	}

	@Override
	public void collect(int doc) throws IOException {
		if (null != reader) {
			String value = reader.document(doc, fieldsToLoad).get(groupField);
			if (null != value) {
				values.add(Long.parseLong(value));
			}
		}
	}

	public Set<Long> getValues() {
		return values;
	}

}
