/*
    * Copyright 2006-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.batch.core.configuration.support;
  
  import java.util.ArrayList;
  import java.util.Collection;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.concurrent.ConcurrentHashMap;
  
  import org.apache.commons.logging.Log;
  import org.apache.commons.logging.LogFactory;
  import org.springframework.batch.core.Job;
  import org.springframework.batch.core.Step;
  import org.springframework.batch.core.configuration.DuplicateJobException;
  import org.springframework.batch.core.configuration.JobFactory;
  import org.springframework.batch.core.configuration.JobRegistry;
  import org.springframework.batch.core.configuration.StepRegistry;
  import org.springframework.batch.core.launch.NoSuchJobException;
  import org.springframework.batch.core.step.StepLocator;
  import org.springframework.beans.factory.InitializingBean;
  import org.springframework.context.ApplicationContext;
  import org.springframework.context.ConfigurableApplicationContext;
  import org.springframework.util.Assert;
  
  /**
   * Default implementation of {@link JobLoader}. Uses a {@link JobRegistry} to
   * manage a population of loaded jobs and clears them up when asked. An optional
   * {@link StepRegistry} might also be set to register the step(s) available for
   * each registered job.
   *
   * @author Dave Syer
   * @author Stephane Nicoll
   */
  public class DefaultJobLoader implements JobLoader, InitializingBean {
  
  	private static Log logger = LogFactory.getLog(DefaultJobLoader.class);
  
  	private JobRegistry jobRegistry;
  	private StepRegistry stepRegistry;
  
  	private Map<ApplicationContextFactory, ConfigurableApplicationContext> contexts = new ConcurrentHashMap<ApplicationContextFactory, ConfigurableApplicationContext>();
  
  	private Map<ConfigurableApplicationContext, Collection<String>> contextToJobNames = new ConcurrentHashMap<ConfigurableApplicationContext, Collection<String>>();
  
  	/**
  	 * Default constructor useful for declarative configuration.
  	 */
  	public DefaultJobLoader() {
  		this(null, null);
  	}
  
  	/**
  	 * Creates a job loader with the job registry provided.
  	 *
  	 * @param jobRegistry a {@link JobRegistry}
  	 */
  	public DefaultJobLoader(JobRegistry jobRegistry) {
  		this(jobRegistry, null);
  	}
  
  	/**
  	 * Creates a job loader with the job and step registries provided.
  	 *
  	 * @param jobRegistry a {@link JobRegistry}
  	 * @param stepRegistry a {@link StepRegistry}
  	 */
  	public DefaultJobLoader(JobRegistry jobRegistry, StepRegistry stepRegistry) {
  		this.jobRegistry = jobRegistry;
  		this.stepRegistry = stepRegistry;
  	}
  
  	/**
  	 * The {@link JobRegistry} to use for jobs created.
  	 *
  	 * @param jobRegistry the job registry
  	 */
  	public void setJobRegistry(JobRegistry jobRegistry) {
  		this.jobRegistry = jobRegistry;
  	}
  
  	/**
  	 * The {@link StepRegistry} to use for the steps of created jobs.
  	 *
  	 * @param stepRegistry the step registry
  	 */
 	public void setStepRegistry(StepRegistry stepRegistry) {
 		this.stepRegistry = stepRegistry;
 	}
 
 	/**
 	 * Unregister all the jobs and close all the contexts created by this
 	 * loader.
 	 *
 	 * @see JobLoader#clear()
 	 */
 	@Override
 	public void clear() {
 		for (ConfigurableApplicationContext context : contexts.values()) {
 			if (context.isActive()) {
 				context.close();
 			}
 		}
 		for (String jobName : jobRegistry.getJobNames()) {
 			doUnregister(jobName);
 		}
 		contexts.clear();
 	}
 
 	@Override
 	public Collection<Job> reload(ApplicationContextFactory factory) {
 
 		// If the same factory is loaded twice the context can be closed
 		if (contexts.containsKey(factory)) {
 			ConfigurableApplicationContext context = contexts.get(factory);
 			for (String name : contextToJobNames.get(context)) {
 				logger.debug("Unregistering job: " + name + " from context: " + context.getDisplayName());
 				doUnregister(name);
 			}
 			context.close();
 		}
 
 		try {
 			return doLoad(factory, true);
 		}
 		catch (DuplicateJobException e) {
 			throw new IllegalStateException("Found duplicate job in reload (it should have been unregistered "
 					+ "if it was previously registered in this loader)", e);
 		}
 	}
 
 	@Override
 	public Collection<Job> load(ApplicationContextFactory factory) throws DuplicateJobException {
 		return doLoad(factory, false);
 	}
 
 	private Collection<Job> doLoad(ApplicationContextFactory factory, boolean unregister) throws DuplicateJobException {
 
 		Collection<String> jobNamesBefore = jobRegistry.getJobNames();
 		ConfigurableApplicationContext context = factory.createApplicationContext();
 		Collection<String> jobNamesAfter = jobRegistry.getJobNames();
 		// Try to detect auto-registration (e.g. through a bean post processor)
 		boolean autoRegistrationDetected = jobNamesAfter.size() > jobNamesBefore.size();
 
 		Collection<String> jobsRegistered = new HashSet<String>();
 		if (autoRegistrationDetected) {
 			for (String name : jobNamesAfter) {
 				if (!jobNamesBefore.contains(name)) {
 					jobsRegistered.add(name);
 				}
 			}
 		}
 
 		contexts.put(factory, context);
 		String[] names = context.getBeanNamesForType(Job.class);
 
 		for (String name : names) {
 
 			if (!autoRegistrationDetected) {
 
 				Job job = (Job) context.getBean(name);
 				String jobName = job.getName();
 
 				// On reload try to unregister first
 				if (unregister) {
 					logger.debug("Unregistering job: " + jobName + " from context: " + context.getDisplayName());
 					doUnregister(jobName);
 				}
 
 				logger.debug("Registering job: " + jobName + " from context: " + context.getDisplayName());
 				doRegister(context, job);
 				jobsRegistered.add(jobName);
 			}
 
 		}
 
 		Collection<Job> result = new ArrayList<Job>();
 		for (String name : jobsRegistered) {
 			try {
 				result.add(jobRegistry.getJob(name));
 			}
 			catch (NoSuchJobException e) {
 				// should not happen;
 				throw new IllegalStateException("Could not retrieve job that was should have been registered", e);
 			}
 
 		}
 
 		contextToJobNames.put(context, jobsRegistered);
 
 		return result;
 
 	}
 
 	/**
 	 * Returns all the {@link Step} instances defined by the specified {@link StepLocator}.
 	 * <p/>
 	 * The specified <tt>jobApplicationContext</tt> is used to collect additional steps that
 	 * are not exposed by the step locator
 	 *
 	 * @param stepLocator the given step locator
 	 * @param jobApplicationContext the application context of the job
 	 * @return all the {@link Step} defined by the given step locator and context
 	 * @see StepLocator
 	 */
 	private Collection<Step> getSteps(final StepLocator stepLocator, final ApplicationContext jobApplicationContext) {
 		final Collection<String> stepNames = stepLocator.getStepNames();
 		final Collection<Step> result = new ArrayList<Step>();
 		for (String stepName : stepNames) {
 			result.add(stepLocator.getStep(stepName));
 		}
 
 		// Because some steps are referenced by name, we need to look in the context to see if there
 		// are more Step instances defined. Right now they are registered as being available in the
 		// context of the job but we have no idea if they are linked to that Job or not.
 		@SuppressWarnings("unchecked")
 		final Map<String, Step> allSteps = jobApplicationContext.getBeansOfType(Step.class);
 		for (Map.Entry<String, Step> entry : allSteps.entrySet()) {
 			if (!stepNames.contains(entry.getKey())) {
 				result.add(entry.getValue());
 			}
 		}
 		return result;
 	}
 
 	/**
 	 * Registers the specified {@link Job} defined in the specified {@link ConfigurableApplicationContext}.
 	 * <p/>
 	 * Makes sure to update the {@link StepRegistry} if it is available.
 	 *
 	 * @param context the context in which the job is defined
 	 * @param job the job to register
 	 * @throws DuplicateJobException if that job is already registered
 	 */
 	private void doRegister(ConfigurableApplicationContext context, Job job) throws DuplicateJobException {
 		final JobFactory jobFactory = new ReferenceJobFactory(job);
 		jobRegistry.register(jobFactory);
 
 		if (stepRegistry != null) {
 			if (!(job instanceof StepLocator)) {
 				throw new UnsupportedOperationException("Cannot locate steps from a Job that is not a StepLocator: job="
 						+ job.getName() + " does not implement StepLocator");
 			}
 			stepRegistry.register(job.getName(), getSteps((StepLocator) job, context));
 		}
 	}
 
 	/**
 	 * Unregisters the job identified by the specified <tt>jobName</tt>.
 	 *
 	 * @param jobName the name of the job to unregister
 	 */
 	private void doUnregister(String jobName)  {
 		jobRegistry.unregister(jobName);
 		if (stepRegistry != null) {
 			stepRegistry.unregisterStepsFromJob(jobName);
 		}
 
 	}
 
 	@Override
 	public void afterPropertiesSet() {
 		Assert.notNull(jobRegistry, "Job registry could not be null.");
 	}
 }