/*
 * Copyright 2014-2018 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.springframework.data.mongodb.test.util;

import java.net.UnknownHostException;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;

import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.MongoClient;

/**
 * {@link CleanMongoDB} is a junit {@link TestRule} implementation to be used as for wiping data from MongoDB instance.
 * MongoDB specific system databases like {@literal admin} and {@literal local} remain untouched. The rule will apply
 * <strong>after</strong> the base {@link Statement}. <br />
 * Use as {@link org.junit.ClassRule} to wipe data after finishing all tests within a class or as {@link org.junit.Rule}
 * to do so after each {@link org.junit.Test}.
 *
 * @author Christoph Strobl
 * @since 1.6
 */
public class CleanMongoDB implements TestRule {

	private static final Logger LOGGER = LoggerFactory.getLogger(CleanMongoDB.class);

	/**
	 * Defines contents of MongoDB.
	 */
	public enum Struct {
		DATABASE, COLLECTION, INDEX;
	}

	@SuppressWarnings("serial")//
	private Set<String> preserveDatabases = new HashSet<String>() {
		{
			add("admin");
			add("local");
		}
	};

	private Set<String> dbNames = new HashSet<String>();
	private Set<String> collectionNames = new HashSet<String>();
	private Set<Struct> types = new HashSet<CleanMongoDB.Struct>();
	private MongoClient client;

	/**
	 * Create new instance using an internal {@link MongoClient}.
	 */
	public CleanMongoDB() {
		this(null);
	}

	/**
	 * Create new instance using an internal {@link MongoClient} connecting to specified instance running at host:port.
	 *
	 * @param host
	 * @param port
	 * @throws UnknownHostException
	 */
	public CleanMongoDB(String host, int port) throws UnknownHostException {
		this(new MongoClient(host, port));
	}

	/**
	 * Create new instance using the given client.
	 *
	 * @param client
	 */
	public CleanMongoDB(MongoClient client) {
		this.client = client;
	}

	/**
	 * Removes everything by dropping every single {@link DB}.
	 *
	 * @return
	 */
	public static CleanMongoDB everything() {

		CleanMongoDB cleanMongoDB = new CleanMongoDB();
		cleanMongoDB.clean(Struct.DATABASE);
		return cleanMongoDB;
	}

	/**
	 * Removes everything from the databases with given name by dropping the according {@link DB}.
	 *
	 * @param dbNames
	 * @return
	 */
	public static CleanMongoDB databases(String... dbNames) {

		CleanMongoDB cleanMongoDB = new CleanMongoDB();
		cleanMongoDB.clean(Struct.DATABASE);
		cleanMongoDB.useDatabases(dbNames);
		return cleanMongoDB;
	}

	/**
	 * Drops the {@link DBCollection} with given names from every single {@link DB} containing them.
	 *
	 * @param collectionNames
	 * @return
	 */
	public static CleanMongoDB collections(String... collectionNames) {
		return collections("", Arrays.asList(collectionNames));
	}

	/**
	 * Drops the {@link DBCollection} with given names from the named {@link DB}.
	 *
	 * @param dbName
	 * @param collectionNames
	 * @return
	 */
	public static CleanMongoDB collections(String dbName, Collection<String> collectionNames) {

		CleanMongoDB cleanMongoDB = new CleanMongoDB();
		cleanMongoDB.clean(Struct.COLLECTION);
		cleanMongoDB.useCollections(dbName, collectionNames);
		return cleanMongoDB;
	}

	/**
	 * Drops all index structures from every single {@link DBCollection}.
	 *
	 * @return
	 */
	public static CleanMongoDB indexes() {
		return indexes(Collections.<String> emptySet());
	}

	/**
	 * Drops all index structures from every single {@link DBCollection}.
	 *
	 * @param collectionNames
	 * @return
	 */
	public static CleanMongoDB indexes(Collection<String> collectionNames) {

		CleanMongoDB cleanMongoDB = new CleanMongoDB();
		cleanMongoDB.clean(Struct.INDEX);
		cleanMongoDB.useCollections(collectionNames);
		return cleanMongoDB;
	}

	/**
	 * Define {@link Struct} to be cleaned.
	 *
	 * @param types
	 * @return
	 */
	public CleanMongoDB clean(Struct... types) {

		this.types.addAll(Arrays.asList(types));
		return this;
	}

	/**
	 * Defines the {@link DB}s to be used. <br />
	 * Impact along with {@link CleanMongoDB#clean(Struct...)}:
	 * <ul>
	 * <li>{@link Struct#DATABASE}: Forces drop of named databases.</li>
	 * <li>{@link Struct#COLLECTION}: Forces drop of collections within named databases.</li>
	 * <li>{@link Struct#INDEX}: Removes index within collections of named databases.</li>
	 * </ul>
	 *
	 * @param dbNames
	 * @return
	 */
	public CleanMongoDB useDatabases(String... dbNames) {

		this.dbNames.addAll(Arrays.asList(dbNames));
		return this;
	}

	/**
	 * Excludes the given {@link DB}s from being processed.
	 *
	 * @param dbNames
	 * @return
	 */
	public CleanMongoDB preserveDatabases(String... dbNames) {
		this.preserveDatabases.addAll(Arrays.asList(dbNames));
		return this;
	}

	/**
	 * Defines the {@link DBCollection}s to be used. <br />
	 * Impact along with {@link CleanMongoDB#clean(Struct...)}:
	 * <ul>
	 * <li>{@link Struct#COLLECTION}: Forces drop of named collections.</li>
	 * <li>{@link Struct#INDEX}: Removes index within named collections.</li>
	 * </ul>
	 *
	 * @param collectionNames
	 * @return
	 */
	public CleanMongoDB useCollections(String... collectionNames) {
		return useCollections(Arrays.asList(collectionNames));
	}

	private CleanMongoDB useCollections(Collection<String> collectionNames) {
		return useCollections("", collectionNames);
	}

	/**
	 * Defines the {@link DBCollection}s and {@link DB} to be used. <br />
	 * Impact along with {@link CleanMongoDB#clean(Struct...)}:
	 * <ul>
	 * <li>{@link Struct#COLLECTION}: Forces drop of named collections in given db.</li>
	 * <li>{@link Struct#INDEX}: Removes index within named collections in given db.</li>
	 * </ul>
	 *
	 * @param collectionNames
	 * @return
	 */
	public CleanMongoDB useCollections(String db, Collection<String> collectionNames) {

		if (StringUtils.hasText(db)) {
			this.dbNames.add(db);
		}

		if (!CollectionUtils.isEmpty(collectionNames)) {
			this.collectionNames.addAll(collectionNames);
		}
		return this;
	}

	Statement apply() {
		return apply(null, null);
	}

	/*
	 * (non-Javadoc)
	 * @see org.junit.rules.TestRule#apply(org.junit.runners.model.Statement, org.junit.runner.Description)
	 */
	public Statement apply(Statement base, Description description) {
		return new MongoCleanStatement(base);
	}

	private void doClean() {

		Collection<String> dbNamesToUse = initDbNames();

		for (String dbName : dbNamesToUse) {

			if (isPreserved(dbName) || dropDbIfRequired(dbName)) {
				continue;
			}

			DB db = client.getDB(dbName);
			dropCollectionsOrIndexIfRequried(db, initCollectionNames(db));
		}
	}

	private boolean dropDbIfRequired(String dbName) {

		if (!types.contains(Struct.DATABASE)) {
			return false;
		}

		client.dropDatabase(dbName);
		LOGGER.debug("Dropping DB '{}'. ", dbName);
		return true;
	}

	private void dropCollectionsOrIndexIfRequried(DB db, Collection<String> collectionsToUse) {

		for (String collectionName : collectionsToUse) {

			if (db.collectionExists(collectionName)) {

				DBCollection collection = db.getCollectionFromString(collectionName);
				if (collection != null) {

					if (types.contains(Struct.COLLECTION)) {
						collection.drop();
						LOGGER.debug("Dropping collection '{}' for DB '{}'. ", collectionName, db.getName());
					} else if (types.contains(Struct.INDEX)) {
						collection.dropIndexes();
						LOGGER.debug("Dropping indexes in collection '{}' for DB '{}'. ", collectionName, db.getName());
					}
				}
			}
		}
	}

	private boolean isPreserved(String dbName) {
		return preserveDatabases.contains(dbName.toLowerCase());
	}

	private Collection<String> initDbNames() {

		Collection<String> dbNamesToUse = dbNames;
		if (dbNamesToUse.isEmpty()) {
			dbNamesToUse = client.getDatabaseNames();
		}
		return dbNamesToUse;
	}

	private Collection<String> initCollectionNames(DB db) {

		Collection<String> collectionsToUse = collectionNames;
		if (CollectionUtils.isEmpty(collectionsToUse)) {
			collectionsToUse = db.getCollectionNames();
		}
		return collectionsToUse;
	}

	/**
	 * @author Christoph Strobl
	 * @since 1.6
	 */
	private class MongoCleanStatement extends Statement {

		private final Statement base;

		public MongoCleanStatement(Statement base) {
			this.base = base;
		}

		@Override
		public void evaluate() throws Throwable {

			if (base != null) {
				base.evaluate();
			}

			boolean isInternal = false;
			if (client == null) {
				client = new MongoClient();
				isInternal = true;
			}

			doClean();

			if (isInternal) {
				client.close();
				client = null;
			}
		}
	}
}
