/*
    * Copyright 2006-2011 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.step.builder;
  
  import org.springframework.batch.core.Step;
  import org.springframework.batch.core.partition.PartitionHandler;
  import org.springframework.batch.core.partition.StepExecutionSplitter;
  import org.springframework.batch.core.partition.support.PartitionStep;
  import org.springframework.batch.core.partition.support.Partitioner;
  import org.springframework.batch.core.partition.support.SimpleStepExecutionSplitter;
  import org.springframework.batch.core.partition.support.StepExecutionAggregator;
  import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler;
  import org.springframework.core.task.SyncTaskExecutor;
  import org.springframework.core.task.TaskExecutor;
  
  /**
   * Step builder for {@link PartitionStep} instances. A partition step executes the same step (possibly remotely)
   * multiple times with different input parameters (in the form of execution context). Useful for parallelization.
   * 
   * @author Dave Syer
   * 
   * @since 2.2
   */
  public class PartitionStepBuilder extends StepBuilderHelper<PartitionStepBuilder> {
  
  	private TaskExecutor taskExecutor;
  
  	private Partitioner partitioner;
  
  	private static final int DEFAULT_GRID_SIZE = 6;
  
  	private Step step;
  
  	private PartitionHandler partitionHandler;
  
  	private int gridSize = DEFAULT_GRID_SIZE;
  
  	private StepExecutionSplitter splitter;
  
  	private StepExecutionAggregator aggregator;
  
  	private String stepName;
  
  	/**
  	 * Create a new builder initialized with any properties in the parent. The parent is copied, so it can be re-used.
  	 * 
  	 * @param parent a parent helper containing common step properties
  	 */
  	public PartitionStepBuilder(StepBuilderHelper<?> parent) {
  		super(parent);
  	}
  
  	/**
  	 * Add a partitioner which can be used to create a {@link StepExecutionSplitter}. Use either this or an explicit
  	 * {@link #splitter(StepExecutionSplitter)} but not both.
  	 * 
  	 * @param slaveStepName the name of the slave step (used to construct step execution names)
  	 * @param partitioner a partitioner to use
  	 * @return this for fluent chaining
  	 */
  	public PartitionStepBuilder partitioner(String slaveStepName, Partitioner partitioner) {
  		this.stepName = slaveStepName;
  		this.partitioner = partitioner;
  		return this;
  	}
  
  	/**
  	 * Provide an actual step instance to execute in parallel. If an explicit
  	 * {@link #partitionHandler(PartitionHandler)} is provided, the step is optional and is only used to extract
  	 * configuration data (name and other basic properties of a step).
  	 * 
  	 * @param step a step to execute in parallel
  	 * @return this for fluent chaining
  	 */
  	public PartitionStepBuilder step(Step step) {
  		this.step = step;
  		return this;
  	}
  
  	/**
  	 * Provide a task executor to use when constructing a {@link PartitionHandler} from the {@link #step(Step)}. Mainly
  	 * used for running a step locally in parallel, but can be used to execute remotely if the step is remote. Not used
  	 * if an explicit {@link #partitionHandler(PartitionHandler)} is provided.
  	 * 
  	 * @param taskExecutor a task executor to use when executing steps in parallel
  	 * @return this for fluent chaining
 	 */
 	public PartitionStepBuilder taskExecutor(TaskExecutor taskExecutor) {
 		this.taskExecutor = taskExecutor;
 		return this;
 	}
 
 	/**
 	 * Provide an explicit partition handler that will carry out the work of the partition step. The partition handler
 	 * is the main SPI for adapting a partition step to a specific distributed computation environment. Optional if you
 	 * only need local or remote processing through the Step interface.
 	 * 
 	 * @see #step(Step) for setting up a default handler that works with a local or remote Step
 	 * 
 	 * @param partitionHandler a partition handler
 	 * @return this for fluent chaining
 	 */
 	public PartitionStepBuilder partitionHandler(PartitionHandler partitionHandler) {
 		this.partitionHandler = partitionHandler;
 		return this;
 	}
 
 	/**
 	 * A hint to the {@link #splitter(StepExecutionSplitter)} about how many step executions are required. If running
 	 * locally or remotely through a {@link #taskExecutor(TaskExecutor)} determines precisely the number of step
 	 * execution sin the first attempt at a partition step execution.
 	 * 
 	 * @param gridSize the grid size
 	 * @return this for fluent chaining
 	 */
 	public PartitionStepBuilder gridSize(int gridSize) {
 		this.gridSize = gridSize;
 		return this;
 	}
 
 	/**
 	 * Provide an explicit {@link StepExecutionSplitter} instead of having one build from the
 	 * {@link #partitioner(String, Partitioner)}. USeful if you need more control over the splitting.
 	 * 
 	 * @param splitter a step execution splitter
 	 * @return this for fluent chaining
 	 */
 	public PartitionStepBuilder splitter(StepExecutionSplitter splitter) {
 		this.splitter = splitter;
 		return this;
 	}
 
 	/**
 	 * Provide a step execution aggregator for aggregating partitioned step executions into a single result for the
 	 * {@link PartitionStep} itself.  Default is a simple implementation that works in most cases.
 	 * 
 	 * @param aggregator a step execution aggregator
 	 * @return this for fluent chaining
 	 */
 	public PartitionStepBuilder aggregator(StepExecutionAggregator aggregator) {
 		this.aggregator = aggregator;
 		return this;
 	}
 
 	public Step build() {
 
 		PartitionStep step = new PartitionStep();
 		step.setName(getName());
 		super.enhance(step);
 
 		if (partitionHandler != null) {
 			step.setPartitionHandler(partitionHandler);
 		}
 		else {
 			TaskExecutorPartitionHandler partitionHandler = new TaskExecutorPartitionHandler();
 			partitionHandler.setStep(this.step);
 			if (taskExecutor == null) {
 				taskExecutor = new SyncTaskExecutor();
 			}
 			partitionHandler.setGridSize(gridSize);
 			partitionHandler.setTaskExecutor(taskExecutor);
 			step.setPartitionHandler(partitionHandler);
 		}
 
 		if (splitter != null) {
 			step.setStepExecutionSplitter(splitter);
 		}
 		else {
 
 			boolean allowStartIfComplete = isAllowStartIfComplete();
 			String name = stepName;
 			if (this.step != null) {
 				try {
 					allowStartIfComplete = this.step.isAllowStartIfComplete();
 					name = this.step.getName();
 				}
 				catch (Exception e) {
 					logger.info("Ignored exception from step asking for name and allowStartIfComplete flag. "
 							+ "Using default from enclosing PartitionStep (" + name + "," + allowStartIfComplete + ").");
 				}
 			}
 			SimpleStepExecutionSplitter splitter = new SimpleStepExecutionSplitter();
 			splitter.setPartitioner(partitioner);
 			splitter.setJobRepository(getJobRepository());
 			splitter.setAllowStartIfComplete(allowStartIfComplete);
 			splitter.setStepName(name);
 			this.splitter = splitter;
 			step.setStepExecutionSplitter(splitter);
 
 		}
 
 		if (aggregator != null) {
 			step.setStepExecutionAggregator(aggregator);
 		}
 
 		try {
 			step.afterPropertiesSet();
 		}
 		catch (Exception e) {
 			throw new StepBuilderException(e);
 		}
 
 		return step;
 
 	}
 
 }