/*
 * 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.item.xml;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.UnsupportedEncodingException;
import java.io.Writer;
import java.nio.channels.FileChannel;
import java.util.List;
import java.util.Map;

import javax.xml.stream.FactoryConfigurationError;
import javax.xml.stream.XMLEventFactory;
import javax.xml.stream.XMLEventWriter;
import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.transform.Result;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.batch.item.ItemStreamException;
import org.springframework.batch.item.ItemWriter;
import org.springframework.batch.item.WriteFailedException;
import org.springframework.batch.item.file.ResourceAwareItemWriterItemStream;
import org.springframework.batch.item.support.AbstractItemStreamItemWriter;
import org.springframework.batch.item.util.FileUtils;
import org.springframework.batch.item.xml.stax.NoStartEndDocumentStreamWriter;
import org.springframework.batch.support.transaction.TransactionAwareBufferedWriter;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.core.io.Resource;
import org.springframework.dao.DataAccessResourceFailureException;
import org.springframework.oxm.Marshaller;
import org.springframework.oxm.XmlMappingException;
import org.springframework.util.Assert;
import org.springframework.util.ClassUtils;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

/**
 * An implementation of {@link ItemWriter} which uses StAX and
 * {@link Marshaller} for serializing object to XML.
 * 
 * This item writer also provides restart, statistics and transaction features
 * by implementing corresponding interfaces.
 * 
 * The implementation is *not* thread-safe.
 * 
 * @author Peter Zozom
 * @author Robert Kasanicky
 * @author Michael Minella
 * 
 */
public class StaxEventItemWriter<T> extends AbstractItemStreamItemWriter<T> implements
ResourceAwareItemWriterItemStream<T>, InitializingBean {

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

	// default encoding
	private static final String DEFAULT_ENCODING = "UTF-8";

	// default encoding
	private static final String DEFAULT_XML_VERSION = "1.0";

	// default root tag name
	private static final String DEFAULT_ROOT_TAG_NAME = "root";

	// restart data property name
	private static final String RESTART_DATA_NAME = "position";

	// restart data property name
	private static final String WRITE_STATISTICS_NAME = "record.count";

	// file system resource
	private Resource resource;

	// xml marshaller
	private Marshaller marshaller;

	// encoding to be used while reading from the resource
	private String encoding = DEFAULT_ENCODING;

	// XML version
	private String version = DEFAULT_XML_VERSION;

	// name of the root tag
	private String rootTagName = DEFAULT_ROOT_TAG_NAME;

	// namespace prefix of the root tag
	private String rootTagNamespacePrefix = "";

	// namespace of the root tag
	private String rootTagNamespace = "";

	// root element attributes
	private Map<String, String> rootElementAttributes = null;

	// TRUE means, that output file will be overwritten if exists - default is
	// TRUE
	private boolean overwriteOutput = true;

	// file channel
	private FileChannel channel;

	// wrapper for XML event writer that swallows StartDocument and EndDocument
	// events
	private XMLEventWriter eventWriter;

	// XML event writer
	private XMLEventWriter delegateEventWriter;

	// current count of processed records
	private long currentRecordCount = 0;

	private boolean saveState = true;

	private StaxWriterCallback headerCallback;

	private StaxWriterCallback footerCallback;

	private Writer bufferedWriter;

	private boolean transactional = true;

	private boolean forceSync;

	private boolean shouldDeleteIfEmpty = false;

	public StaxEventItemWriter() {
		setExecutionContextName(ClassUtils.getShortName(StaxEventItemWriter.class));
	}

	/**
	 * Set output file.
	 * 
	 * @param resource the output file
	 */
	@Override
	public void setResource(Resource resource) {
		this.resource = resource;
	}

	/**
	 * Set Object to XML marshaller.
	 * 
	 * @param marshaller the Object to XML marshaller
	 */
	public void setMarshaller(Marshaller marshaller) {
		this.marshaller = marshaller;
	}

	/**
	 * headerCallback is called before writing any items.
	 */
	public void setHeaderCallback(StaxWriterCallback headerCallback) {
		this.headerCallback = headerCallback;
	}

	/**
	 * footerCallback is called after writing all items but before closing the
	 * file
	 */
	public void setFooterCallback(StaxWriterCallback footerCallback) {
		this.footerCallback = footerCallback;
	}

	/**
	 * Flag to indicate that writes should be deferred to the end of a
	 * transaction if present. Defaults to true.
	 * 
	 * @param transactional the flag to set
	 */
	public void setTransactional(boolean transactional) {
		this.transactional = transactional;
	}

	/**
	 * Flag to indicate that changes should be force-synced to disk on flush.
	 * Defaults to false, which means that even with a local disk changes could
	 * be lost if the OS crashes in between a write and a cache flush. Setting
	 * to true may result in slower performance for usage patterns involving
	 * many frequent writes.
	 * 
	 * @param forceSync the flag value to set
	 */
	public void setForceSync(boolean forceSync) {
		this.forceSync = forceSync;
	}

	/**
	 * Flag to indicate that the target file should be deleted if no items have
	 * been written (other than header and footer) on close. Defaults to false.
	 * 
	 * @param shouldDeleteIfEmpty the flag value to set
	 */
	public void setShouldDeleteIfEmpty(boolean shouldDeleteIfEmpty) {
		this.shouldDeleteIfEmpty = shouldDeleteIfEmpty;
	}

	/**
	 * Get used encoding.
	 * 
	 * @return the encoding used
	 */
	public String getEncoding() {
		return encoding;
	}

	/**
	 * Set encoding to be used for output file.
	 * 
	 * @param encoding the encoding to be used
	 */
	public void setEncoding(String encoding) {
		this.encoding = encoding;
	}

	/**
	 * Get XML version.
	 * 
	 * @return the XML version used
	 */
	public String getVersion() {
		return version;
	}

	/**
	 * Set XML version to be used for output XML.
	 * 
	 * @param version the XML version to be used
	 */
	public void setVersion(String version) {
		this.version = version;
	}

	/**
	 * Get the tag name of the root element.
	 * 
	 * @return the root element tag name
	 */
	public String getRootTagName() {
		return rootTagName;
	}

	/**
	 * Set the tag name of the root element. If not set, default name is used
	 * ("root"). Namespace URI and prefix can also be set optionally using the
	 * notation:
	 * 
	 * <pre>
	 * {uri}prefix:root
	 * </pre>
	 * 
	 * The prefix is optional (defaults to empty), but if it is specified then
	 * the uri must be provided. In addition you might want to declare other
	 * namespaces using the {@link #setRootElementAttributes(Map) root
	 * attributes}.
	 * 
	 * @param rootTagName the tag name to be used for the root element
	 */
	public void setRootTagName(String rootTagName) {
		this.rootTagName = rootTagName;
	}

	/**
	 * Get the namespace prefix of the root element. Empty by default.
	 * 
	 * @return the rootTagNamespacePrefix
	 */
	public String getRootTagNamespacePrefix() {
		return rootTagNamespacePrefix;
	}

	/**
	 * Get the namespace of the root element.
	 * 
	 * @return the rootTagNamespace
	 */
	public String getRootTagNamespace() {
		return rootTagNamespace;
	}

	/**
	 * Get attributes of the root element.
	 * 
	 * @return attributes of the root element
	 */
	public Map<String, String> getRootElementAttributes() {
		return rootElementAttributes;
	}

	/**
	 * Set the root element attributes to be written. If any of the key names
	 * begin with "xmlns:" then they are treated as namespace declarations.
	 * 
	 * @param rootElementAttributes attributes of the root element
	 */
	public void setRootElementAttributes(Map<String, String> rootElementAttributes) {
		this.rootElementAttributes = rootElementAttributes;
	}

	/**
	 * Set "overwrite" flag for the output file. Flag is ignored when output
	 * file processing is restarted.
	 * 
	 * @param overwriteOutput
	 */
	public void setOverwriteOutput(boolean overwriteOutput) {
		this.overwriteOutput = overwriteOutput;
	}

	public void setSaveState(boolean saveState) {
		this.saveState = saveState;
	}

	/**
	 * @throws Exception
	 * @see org.springframework.beans.factory.InitializingBean#afterPropertiesSet()
	 */
	@Override
	public void afterPropertiesSet() throws Exception {
		Assert.notNull(marshaller);
		if (rootTagName.contains("{")) {
			rootTagNamespace = rootTagName.replaceAll("\\{(.*)\\}.*", "$1");
			rootTagName = rootTagName.replaceAll("\\{.*\\}(.*)", "$1");
			if (rootTagName.contains(":")) {
				rootTagNamespacePrefix = rootTagName.replaceAll("(.*):.*", "$1");
				rootTagName = rootTagName.replaceAll(".*:(.*)", "$1");
			}
		}
	}

	/**
	 * Open the output source
	 * 
	 * @see org.springframework.batch.item.ItemStream#open(ExecutionContext)
	 */
	@Override
	public void open(ExecutionContext executionContext) {
		super.open(executionContext);

		Assert.notNull(resource, "The resource must be set");

		long startAtPosition = 0;
		boolean restarted = false;

		// if restart data is provided, restart from provided offset
		// otherwise start from beginning
		if (executionContext.containsKey(getExecutionContextKey(RESTART_DATA_NAME))) {
			startAtPosition = executionContext.getLong(getExecutionContextKey(RESTART_DATA_NAME));
			currentRecordCount = executionContext.getLong(getExecutionContextKey(WRITE_STATISTICS_NAME));
			restarted = true;
			if (shouldDeleteIfEmpty && currentRecordCount == 0) {
				// previous execution deleted the output file because no items were written
				restarted = false;
				startAtPosition = 0;
			} else {
				restarted = true;
			}
		}

		open(startAtPosition, restarted);

		if (startAtPosition == 0) {
			try {
				if (headerCallback != null) {
					headerCallback.write(delegateEventWriter);
				}
			}
			catch (IOException e) {
				throw new ItemStreamException("Failed to write headerItems", e);
			}
		}

	}

	/**
	 * Helper method for opening output source at given file position
	 */
	private void open(long position, boolean restarted) {

		File file;
		FileOutputStream os = null;
		FileChannel fileChannel = null;

		try {
			file = resource.getFile();
			FileUtils.setUpOutputFile(file, restarted, false, overwriteOutput);
			Assert.state(resource.exists(), "Output resource must exist");
			os = new FileOutputStream(file, true);
			fileChannel = os.getChannel();
			channel = os.getChannel();
			setPosition(position);
		}
		catch (IOException ioe) {
			throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", ioe);
		}

		XMLOutputFactory outputFactory = createXmlOutputFactory();

		if (outputFactory.isPropertySupported("com.ctc.wstx.automaticEndElements")) {
			// If the current XMLOutputFactory implementation is supplied by
			// Woodstox >= 3.2.9 we want to disable its
			// automatic end element feature (see:
			// http://jira.codehaus.org/browse/WSTX-165) per
			// http://jira.springframework.org/browse/BATCH-761).
			outputFactory.setProperty("com.ctc.wstx.automaticEndElements", Boolean.FALSE);
		}
		if (outputFactory.isPropertySupported("com.ctc.wstx.outputValidateStructure")) {
			// On restart we don't write the root element so we have to disable
			// structural validation (see:
			// http://jira.springframework.org/browse/BATCH-1681).
			outputFactory.setProperty("com.ctc.wstx.outputValidateStructure", Boolean.FALSE);
		}

		try {
			final FileChannel channel = fileChannel;
			if (transactional) {
				TransactionAwareBufferedWriter writer = new TransactionAwareBufferedWriter(channel, new Runnable() {
					@Override
					public void run() {
						closeStream();
					}
				});

				writer.setEncoding(encoding);
				writer.setForceSync(forceSync);
				bufferedWriter = writer;
			}
			else {
				bufferedWriter =  new BufferedWriter(new OutputStreamWriter(os, encoding));
			}
			delegateEventWriter = createXmlEventWriter(outputFactory, bufferedWriter);
			eventWriter = new NoStartEndDocumentStreamWriter(delegateEventWriter);
			initNamespaceContext(delegateEventWriter);
			if (!restarted) {
				startDocument(delegateEventWriter);
				if (forceSync) {
					channel.force(false);
				}
			}
		}
		catch (XMLStreamException xse) {
			throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", xse);
		}
		catch (UnsupportedEncodingException e) {
			throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource
					+ "] with encoding=[" + encoding + "]", e);
		} 
		catch (IOException e) {
			throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", e);
		}
	}

	/**
	 * Subclasses can override to customize the writer.
	 * @param outputFactory
	 * @param writer
	 * @return an xml writer
	 * @throws XMLStreamException
	 */
	protected XMLEventWriter createXmlEventWriter(XMLOutputFactory outputFactory, Writer writer)
			throws XMLStreamException {
		return outputFactory.createXMLEventWriter(writer);
	}

	/**
	 * Subclasses can override to customize the factory.
	 * @return a factory for the xml output
	 * @throws FactoryConfigurationError
	 */
	protected XMLOutputFactory createXmlOutputFactory() throws FactoryConfigurationError {
		return XMLOutputFactory.newInstance();
	}

	/**
	 * Subclasses can override to customize the event factory.
	 * @return a factory for the xml events
	 * @throws FactoryConfigurationError
	 */
	protected XMLEventFactory createXmlEventFactory() throws FactoryConfigurationError {
		XMLEventFactory factory = XMLEventFactory.newInstance();
		return factory;
	}

	/**
	 * Subclasses can override to customize the stax result.
	 * @return a result for writing to
	 * @throws Exception
	 */
	protected Result createStaxResult() throws Exception {
		return StaxUtils.getResult(eventWriter);
	}

	/**
	 * Inits the namespace context of the XMLEventWriter:
	 * <ul>
	 * <li>rootTagNamespacePrefix for rootTagName</li>
	 * <li>any other xmlns namespace prefix declarations in the root element attributes</li>
	 * </ul>
	 * 
	 * @param writer XML event writer
	 * @throws XMLStreamException
	 */
	protected void initNamespaceContext(XMLEventWriter writer) throws XMLStreamException {
		if (StringUtils.hasText(getRootTagNamespace())) {
			if(StringUtils.hasText(getRootTagNamespacePrefix())) {
				writer.setPrefix(getRootTagNamespacePrefix(), getRootTagNamespace());
			} else {
				writer.setDefaultNamespace(getRootTagNamespace());
			}
		}
		if (!CollectionUtils.isEmpty(getRootElementAttributes())) {
			for (Map.Entry<String, String> entry : getRootElementAttributes().entrySet()) {
				String key = entry.getKey();
				if (key.startsWith("xmlns")) {
					String prefix = "";
					if (key.contains(":")) {
						prefix = key.substring(key.indexOf(":") + 1);
					}
					log.debug("registering prefix: " +prefix + "=" + entry.getValue());
					writer.setPrefix(prefix, entry.getValue());
				}
			}
		}
	}

	/**
	 * Writes simple XML header containing:
	 * <ul>
	 * <li>xml declaration - defines encoding and XML version</li>
	 * <li>opening tag of the root element and its attributes</li>
	 * </ul>
	 * If this is not sufficient for you, simply override this method. Encoding,
	 * version and root tag name can be retrieved with corresponding getters.
	 * 
	 * @param writer XML event writer
	 * @throws XMLStreamException
	 */
	protected void startDocument(XMLEventWriter writer) throws XMLStreamException {

		XMLEventFactory factory = createXmlEventFactory();

		// write start document
		writer.add(factory.createStartDocument(getEncoding(), getVersion()));

		// write root tag
		writer.add(factory.createStartElement(getRootTagNamespacePrefix(), getRootTagNamespace(), getRootTagName()));
		if (StringUtils.hasText(getRootTagNamespace())) {
			if (StringUtils.hasText(getRootTagNamespacePrefix())) {
				writer.add(factory.createNamespace(getRootTagNamespacePrefix(), getRootTagNamespace()));
			}
			else {
				writer.add(factory.createNamespace(getRootTagNamespace()));
			}
		}

		// write root tag attributes
		if (!CollectionUtils.isEmpty(getRootElementAttributes())) {

			for (Map.Entry<String, String> entry : getRootElementAttributes().entrySet()) {
				String key = entry.getKey();
				if (key.startsWith("xmlns")) {
					String prefix = "";
					if (key.contains(":")) {
						prefix = key.substring(key.indexOf(":") + 1);
					}
					writer.add(factory.createNamespace(prefix, entry.getValue()));
				}
				else {
					writer.add(factory.createAttribute(key, entry.getValue()));
				}
			}

		}

		/*
		 * This forces the flush to write the end of the root element and avoids
		 * an off-by-one error on restart.
		 */
		writer.add(factory.createIgnorableSpace(""));
		writer.flush();

	}

	/**
	 * Writes the EndDocument tag manually.
	 * 
	 * @param writer XML event writer
	 * @throws XMLStreamException
	 */
	protected void endDocument(XMLEventWriter writer) throws XMLStreamException {

		// writer.writeEndDocument(); <- this doesn't work after restart
		// we need to write end tag of the root element manually

		String nsPrefix = !StringUtils.hasText(getRootTagNamespacePrefix()) ? "" : getRootTagNamespacePrefix() + ":";
		try {
			bufferedWriter.write("</" + nsPrefix + getRootTagName() + ">");
		}
		catch (IOException ioe) {
			throw new DataAccessResourceFailureException("Unable to close file resource: [" + resource + "]", ioe);
		}
	}

	/**
	 * Flush and close the output source.
	 * 
	 * @see org.springframework.batch.item.ItemStream#close()
	 */
	@Override
	public void close() {
		super.close();

		XMLEventFactory factory = createXmlEventFactory();
		try {
			delegateEventWriter.add(factory.createCharacters(""));
		}
		catch (XMLStreamException e) {
			log.error(e);
		}

		try {
			if (footerCallback != null) {
				footerCallback.write(delegateEventWriter);
			}
			delegateEventWriter.flush();
			endDocument(delegateEventWriter);
		}
		catch (IOException e) {
			throw new ItemStreamException("Failed to write footer items", e);
		}
		catch (XMLStreamException e) {
			throw new ItemStreamException("Failed to write end document tag", e);
		}
		finally {

			try {
				delegateEventWriter.close();
			}
			catch (XMLStreamException e) {
				log.error("Unable to close file resource: [" + resource + "] " + e);
			}
			finally {
				try {
					bufferedWriter.close();
				}
				catch (IOException e) {
					log.error("Unable to close file resource: [" + resource + "] " + e);
				}
				finally {
					if (!transactional) {
						closeStream();
					}
				}
			}
			if (currentRecordCount == 0 && shouldDeleteIfEmpty) {
				try {
					resource.getFile().delete();
				}
				catch (IOException e) {
					throw new ItemStreamException("Failed to delete empty file on close", e);
				}
			}
		}
	}

	private void closeStream() {
		try {
			channel.close();
		}
		catch (IOException ioe) {
			log.error("Unable to close file resource: [" + resource + "] " + ioe);
		}
	}

	/**
	 * Write the value objects and flush them to the file.
	 * 
	 * @param items the value object
	 * @throws IOException
	 * @throws XmlMappingException
	 */
	@Override
	public void write(List<? extends T> items) throws XmlMappingException, Exception {

		currentRecordCount += items.size();

		for (Object object : items) {
			Assert.state(marshaller.supports(object.getClass()),
					"Marshaller must support the class of the marshalled object");
			Result result = createStaxResult();
			marshaller.marshal(object, result);
		}
		try {
			eventWriter.flush();
			if (forceSync) {
				channel.force(false);
			}			
		}
		catch (XMLStreamException e) {
			throw new WriteFailedException("Failed to flush the events", e);
		} 
		catch (IOException e) {
			throw new WriteFailedException("Failed to flush the events", e);
		}

	}

	/**
	 * Get the restart data.
	 * 
	 * @see org.springframework.batch.item.ItemStream#update(ExecutionContext)
	 */
	@Override
	public void update(ExecutionContext executionContext) {
		super.update(executionContext);
		if (saveState) {
			Assert.notNull(executionContext, "ExecutionContext must not be null");
			executionContext.putLong(getExecutionContextKey(RESTART_DATA_NAME), getPosition());
			executionContext.putLong(getExecutionContextKey(WRITE_STATISTICS_NAME), currentRecordCount);
		}
	}

	/*
	 * Get the actual position in file channel. This method flushes any buffered
	 * data before position is read.
	 * 
	 * @return byte offset in file channel
	 */
	private long getPosition() {

		long position;

		try {
			eventWriter.flush();
			position = channel.position();
			if (bufferedWriter instanceof TransactionAwareBufferedWriter) {
				position += ((TransactionAwareBufferedWriter) bufferedWriter).getBufferSize();
			}
		}
		catch (Exception e) {
			throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", e);
		}

		return position;
	}

	/**
	 * Set the file channel position.
	 * 
	 * @param newPosition new file channel position
	 */
	private void setPosition(long newPosition) {

		try {
			channel.truncate(newPosition);
			channel.position(newPosition);
		}
		catch (IOException e) {
			throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", e);
		}

	}

}
