D-Net

package eu.dnetlib.data.mapreduce.hbase.oai.utils;

import java.text.Normalizer;

import java.util.List;

import com.google.common.collect.Iterables;

import com.google.common.collect.Lists;

import com.mongodb.BasicDBObject;

import com.mongodb.BasicDBObjectBuilder;

import com.mongodb.DBObject;

import com.mongodb.MongoClient;

import com.mongodb.client.FindIterable;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.model.Filters;

import com.mongodb.client.model.FindOneAndReplaceOptions;

import com.mongodb.client.model.FindOneAndUpdateOptions;

import com.mongodb.client.model.IndexOptions;

import org.apache.commons.lang.StringEscapeUtils;

import org.apache.commons.lang.StringUtils;

import org.bson.conversions.Bson;

public class MongoSetCollection implements SetCollection {

	public static String DEFAULT_SET = "OTHER";

	private MongoClient publisherMongoClient;

	private String setCollection = "sets";

	private String setCountCollection = "setsCount";

	public MongoSetCollection(MongoClient c){

		this.publisherMongoClient = c;

	}

	public void ensureIndexes(final String dbName) {

		this.ensureIndexesOnSets(dbName);

		this.ensureIndexesOnCount(dbName);

	}

	@Override

	public List<SetInfo> getAllSets(final boolean enabledOnly, final String dbName) {

		FindIterable<DBObject> iter = null;

		if (!enabledOnly) {

			iter = this.getSetsCollection(dbName).find();

		} else {

			Bson where = Filters.eq("enabled", true);

			iter = this.getSetsCollection(dbName).find(where);

		}

		return Lists.newArrayList(Iterables.transform(iter, new com.google.common.base.Function<DBObject, SetInfo>() {

			@Override

			public SetInfo apply(final DBObject dbObject) {

				return getSetFromDBObject(dbObject);

			}

		}));

	}

	@Override

	public boolean containSet(final String set, final String dbName) {

		Bson query = Filters.eq("spec", set);

		return this.getSetsCollection(dbName).count(query) != 0;

	}

	@Override

	public boolean containEnabledSet(final String set, final String publisherDBName) {

		Bson query = Filters.and(Filters.eq("spec", set), Filters.eq("enabled", true));

		return this.getSetsCollection(publisherDBName).count(query) != 0;

	}

	@Override

	public String getSetQuery(final String set, final String dbName) {

		Bson query = Filters.eq("spec", set);

		BasicDBObject returnField = new BasicDBObject("query", 1);

		DBObject obj = this.getSetsCollection(dbName).find(query).projection(returnField).first();

		return (String) obj.get("query");

	}

	@Override

	public int count(final String setSpec, final String mdPrefix, final String dbName) {

		Bson query = Filters.and(Filters.eq("spec", setSpec), Filters.eq("mdPrefix", mdPrefix));

		BasicDBObject returnField = new BasicDBObject("count", 1);

		DBObject obj = this.getSetsCountCollection(dbName).find(query).projection(returnField).first();

		if (obj == null) return 0;

		return (Integer) obj.get("count");

	}

	public void updateCounts(final String setSpec, final String mdPrefix, final int count, final String dbName) {

		BasicDBObject countUpdate = new BasicDBObject("$set", new BasicDBObject("count", count));

		Bson query = Filters.and(Filters.eq("spec", setSpec), Filters.eq("mdPrefix", mdPrefix));

		this.getSetsCountCollection(dbName).findOneAndUpdate(query, countUpdate, new FindOneAndUpdateOptions().upsert(true));

	}

	public void upsertSet(final SetInfo setInfo, final boolean fromConfiguration, final String dbName) {

		DBObject obj = this.getObjectFromSet(setInfo);

		obj.put("fromConfiguration", fromConfiguration);

		//this.getSetsCollection(dbName).update(new BasicDBObject("spec", setInfo.getSetSpec()), obj, true, false);

		this.getSetsCollection(dbName).findOneAndReplace(Filters.eq("spec", setInfo.getSetSpec()), obj, new FindOneAndReplaceOptions().upsert(true));

	}

	public String normalizeSetSpec(final String setName) {

		String s = StringEscapeUtils.unescapeXml(setName);

		s = Normalizer.normalize(s, Normalizer.Form.NFD);

		// replace spaces with underscores

		s = s.replaceAll(" ", "_");

		// remove tilde, dots... over letters

		s = s.replaceAll("[\\p{InCombiningDiacriticalMarks}&&[^-_]]", "");

		// change punctuation into an underscore

		s = s.replaceAll("[\\p{Punct}&&[^-_]]", "_");

		// remove all non-word characters

		s = s.replaceAll("[\\W&&[^-_]]", "");

		// Avoiding set '___' generated when we have "strange" set names such as those in cyrillic/ukrain

		// strips _ from the beginning and the end

		String stripped = StringUtils.strip(s, "_ ");

		if (StringUtils.isBlank(stripped)) {

			stripped = DEFAULT_SET;

		}

		return stripped;

	}

	public List<SetInfo> getConfiguredSets(final String dbName) {

		Bson query = Filters.eq("fromConfiguration", true);

		return this.findSets(query, dbName);

	}

	public List<SetInfo> getSetsFromData(final String dbName) {

		Bson query = Filters.eq("fromConfiguration", false);

		return this.findSets(query, dbName);

	}

	public void dropOAISets(final String dbName) {

		this.getSetsCountCollection(dbName).drop();

		this.getSetsCollection(dbName).drop();

	}

	public void dropSet(final String dbName, final String setSpec) {

		Bson query = Filters.eq("spec", setSpec);

		this.getSetsCollection(dbName).deleteMany(query);

		this.getSetsCountCollection(dbName).deleteMany(query);

	}

	public void dropConfigurationSets(final String dbName) {

		this.getSetsCollection(dbName).deleteMany(Filters.eq("fromConfiguration", true));

	}

	protected List<SetInfo> findSets(final Bson query, final String dbName) {

		final FindIterable<DBObject> sets = this.getSetsCollection(dbName).find(query);

		List<SetInfo> res = Lists.newArrayList();

		for (DBObject obj : sets) {

			res.add(this.getSetFromDBObject(obj));

		}

		return res;

	}

	private SetInfo getSetFromDBObject(final DBObject obj) {

		SetInfo setInfo = new SetInfo();

		setInfo.setEnabled((Boolean) obj.get("enabled"));

		setInfo.setQuery((String) obj.get("query"));

		setInfo.setSetDescription((String) obj.get("description"));

		setInfo.setSetName((String) obj.get("name"));

		setInfo.setSetSpec((String) obj.get("spec"));

		return setInfo;

	}

	private DBObject getObjectFromSet(final SetInfo s) {

		DBObject obj = BasicDBObjectBuilder.start("spec", s.getSetSpec()).add("name", s.getSetName()).add("description", s.getSetDescription())

				.add("query", s.getQuery()).add("enabled", s.isEnabled()).get();

		return obj;

	}

	private void ensureIndexesOnSets(final String dbName) {

		this.getSetsCollection(dbName).createIndex(new BasicDBObject("spec", 1), new IndexOptions().background(true));

		this.getSetsCollection(dbName).createIndex(new BasicDBObject("fromConfiguration", 1), new IndexOptions().background(true));

	}

	private void ensureIndexesOnCount(final String dbName) {

		BasicDBObject index = (BasicDBObject) BasicDBObjectBuilder.start("spec", 1).add("mdPrefix", 1).get();

		this.getSetsCountCollection(dbName).createIndex(index, new IndexOptions().background(true));

	}

	public MongoCollection<DBObject> getSetsCollection(final String dbName) {

		return this.getCollection(this.setCollection, dbName);

	}

	public MongoCollection<DBObject> getSetsCountCollection(final String dbName) {

		return this.getCollection(this.setCountCollection, dbName);

	}

	private MongoCollection<DBObject> getCollection(final String collectionName, final String dbName) {

		return publisherMongoClient.getDatabase(dbName).getCollection(collectionName, DBObject.class);

	}

	public String getSetCollection() {

		return setCollection;

	}

	public void setSetCollection(final String setCollection) {

		this.setCollection = setCollection;

	}

	public String getSetCountCollection() {

		return setCountCollection;

	}

	public void setSetCountCollection(final String setCountCollection) {

		this.setCountCollection = setCountCollection;

	}

	public MongoClient getPublisherMongoClient() {

		return publisherMongoClient;

	}

	public void setPublisherMongoClient(final MongoClient publisherMongoClient) {

		this.publisherMongoClient = publisherMongoClient;

	}

}