/*
 * Copyright 2013 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.yarn.test.support;

import java.util.Hashtable;
import java.util.Map.Entry;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.yarn.test.context.YarnCluster;

/**
 * Manager handling running mini clusters for tests.
 *
 * @author Janne Valkealahti
 *
 */
public class YarnClusterManager {

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

	/** Singleton instance */
	private static volatile YarnClusterManager instance = null;

	/** Synchronization monitor for the "refresh" and "destroy" */
	private final Object startupShutdownMonitor = new Object();

	/** Reference to the JVM shutdown hook, if registered */
	private Thread shutdownHook;

	/** Cached clusters*/
	private Hashtable<ClusterInfo, YarnCluster> clusters = new Hashtable<ClusterInfo, YarnCluster>();

	/**
	 * Private initializer for singleton access.
	 */
	private YarnClusterManager() {}

	/**
	 * Gets the singleton instance of {@link YarnClusterManager}.
	 *
	 * @return the singleton instance
	 */
	public static YarnClusterManager getInstance() {
		if (instance == null) {
			synchronized (YarnClusterManager.class) {
				if (instance == null) {
					instance = new YarnClusterManager();
				}
			}
		}
		return instance;
	}

	/**
	 * Gets the singleton instance of {@link YarnClusterManager}.
	 *
	 * @param registerShutdownHook if true register shutdown hook
	 * @return the singleton instance
	 */
	public static YarnClusterManager getInstance(boolean registerShutdownHook) {
		YarnClusterManager instance2 = getInstance();
		if (registerShutdownHook) {
			instance2.registerShutdownHook();
		}
		return instance2;
	}

	/**
	 * Gets and starts the mini cluster.
	 *
	 * @param clusterInfo the info about the cluster
	 * @return the running mini cluster
	 */
	public YarnCluster getCluster(ClusterInfo clusterInfo) {
		YarnCluster cluster = clusters.get(clusterInfo);
		if (cluster == null) {
			log.info("Building new cluster for ClusterInfo=" + clusterInfo);
			try {
				YarnClusterFactoryBean fb = new YarnClusterFactoryBean();
				fb.setClusterId("yarn-" + clusterInfo.hashCode());
				fb.setNodes(clusterInfo.getNumYarn());
				fb.setAutoStart(true);
				fb.afterPropertiesSet();
				cluster = fb.getObject();
				clusters.put(clusterInfo, cluster);
			} catch (Exception e) {
			}
		} else {
			log.info("Found cached cluster for ClusterInfo=" + clusterInfo);
		}
		return cluster;
	}

	/**
	 * Closes and remove {@link YarnCluster} from
	 * a manager cache.
	 *
	 * @param cluster the Yarn cluster
	 * @return true if cluster was closed, false otherwise
	 */
	public boolean closeCluster(YarnCluster cluster) {
		for (Entry<ClusterInfo, YarnCluster> entry : clusters.entrySet()) {
			if (entry.getValue().equals(cluster)) {
				entry.getValue().stop();
				clusters.remove(entry.getKey());
				return true;
			}
		}
		return false;
	}

	/**
	 * Close the manager, removes shutdown hook and
	 * closes all running clusters.
	 */
	public void close() {
		synchronized (this.startupShutdownMonitor) {
			doClose();
			// If we registered a JVM shutdown hook, we don't need it anymore now:
			// We've already explicitly closed the context.
			if (this.shutdownHook != null) {
				try {
					Runtime.getRuntime().removeShutdownHook(this.shutdownHook);
				}
				catch (IllegalStateException ex) {
					// ignore - VM is already shutting down
				}
			}
		}
	}

	/**
	 * Register a jvm shutdown hook allowing manager
	 * to gracefully shutdown clusters in case that
	 * hasn't already happened. This is usually the
	 * scenario in tests.
	 */
	public synchronized void registerShutdownHook() {
		if (this.shutdownHook == null) {
			// No shutdown hook registered yet.
			this.shutdownHook = new Thread() {
				@Override
				public void run() {
					doClose();
				}
			};
			Runtime.getRuntime().addShutdownHook(this.shutdownHook);
		}
	}

	/**
	 * Bring down and un-register all the running clusters.
	 */
	protected void doClose() {
		for (Entry<ClusterInfo, YarnCluster> entry : clusters.entrySet()) {
			entry.getValue().stop();
		}
		clusters.clear();
	}

}