package com.uas.search.grouping;

import com.uas.search.util.ArrayUtils;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.LeafReader;
import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.search.SimpleCollector;

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

public class DistinctGroupCollector extends SimpleCollector {

	private final String groupField;

	private final CollectField[] collectFields;

	private final Set<String> fieldsToLoad = new HashSet<>();

	private LeafReader reader;

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

	private Set<Map<String, Object>> collectValues = new HashSet<>();

	public DistinctGroupCollector(String groupField, CollectField... collectFields) {
		this.groupField = groupField;
		this.collectFields = collectFields;
		this.fieldsToLoad.add(groupField);
		if(collectFields!=null){
			for(CollectField collectField : collectFields){
				fieldsToLoad.add(collectField.getFieldName());
			}
		}
	}

	@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 (reader != null) {
			Document document = reader.document(doc, fieldsToLoad);
			String value = document.get(groupField);
			if (value != null) {
				values.add(Long.parseLong(value));
				// 统计指定信息
				if(!ArrayUtils.isEmpty(collectFields)){
					Map<String, Object> collectValue = new HashMap<>();
					for(CollectField collectField : collectFields){
						String v = document.get(collectField.getFieldName());
						if(v != null){
							collectValue.put(collectField.getDisplayName(), v);
						}
					}
					collectValues.add(collectValue);
				}
			}
		}
	}

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

	public Set<Map<String, Object>> getCollectValues(){
		return collectValues;
	}

	/**
	 * 用于统计的字段
	 */
	public static class CollectField {
		private String fieldName;
		private String displayName;

		public CollectField(String fieldName, String displayName) {
			this.fieldName = fieldName;
			this.displayName = displayName;
		}

		public String getFieldName() {
			return fieldName;
		}

		public String getDisplayName() {
			return displayName;
		}
	}
}
