package org.wso2.carbon.mediator;

import org.apache.axiom.om.*;
import org.apache.axiom.om.impl.builder.StAXOMBuilder;
import org.apache.axiom.om.util.StAXUtils;
import org.apache.axiom.soap.SOAPEnvelope;
import org.apache.axiom.soap.SOAPFactory;
import org.apache.synapse.MessageContext;
import org.apache.synapse.SynapseLog;
import org.apache.synapse.config.SynapseConfigUtils;
import org.apache.synapse.core.SynapseEnvironment;
import org.apache.synapse.mediators.AbstractMediator;
import org.wso2.carbon.relay.StreamingOnRequestDataSource;

import javax.activation.DataHandler;
import javax.activation.DataSource;
import javax.xml.namespace.QName;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamReader;
import javax.xml.transform.*;
import javax.xml.transform.stream.StreamResult;
import javax.xml.transform.stream.StreamSource;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

public class XMLConvertor extends AbstractMediator {

    private String key=null;
    private TransformerFactory transFact=TransformerFactory.newInstance();
    private MessageDataSource resulMessage=null;
    private InputStream inMessage=null;
    private OutputStream outMessage=null;

    public boolean mediate(MessageContext context) {

        Templates cTemplate=null;

        SynapseLog synLog = getLog(context);

        if (synLog.isTraceOrDebugEnabled()) {
            synLog.traceOrDebug("XMLConverter mediator : start");
        }

        try {
            inMessage = getMessageInputStream(context);
            outMessage=getMessageOutPutStream(context);
        } catch (IOException e) {
            handleException("Error while reading the input stream ", e, context);
        }

        String generatedXsltKey = key;
        cTemplate=createTemplate(context,generatedXsltKey);

        System.setProperty("javax.xml.transform.TransformerFactory","net.sf.saxon.TransformerFactoryImpl");
        System.out.println("Transforming On Progress.....");

        try {
            transform(inMessage,outMessage,cTemplate);
            outMessage.close();
            writeResult(context);

            if (synLog.isTraceOrDebugEnabled()) {
                synLog.traceOrDebug("XMLConverter mediator : Done");
            }

            return  true;
        } catch (Exception e) {
            handleException("Error while transforming the Stream ", e, context);
        }

        handleException("Unexpected SOAP message content found. " + this.getClass().getName() +
                " mediator can only be used with messages built with BinaryRelayBuilder", context);
        return false;
    }

    private void transform(InputStream xmlIn,OutputStream out, Templates templates)
            throws Exception {

        //Source xslt = new StreamSource(xsltIn);
        Transformer trans = templates.newTransformer();
        Source source = new StreamSource(xmlIn);
        Result resultXML = new StreamResult(out);
        trans.transform(source,resultXML);

    }

    /**
     * Create a XSLT template object and assign it to the cachedTemplates variable
     * @param synCtx current message
     * @param xsltKey evaluated xslt key(real key value) for dynamic or static key
     * @return cached template
     */
    private Templates createTemplate(MessageContext synCtx, String xsltKey) {
        Templates cachedTemplates = null;

        try {
            cachedTemplates = transFact.newTemplates(
                    SynapseConfigUtils.getStreamSource(synCtx.getEntry(xsltKey)));
        } catch (Exception e) {
            handleException("Error creating XSLT transformer using : " + xsltKey, e, synCtx);
        }
        return cachedTemplates;
    }

    private InputStream getMessageInputStream(MessageContext context) throws IOException {
        InputStream temp;
        SOAPEnvelope envelope = context.getEnvelope();
        OMElement contentEle = envelope.getBody().getFirstElement();

        if (contentEle != null) {

            OMNode node = contentEle.getFirstOMChild();

            if (node != null && (node instanceof OMText)) {

                OMText binaryDataNode = (OMText) node;
                DataHandler dh = (DataHandler) binaryDataNode.getDataHandler();
                DataSource dataSource = dh.getDataSource();

                if (dataSource instanceof StreamingOnRequestDataSource) {
                    // preserve the content while reading the incoming data stream
                    ((StreamingOnRequestDataSource) dataSource).setLastUse(false);
                    // forcing to consume the incoming data stream
                    temp=dataSource.getInputStream();
                    return temp;
                }
            }
        }
        return null;
    }

    private OutputStream getMessageOutPutStream(MessageContext synCtx) throws IOException {
        SynapseEnvironment synEnv = synCtx.getEnvironment();
        resulMessage=new MessageDataSource(synEnv.createOverflowBlob());
        outMessage = resulMessage.getOutputStream();
        return outMessage;
    }



    private boolean writeResult(MessageContext synCtx) throws IOException, XMLStreamException {

        SOAPFactory factory = OMAbstractFactory.getSOAP12Factory();
        SOAPEnvelope env = synCtx.getEnvelope();

        if (resulMessage != null) {
            XMLStreamReader reader = StAXUtils.createXMLStreamReader(resulMessage.getInputStream());
            StAXOMBuilder sb=new StAXOMBuilder(reader);
            System.out.println("Transformed Stream....................");
            System.out.println(sb.getDocumentElement());
            System.out.println("......................................");

            OMDataSource omDataSource=new MessageOMDataSource(resulMessage.getOverflowBlob());
            OMElement omEle = factory.createOMElement(omDataSource,sb.getDocumentElement().getQName());
            synCtx.setDoingMTOM(false);
            env.getBody().addChild(omEle);

            System.out.println("Child Added Body......................");
            System.out.println(env);
            System.out.println("......................................");
            return true;
        }
        return false;
    }

    public String getKey() {
        return key;
    }

    public void setKey(String key) {
        this.key = key;
    }

}