package eu.dnetlib.data.mdstore.modular.mongodb;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Pattern;

import com.google.common.base.Function;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.mongodb.BasicDBObject;
import com.mongodb.DBObject;
import com.mongodb.MongoClient;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.Sorts;
import eu.dnetlib.enabling.resultset.ResultSet;
import eu.dnetlib.enabling.resultset.ResultSetAware;
import eu.dnetlib.enabling.resultset.ResultSetListener;
import eu.dnetlib.miscutils.maps.ConcurrentSizedMap;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.bson.BsonDocument;
import org.bson.conversions.Bson;

import static com.mongodb.client.model.Filters.*;

public class MongoResultSetListener implements ResultSetListener, ResultSetAware {

	private static final Log log = LogFactory.getLog(MongoResultSetListener.class);

	private ConcurrentSizedMap<Integer, String> lastKeys = new ConcurrentSizedMap<>();
	private Bson sortByIdAsc = Sorts.orderBy(Sorts.ascending("id"));

	private Function<DBObject, String> serializer;
	private MongoCollection<DBObject> collection;
	private Bson query;

	public MongoResultSetListener(final MongoCollection<DBObject> collection, final Long from, final Long until, final Pattern filter, final Function<DBObject, String> serializer) {
		this.collection = collection;
		this.serializer = serializer;
		this.query = query(from, until, filter);
		log.debug("Query on mongo: "+this.query.toBsonDocument(BsonDocument.class, MongoClient.getDefaultCodecRegistry()));
	}

	@Override
	public List<String> getResult(final int fromPosition, final int toPosition) {

		ArrayList<DBObject> page = null;

		String lastKey = lastKeys.get(fromPosition);
		if (lastKey != null) {
			page = continueFrom(lastKey, (toPosition - fromPosition) + 1);
		} else {
			page = fetchNew(fromPosition - 1, (toPosition - fromPosition) + 1);
		}

		if (!page.isEmpty()) {
			DBObject last = page.get(page.size() - 1);
			lastKeys.put(toPosition + 1, (String) last.get("id"));
		}

		if (log.isDebugEnabled()) {
			log.info(String.format("got %s records from %s to %s", page.size(), fromPosition, toPosition));
		}

		return Lists.newArrayList(Iterables.transform(page, serializer));
	}

	private ArrayList<DBObject> fetchNew(final int from, final int size) {
		final FindIterable<DBObject> it = collection.find(query).batchSize(size);
		return Lists.newArrayList(it.sort(sortByIdAsc).skip(from).limit(size));
	}

	private ArrayList<DBObject> continueFrom(final String lastKey, final int size) {
		if (log.isDebugEnabled()) {
			log.debug("trying to continue from previous key: " + lastKey);
		}
		final Bson q = and(query, gt("id", lastKey));
		final FindIterable<DBObject> it = collection.find(q).batchSize(size).sort(sortByIdAsc).limit(size);
		return Lists.newArrayList(it);
	}

	private Bson query(final Long from, final Long until, final Pattern pattern) {
		final Bson dateFilter = dateQuery(from, until);
		final Bson regexFilter = regexQuery(pattern);
		if (dateFilter != null & regexFilter != null) {
			return and(dateFilter, regexFilter);
		} else if (dateFilter != null) {
			return dateFilter;
		} else if (regexFilter != null) {
			return regexFilter;
		}
		return new BasicDBObject();
	}

	private Bson dateQuery(final Long from, final Long until) {
		if (from != null & until != null) {
			return and(gt("timestamp", from), lt("timestamp", until));
		}
		if (from != null) {
			return gt("timestamp", from);
		}
		if (until != null) {
			return lt("timestamp", until);
		}
		return null;
	}

	private Bson regexQuery(final Pattern pattern) {
		if (pattern != null) {
			return regex("body", pattern);
		}
		return null;
	}

	@Override
	public int getSize() {
		return (int) collection.count(query);
	}

	@Override
	public void setResultSet(final ResultSet resultSet) {
		resultSet.close();
	}

}
