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.SOAPBody;
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.*;

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");
        }

        System.out.println();
        System.out.println(".....................Original SOAP Envelop.......................... ");
        System.out.println(context.getEnvelope());
        System.out.println(".................................................................... ");
        System.out.println();


        try {
            inMessage = getMessageInputStream(context);
            outMessage=getMessageOutPutStream(context);

            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            byte[] buffer = new byte[1024];
            int len;
            while ((len = inMessage.read(buffer)) > -1 ) {
                baos.write(buffer, 0, len);
            }
            baos.flush();

            inMessage = new ByteArrayInputStream(baos.toByteArray());
            InputStream temp = new ByteArrayInputStream(baos.toByteArray());

            System.out.println();
            System.out.println(".....................Original SOAP Message.......................... ");
            XMLStreamReader reader = StAXUtils.createXMLStreamReader(temp);
            System.out.println(new StAXOMBuilder(reader).getDocumentElement());
            System.out.println(".................................................................... ");
            System.out.println();

        } catch (IOException e) {
            handleException("Error while reading the input stream ", e, context);
        } catch (XMLStreamException 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 {



        if (resulMessage != null) {

            System.out.println();
            System.out.println("...................Transformed SOAP Message......................... ");
            InputStream temp=resulMessage.getInputStream();
            XMLStreamReader reader = StAXUtils.createXMLStreamReader(temp);
            System.out.println(new StAXOMBuilder(reader).getDocumentElement());
            System.out.println(".................................................................... ");
            System.out.println();

            SOAPFactory factory = OMAbstractFactory.getSOAP12Factory();

            QName BINARY_CONTENT_QNAME = new QName("http://ws.apache.org/commons/ns/payload", "binary");
            OMNamespace ns = factory.createOMNamespace(
                    BINARY_CONTENT_QNAME.getNamespaceURI(), "ns");
            OMElement omEle = factory.createOMElement(
                    BINARY_CONTENT_QNAME.getLocalPart(), ns);

            DataHandler dataHandler = new DataHandler(resulMessage);
            OMText textData = factory.createOMText(dataHandler, true);

            omEle.addChild(textData);

            SOAPBody body=synCtx.getEnvelope().getBody();
            body.getFirstElement().detach();
            body.addChild(omEle);

            System.out.println();
            System.out.println("...................Transformed SOAP Envelop......................... ");
            System.out.println(synCtx.getEnvelope());
            System.out.println(".................................................................... ");
            System.out.println();

            return true;
        }
        return false;
    }

    public String getKey() {
        return key;
    }

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

}