/*
 *  Licensed to the Apache Software Foundation (ASF) under one
 *  or more contributor license agreements.  See the NOTICE file
 *  distributed with this work for additional information
 *  regarding copyright ownership.  The ASF licenses this file
 *  to you 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.apache.synapse.mediators.xquery;

import net.sf.saxon.s9api.*;
import org.apache.axiom.om.OMElement;
import org.apache.axiom.om.OMNode;
import org.apache.axiom.om.OMText;
import org.apache.axiom.om.impl.builder.StAXOMBuilder;
import org.apache.axiom.om.impl.dom.DOOMAbstractFactory;
import org.apache.axiom.om.util.ElementHelper;
import org.apache.synapse.MessageContext;
import org.apache.synapse.SynapseException;
import org.apache.synapse.SynapseLog;
import org.apache.synapse.config.Entry;
import org.apache.synapse.config.SynapseConfigUtils;
import org.apache.synapse.mediators.AbstractMediator;
import org.apache.synapse.mediators.MediatorProperty;
import org.apache.synapse.mediators.Value;
import org.apache.synapse.util.xpath.SourceXPathSupport;
import org.apache.synapse.util.xpath.SynapseXPath;
import org.w3c.dom.Element;

import javax.activation.DataHandler;
import javax.xml.stream.XMLInputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamSource;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringReader;
import java.util.*;


/**
 * The XQueryMediator  provides the means to extract and manipulate data from XML documents  using
 * XQuery . It is possible to query against the current  SOAP Message or external XML. To query
 * against the current  SOAP Message ,it is need to define custom variable with any name and type as
 * element,document,document_element By providing a expression ,It is possible to select a custom
 * node for querying.The all the variable  that  have defined  in the mediator will be available
 * during the query process .Basic variable can use bind basic type.
 * currently only support * string,int,byte,short,double,long,float and boolean * types.
 * Custom Variable can use to bind XML documents ,SOAP payload and any basic type which create
 * through the XPath expression .
 */

public class XQueryMediator extends AbstractMediator {

    /* Properties that must set to the Processor  */
    private final List<MediatorProperty> processorProperties = new ArrayList<MediatorProperty>();

    /* The key for lookup the xquery (Supports both static and dynamic keys)*/
    private Value queryKey;

    /* The source of the xquery */
    private String querySource;

    /*The target node*/
    private final SourceXPathSupport target = new SourceXPathSupport();

    /* The list of variables for binding to the DyanamicContext in order to available for querying */
    private final List<MediatorVariable> variables = new ArrayList<MediatorVariable>();

    /*Lock used to ensure thread-safe lookup of the object from the registry */
    private final Object resourceLock = new Object();

    /* Is it need to use DOMSource and DOMResult? */
    private boolean useDOMSource = false;

    /*The Processor allows global Saxon configuration options to be set &  it acts as a factory for generating XQuery compiler */
    private Processor cachedProcessor = null;

    /* XQueryCompiler allows to compile XQuery 1.0 queries */
    private XQueryCompiler cachedQueryCompiler = null;

    /* An XQueryEvaluator that use for represents a compiled and loaded query ready for execution. XQueryEvaluator will recreate if query has changed */
    private Map<String, XQueryEvaluator> cachedXQueryEvaluatorMap = new Hashtable<String, XQueryEvaluator>();

    public XQueryMediator() {
    }

    /**
     * Performs the query and attached the result to the target Node
     *
     * @param synCtx The current message
     * @return true always
     */
    public boolean mediate(MessageContext synCtx) {

        try {

            if (synCtx.getEnvironment().isDebuggerEnabled()) {
                if (super.divertMediationRoute(synCtx)) {
                    return true;
                }
            }

            SynapseLog synLog = getLog(synCtx);

            if (synLog.isTraceOrDebugEnabled()) {
                synLog.traceOrDebug("Start : XQuery mediator");

                if (synLog.isTraceTraceEnabled()) {
                    synLog.traceTrace("Message : " + synCtx.getEnvelope());
                }
                synLog.traceOrDebug("Performing XQuery using query resource with key : " +
                        queryKey);
            }

            // perform the xquery
            performQuery(synCtx, synLog);

            synLog.traceOrDebug("End : XQuery mediator");

            return true;

        } catch (Exception e) {
            handleException("Unable to execute the query ", e);
        }
        return false;
    }

    /**
     * Perform the quering and get the result and attached to the target node
     *
     * @param synCtx The current MessageContext
     * @param synLog the Synapse log to use
     */
    private void performQuery(MessageContext synCtx, SynapseLog synLog) {

        boolean reLoad = false;
        boolean needSet = false;
        XQueryEvaluator queryEvaluator = null;
        String generatedQueryKey = null;
        XQueryExecutable xQueryExecutable = null;
        XdmValue xdmValue;
        boolean isQueryKeyGenerated = false;

        if (queryKey != null) {
            // Derive actual key from xpath or get static key
            generatedQueryKey = queryKey.evaluateValue(synCtx);
        }

        if (generatedQueryKey != null) {
            isQueryKeyGenerated = true;
        }

        if (generatedQueryKey != null && !"".equals(generatedQueryKey)) {

            Entry dp = synCtx.getConfiguration().getEntryDefinition(generatedQueryKey);
            // if the queryKey refers to a dynamic resource
            if (dp != null && dp.isDynamic()) {
                if (!dp.isCached() || dp.isExpired()) {
                    reLoad = true;
                }
            }
        }

        try {
            synchronized (resourceLock) {
                //creating processor
                if (cachedProcessor == null) {
                    cachedProcessor = new Processor(false);
                    //setting up the properties to the Processor
                    if (processorProperties != null && !processorProperties.isEmpty()) {
                        synLog.traceOrDebug("Setting up properties to the XQDataSource");
                        for (MediatorProperty processorProperty : processorProperties) {
                            if (processorProperty != null) {
                                cachedProcessor.setConfigurationProperty(processorProperty.getName(),
                                        processorProperty.getValue());
                            }
                        }
                    }
                }

                //creating XQueryCompiler
                if (cachedQueryCompiler == null) {
                    synLog.traceOrDebug("Creating a compiler from the Processor ");
                    cachedQueryCompiler = cachedProcessor.newXQueryCompiler();
                }

                //If already cached evaluator then load it from cachedXQueryEvaluatorMap
                if (isQueryKeyGenerated) {
                    queryEvaluator = cachedXQueryEvaluatorMap.get(generatedQueryKey);
                }

                if (reLoad || queryEvaluator == null) {

                    if (querySource != null && !"".equals(querySource)) {

                        if (synLog.isTraceOrDebugEnabled()) {
                            synLog.traceOrDebug("Using in-lined query source - " + querySource);
                            synLog.traceOrDebug("Prepare an expression for the query ");
                        }

                        xQueryExecutable = cachedQueryCompiler.compile(querySource);
                        queryEvaluator = xQueryExecutable.load();

                        // if queryEvaluator is created then put it in to cachedXQueryEvaluatorMap
                        if (isQueryKeyGenerated) {
                            cachedXQueryEvaluatorMap.put(generatedQueryKey, queryEvaluator);
                        }

                        // need set because the expression just has recreated
                        needSet = true;


                    } else {

                        Object o = synCtx.getEntry(generatedQueryKey);
                        if (o == null) {
                            if (synLog.isTraceOrDebugEnabled()) {
                                synLog.traceOrDebug("Couldn't find the xquery source with a key "
                                        + queryKey);
                            }
                            throw new SynapseException("No object found for the key '" + generatedQueryKey + "'");
                        }

                        String sourceCode = null;
                        InputStream inputStream = null;
                        if (o instanceof OMElement) {
                            sourceCode = ((OMElement) (o)).getText();
                        } else if (o instanceof String) {
                            sourceCode = (String) o;
                        } else if (o instanceof OMText) {
                            DataHandler dataHandler = (DataHandler) ((OMText) o).getDataHandler();
                            if (dataHandler != null) {
                                try {
                                    inputStream = dataHandler.getInputStream();
                                    if (inputStream == null) {
                                        if (synLog.isTraceOrDebugEnabled()) {
                                            synLog.traceOrDebug("Couldn't get" +
                                                    " the stream from the xquery source with a key "
                                                    + queryKey);
                                        }
                                        return;
                                    }

                                } catch (IOException e) {
                                    handleException("Error in reading content as a stream ");
                                }
                            }
                        }

                        if ((sourceCode == null || "".equals(sourceCode)) && inputStream == null) {
                            if (synLog.isTraceOrDebugEnabled()) {
                                synLog.traceOrDebug("Couldn't find the xquery source with a key "
                                        + queryKey);
                            }
                            return;
                        }

                        if (synLog.isTraceOrDebugEnabled()) {
                            synLog.traceOrDebug("Picked up the xquery source from the " +
                                    "key " + queryKey);
                            synLog.traceOrDebug("Prepare an expression for the query ");
                        }

                        try {
                            if (sourceCode != null) {
                                //create an xQueryExecutable using the query source
                                xQueryExecutable = cachedQueryCompiler.compile(sourceCode);
                            } else {
                                xQueryExecutable = cachedQueryCompiler.compile(inputStream);
                            }

                        } catch (IOException e) {
                            handleException("Error during the query inputStream compilation");
                        }
                        queryEvaluator = xQueryExecutable.load();

                        // if queryEvaluator is created then put it in to cachedXQueryEvaluatorMap
                        if (isQueryKeyGenerated) {
                            cachedXQueryEvaluatorMap.put(generatedQueryKey, queryEvaluator);
                        }

                        // need set because the evaluator just has recreated
                        needSet = true;
                    }
                }

                //Set the external variables to the queryEvaluator
                if (variables != null && !variables.isEmpty()) {
                    synLog.traceOrDebug("Binding  external variables to the DynamicContext");
                    for (MediatorVariable variable : variables) {
                        if (variable != null) {
                            boolean hasValueChanged = variable.evaluateValue(synCtx);
                            //if the value has changed or need set because the evaluator has recreated
                            if (hasValueChanged || needSet) {
                                //Set the external variable to the queryEvaluator
                                setVariable(queryEvaluator, variable, synLog);
                            }
                        }
                    }
                }

                //executing the query
                xdmValue = queryEvaluator.evaluate();

            }

            if (queryEvaluator == null) {
                synLog.traceOrDebug("Result Sequence is null");
                return;
            }

            //processing the result
            for (XdmItem xdmItem : xdmValue) {

                if (xdmItem == null) {
                    return;
                }

                XdmNodeKind xdmNodeKind = null;
                ItemType itemType = null;
                if (xdmItem.isAtomicValue()) {
                    itemType = getItemType(xdmItem, cachedProcessor);
                    if (itemType == null) {
                        return;
                    }
                } else {
                    xdmNodeKind = ((XdmNode) xdmItem).getNodeKind();
                }

                if (synLog.isTraceOrDebugEnabled()) {
                    synLog.traceOrDebug("The XQuery Result " + xdmItem.toString());
                }

                //The target node that is going to modify
                OMNode destination = target.selectOMNode(synCtx, synLog);
                if (destination != null) {
                    if (synLog.isTraceOrDebugEnabled()) {
                        synLog.traceOrDebug("The target node " + destination);
                    }

                    //If the result is XML
                    if (XdmNodeKind.DOCUMENT == xdmNodeKind || XdmNodeKind.ELEMENT == xdmNodeKind) {

                        StAXOMBuilder builder = new StAXOMBuilder(XMLInputFactory.newInstance().createXMLStreamReader(
                                new StringReader(xdmItem.toString())));
                        OMElement resultOM = builder.getDocumentElement();

                        if (resultOM != null) {
                            //replace the target node from the result
                            destination.insertSiblingAfter(resultOM);
                            destination.detach();
                        }
                    } else if (ItemType.INTEGER == itemType || ItemType.INT == itemType) {
                        //replace the text value of the target node by the result ,If the result is
                        // a basic type
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getDecimalValue().intValue()));
                    } else if (ItemType.BOOLEAN == itemType) {
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getBooleanValue()));
                    } else if (ItemType.DOUBLE == itemType) {
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getDoubleValue()));
                    } else if (ItemType.FLOAT == itemType) {
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getDecimalValue().floatValue()));
                    } else if (ItemType.LONG == itemType) {
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getLongValue()));
                    } else if (ItemType.SHORT == itemType) {
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getDecimalValue().shortValue()));
                    } else if (ItemType.BYTE == itemType) {
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getDecimalValue().byteValue()));
                    } else if (ItemType.STRING == itemType) {
                        ((OMElement) destination).setText(String.valueOf(((XdmAtomicValue) xdmItem).getValue()));
                    }

                } else if (null == target.getXPath() && null == destination) {
                    //In the case soap body doesn't have the first element --> Empty soap body
                    destination = synCtx.getEnvelope().getBody();

                    if (synLog.isTraceOrDebugEnabled()) {
                        synLog.traceOrDebug("The target node " + destination);
                    }

                    //If the result is XML
                    if (XdmNodeKind.ELEMENT == xdmNodeKind || XdmNodeKind.DOCUMENT == xdmNodeKind) {
                        StAXOMBuilder builder = new StAXOMBuilder(
                                XMLInputFactory.newInstance().createXMLStreamReader(
                                        new StringReader(xdmItem.toString())));
                        OMElement resultOM = builder.getDocumentElement();
                        if (resultOM != null) {
                            ((OMElement) destination).addChild(resultOM);
                        }
                    }
                    //No else part since soap body could have only XML part not text values

                }
                break;   // Only take the *first* value of the result sequence
            }
            queryEvaluator.close();  // closing the result sequence
        } catch (SaxonApiException e) {
            handleException("Error during the querying " + e.getMessage(), e);
        } catch (XMLStreamException e) {
            handleException("Error during retrieving  the Document Node as  the result "
                    + e.getMessage(), e);
        }
    }

    /**
     * Binding a variable to the Dynamic Context in order to available during doing the querying
     *
     * @param queryEvaluator   The XQuery evaluator to which the variable will be added
     * @param variable         The variable which contains the name and vaule for adding
     * @param synLog           the Synapse log to use
     * @throws SaxonApiException throws if any error occurs when adding the variable
     */
    private void setVariable(XQueryEvaluator queryEvaluator, MediatorVariable variable,
                             SynapseLog synLog) throws SaxonApiException {
        QName name = new QName(variable.getName().getLocalPart());

        if (variable != null) {

            ItemType type = variable.getType();
            XdmNodeKind nodeKind = variable.getNodeKind();
            Object value = variable.getValue();

            if (value != null && (type != null || nodeKind != null)) {

                if (synLog.isTraceOrDebugEnabled()) {
                    synLog.traceOrDebug("Binding a variable to the DynamicContext with a name : "
                            + name + " and a value : " + value);
                }

                //Binding the basic type As-Is and XML element as an InputSource
                if (ItemType.BOOLEAN == type) {
                    boolean booleanValue = false;
                    if (value instanceof String) {
                        booleanValue = Boolean.parseBoolean((String) value);
                    } else if (value instanceof Boolean) {
                        booleanValue = (Boolean) value;
                    } else {
                        handleException("Incompatible type for the Boolean");
                    }
                    queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(booleanValue), ItemType.BOOLEAN));
                } else if (ItemType.INTEGER == type) {
                    int intValue = -1;
                    if (value instanceof String) {
                        try {
                            intValue = Integer.parseInt((String) value);
                        } catch (NumberFormatException e) {
                            handleException("Incompatible value '" + value + "' " +
                                    "for the Integer", e);
                        }
                    } else if (value instanceof Integer) {
                        intValue = (Integer) value;
                    } else {
                        handleException("Incompatible type for the Integer");
                    }
                    if (intValue != -1) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(intValue), ItemType.INTEGER));
                    }
                } else if (ItemType.INT == type) {
                    int intValue = -1;
                    if (value instanceof String) {
                        try {
                            intValue = Integer.parseInt((String) value);
                        } catch (NumberFormatException e) {
                            handleException("Incompatible value '" + value +
                                    "' for the Int", e);
                        }
                    } else if (value instanceof Integer) {
                        intValue = (Integer) value;
                    } else {
                        handleException("Incompatible type for the Int");
                    }
                    if (intValue != -1) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(intValue), ItemType.INT));
                    }

                } else if (ItemType.LONG == type) {
                    long longValue = -1;
                    if (value instanceof String) {
                        try {
                            longValue = Long.parseLong((String) value);
                        } catch (NumberFormatException e) {
                            handleException("Incompatible value '" + value + "' " +
                                    "for the long ", e);
                        }
                    } else if (value instanceof Long) {
                        longValue = (Long) value;
                    } else {
                        handleException("Incompatible type for the Long");
                    }
                    if (longValue != -1) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(longValue), ItemType.LONG));
                    }
                } else if (ItemType.SHORT == type) {
                    short shortValue = -1;
                    if (value instanceof String) {
                        try {
                            shortValue = Short.parseShort((String) value);
                        } catch (NumberFormatException e) {
                            handleException("Incompatible value '" + value + "' " +
                                    "for the short ", e);
                        }
                    } else if (value instanceof Short) {
                        shortValue = (Short) value;
                    } else {
                        handleException("Incompatible type for the Short");
                    }
                    if (shortValue != -1) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(shortValue), ItemType.SHORT));
                    }
                } else if (ItemType.DOUBLE == type) {
                    double doubleValue = -1;
                    if (value instanceof String) {
                        try {
                            doubleValue = Double.parseDouble((String) value);
                        } catch (NumberFormatException e) {
                            handleException("Incompatible value '" + value + "' " +
                                    "for the double ", e);
                        }
                    } else if (value instanceof Double) {
                        doubleValue = (Double) value;
                    } else {
                        handleException("Incompatible type for the Double");
                    }
                    if (doubleValue != -1) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(doubleValue), ItemType.DOUBLE));
                    }
                } else if (ItemType.FLOAT == type) {
                    float floatValue = -1;
                    if (value instanceof String) {
                        try {
                            floatValue = Float.parseFloat((String) value);
                        } catch (NumberFormatException e) {
                            handleException("Incompatible value '" + value + "' " +
                                    "for the float ", e);
                        }
                    } else if (value instanceof Float) {
                        floatValue = (Float) value;
                    } else {
                        handleException("Incompatible type for the Float");
                    }
                    if (floatValue != -1) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(floatValue), ItemType.FLOAT));
                    }
                } else if (ItemType.BYTE == type) {
                    byte byteValue = -1;
                    if (value instanceof String) {
                        try {
                            byteValue = Byte.parseByte((String) value);
                        } catch (NumberFormatException e) {
                            handleException("Incompatible value '" + value + "' " +
                                    "for the byte ", e);
                        }
                    } else if (value instanceof Byte) {
                        byteValue = (Byte) value;
                    } else {
                        handleException("Incompatible type for the Byte");
                    }
                    if (byteValue != -1) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(byteValue), ItemType.BYTE));
                    }
                } else if (ItemType.STRING == type) {
                    if (value instanceof String) {
                        queryEvaluator.setExternalVariable(name, new XdmAtomicValue(String.valueOf(value), ItemType.STRING));
                    } else {
                        handleException("Incompatible type for the String");
                    }
                } else if (XdmNodeKind.DOCUMENT == nodeKind || XdmNodeKind.ELEMENT == nodeKind) {
                    setOMNode(name, value, queryEvaluator, cachedProcessor);
                } else {
                    handleException("Unsupported  type for the binding type" + type +
                            " in the variable name " + name);
                }
            }
        } else {
                /*
                The following block will invariably result in  "javax.xml.xquery.XQException: Argument value is null",
                    which is the proper behaviour for this case.
                This block was added to fix the issue where if a variable with non-null value is passed for evaluation
                    and if the subsequent request has a null value for the same variable (i.e. with the same name), the
                    previous non-null value is given out instead of considering the null-case.
                */
            if (synLog.isTraceOrDebugEnabled()) {
                synLog.traceOrDebug("Null variable value encountered for variable name: " + name);
            }
            queryEvaluator.setExternalVariable(name, null);
        }

    }

    private void setOMNode(QName name, Object value,
                           XQueryEvaluator queryEvaluator, Processor processor) throws SaxonApiException{

        OMElement variableValue = null;
        if (value instanceof String) {
            variableValue = SynapseConfigUtils.stringToOM((String) value);
        } else if (value instanceof OMElement) {
            variableValue = (OMElement) value;
        }

        if (variableValue != null) {
            DocumentBuilder documentBuilder = processor.newDocumentBuilder();
            if (useDOMSource) {
                XdmNode xdmNode = documentBuilder.build(new DOMSource(((Element) ElementHelper.
                        importOMElement(variableValue,
                                DOOMAbstractFactory.getOMFactory())).
                        getOwnerDocument()));
                queryEvaluator.setExternalVariable(name, xdmNode);
            } else {
                StreamSource streamSource = new StreamSource(SynapseConfigUtils.getInputStream(variableValue));
                XdmNode xdmNode = documentBuilder.build(streamSource);
                queryEvaluator.setExternalVariable(name, xdmNode);
            }
        }
    }

    private static ItemType getItemType(XdmItem item, Processor process)throws SaxonApiException{
        return new ItemTypeFactory(process).getAtomicType(((XdmAtomicValue) item).getPrimitiveTypeName());
    }

    private void handleException(String msg, Exception e) {
        log.error(msg, e);
        throw new SynapseException(msg, e);
    }

    private void handleException(String msg) {
        log.error(msg);
        throw new SynapseException(msg);
    }

    public Value getQueryKey() {
        return queryKey;
    }

    public void setQueryKey(Value queryKey) {
        this.queryKey = queryKey;
    }

    public String getQuerySource() {
        return querySource;
    }

    public void setQuerySource(String querySource) {
        this.querySource = querySource;
    }

    public void addAllVariables(List<MediatorVariable> list) {
        this.variables.addAll(list);
    }

    public void addVariable(MediatorVariable variable) {
        this.variables.add(variable);
    }

    public List<MediatorProperty> getProcessorProperties() {
        return processorProperties;
    }

    public List<MediatorVariable> getVariables() {
        return variables;
    }

    public SynapseXPath getTarget() {
        return target.getXPath();
    }

    public void setTarget(SynapseXPath source) {
        this.target.setXPath(source);
    }

    public void addAllDataSourceProperties(List<MediatorProperty> list) {
        this.processorProperties.addAll(list);
    }

    public boolean isUseDOMSource() {
        return useDOMSource;
    }

    public void setUseDOMSource(boolean useDOMSource) {
        this.useDOMSource = useDOMSource;
    }

    public boolean isContentAltering() {
        return true;
    }

}
