package com.mulesoft.support;


import java.io.ByteArrayOutputStream;
import java.util.ArrayList;
import java.util.List;
import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;
import javax.xml.stream.XMLStreamWriter;
import javax.xml.transform.Source;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.stax.StAXResult;
import javax.xml.transform.stax.StAXSource;


import org.mule.api.MuleEvent;
import org.mule.api.MuleMessage;
import org.mule.routing.AbstractSplitter;

public class CustomXmlMessageSplitter extends AbstractSplitter {

	private String targetElement = null;
	
	@SuppressWarnings({ "unchecked", "rawtypes" })
	@Override
	protected List<MuleMessage> splitMessage(MuleEvent event){
			
		MuleMessage message = event.getMessage();
	
		if (!(message .getPayload() instanceof XMLStreamReader)){
			throw new IllegalArgumentException("Payload must be an instance of a class implementing javax.xml.stream.XMLStreamReader. Please use org.mule.module.xml.transformer.XmlToXMLStreamReader to transform the payload to an instance of org.mule.module.xml.stax.DelegateXMLStreamReader");
		}
		
		List splittedMessages = new ArrayList();

		System.setProperty("javax.xml.transform.TransformerFactory", "com.sun.org.apache.xalan.internal.xsltc.trax.TransformerFactoryImpl");
		XMLStreamReader xsr = (XMLStreamReader) message.getPayload();
    	XMLOutputFactory xmlOutputFactory = XMLOutputFactory.newInstance();
    	XMLStreamWriter xsw = null;
    	ByteArrayOutputStream baos = new ByteArrayOutputStream();

    	try {
        	
	        TransformerFactory transformerFactory = TransformerFactory.newInstance();
	        Transformer transformer = transformerFactory.newTransformer();
	        
	        while(xsr.hasNext()) {
	        	if(xsr.isStartElement() && xsr.getLocalName().equals(this.targetElement)){
	        		xsw = xmlOutputFactory.createXMLStreamWriter(baos);
		        	Source source = new StAXSource(xsr);
		        	StAXResult result = new StAXResult(xsw);
		        	
		            transformer.transform(source, result);
		            
		            splittedMessages.add(baos.toByteArray());
		            
		            baos.flush();
		            baos.reset();
		            
	        	}
	        	xsr.next();
	        }

		} catch (Exception e) {
			logger.error(e);
		} finally{
			try {
				xsw.close();
				xsr.close();
			} catch (XMLStreamException e) {
				logger.error("Error disposing XMLStreamReader and XMLStreamWriter: " + e.getMessage());
			}
		}
		
		return splittedMessages;
		
	}

	public String getTargetElement() {
		return targetElement;
	}

	public void setTargetElement(String targetElement) {
		this.targetElement = targetElement;
	}
	
}