/*
 * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.org) All Rights Reserved.
 *
 * WSO2 Inc. 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.bsf;

import org.apache.axiom.om.OMAttribute;
import org.apache.axiom.om.OMElement;
import org.apache.axiom.om.OMException;
import org.apache.axiom.om.OMNode;
import org.apache.axiom.om.OMXMLBuilderFactory;
import org.apache.axiom.om.OMXMLParserWrapper;
import org.apache.axiom.om.xpath.AXIOMXPath;
import org.apache.axiom.soap.SOAPBody;
import org.apache.axiom.soap.SOAPEnvelope;
import org.apache.axiom.soap.SOAPFactory;
import org.apache.axiom.soap.SOAPHeader;
import org.apache.axiom.soap.SOAPHeaderBlock;
import org.apache.axis2.AxisFault;
import org.apache.axis2.addressing.EndpointReference;
import org.apache.axis2.addressing.RelatesTo;
import org.apache.axis2.context.OperationContext;
import org.apache.bsf.xml.XMLHelper;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.http.protocol.HTTP;
import org.apache.synapse.ContinuationState;
import org.apache.synapse.FaultHandler;
import org.apache.synapse.Mediator;
import org.apache.synapse.MessageContext;
import org.apache.synapse.SynapseConstants;
import org.apache.synapse.commons.json.JsonUtil;
import org.apache.synapse.config.SynapseConfiguration;
import org.apache.synapse.config.xml.XMLConfigConstants;
import org.apache.synapse.core.SynapseEnvironment;
import org.apache.synapse.core.axis2.Axis2MessageContext;
import org.apache.synapse.endpoints.Endpoint;
import org.apache.synapse.mediators.bsf.utils.ScriptUtils;
import org.graalvm.polyglot.Context;
import org.graalvm.polyglot.Value;
import org.jaxen.JaxenException;
import org.w3c.dom.Document;
import org.xml.sax.SAXException;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.Stack;

import javax.script.ScriptEngine;
import javax.script.ScriptException;

public class GraalVMJavaScriptMessageContext implements ScriptMessageContext {

    private static final Log logger = LogFactory.getLog(GraalVMJavaScriptMessageContext.class.getName());

    private static final String JSON_OBJECT = "JSON_OBJECT";
    private static final String JSON_TEXT = "JSON_TEXT";

    /**
     * The actual Synapse message context reference.
     */
    private final MessageContext mc;

    /**
     * The OMElement to scripting language object converter for the selected language.
     */
    private final XMLHelper xmlHelper;

    /**
     * To keep Script Engine instance.
     */
    private ScriptEngine scriptEngine;

    /**
     * Reference to an empty JSON object.
     */
    private Object emptyJsonObject;

    /**
     * Reference to JSON object which is used to serialize json.
     */
    private Value jsonSerializer;

    public GraalVMJavaScriptMessageContext(MessageContext mc, XMLHelper xmlHelper, Context context) throws
            ScriptException {
        this.mc = mc;
        this.xmlHelper = xmlHelper;
        this.emptyJsonObject = context.eval("js", "({})");
        this.jsonSerializer = context.eval("js", "JSON");
    }

    public Object jsonSerializerCallMember(String key, String value) {
        return jsonSerializer.invokeMember(key, value);
    }

    /**
     * Get the XML representation of SOAP Body payload.
     * The payload is the first element inside the SOAP <Body> tags
     *
     * @return the XML SOAP Body
     */
    public Object getPayloadXML() throws ScriptException {
        return xmlHelper.toScriptXML(mc.getEnvelope().getBody().getFirstElement());
//        return mc.getEnvelope().getBody().getFirstElement();
    }

    /**
     * Set the SOAP body payload from XML.
     *
     * @param payload Message payload
     * @throws ScriptException For errors in converting xml To OM
     * @throws OMException     For errors in OM manipulation
     */
    public void setPayloadXML(Object payload) throws OMException, ScriptException {
        SOAPBody body = mc.getEnvelope().getBody();
        OMElement firstChild = body.getFirstElement();
        OMElement omElement = xmlHelper.toOMElement(payload);
        if (firstChild == null) {
            body.addChild(omElement);
        } else {
            firstChild.insertSiblingAfter(omElement);
            firstChild.detach();
        }
    }

    /**
     * Get the JSON object representation of the JSON message body of the request.
     *
     * @return JSON object of the message body
     */
    public Object getPayloadJSON() {
        return jsonObject(mc);
    }

    /**
     * Saves the payload of this message context as a JSON payload.
     *
     * @param jsonPayload Javascript native object to be set as the message body
     * @throws ScriptException in case of creating a JSON object out of the javascript native object.
     */
    public void setPayloadJSON(Object jsonPayload) throws ScriptException {
        try {
            String jsonString = jsonSerializer.invokeMember("stringify", jsonPayload).toString();
            InputStream stream = new ByteArrayInputStream(jsonString.getBytes(StandardCharsets.UTF_8));
            org.apache.axis2.context.MessageContext messageContext;
            messageContext = ((Axis2MessageContext) mc).getAxis2MessageContext();
            JsonUtil.getNewJsonPayload(messageContext, stream, true, true);
            messageContext.setProperty(JSON_OBJECT, jsonPayload);
        } catch (AxisFault axisFault) {
            throw new ScriptException(axisFault);
        }
    }

    /**
     * Get the Message Payload as a text.
     *
     * @return Payload as text
     */
    public Object getJsonText() {
        if (mc == null) {
            return "";
        }
        Object text = mc.getProperty(JSON_TEXT);
        return text == null ? "{}" : text;
    }

    /**
     * Get the Message Payload as a text.
     *
     * @return Payload as text
     */
    public String getPayloadText() {
        if (JsonUtil.hasAJsonPayload(((Axis2MessageContext) mc).getAxis2MessageContext())) {
            return JsonUtil.jsonPayloadToString(((Axis2MessageContext) mc).getAxis2MessageContext());
        } else {
            return mc.getEnvelope().toString();
        }
    }

    /**
     * Saves the JavaScript Object to the message context.
     *
     * @param messageContext The message context of the sequence
     * @param jsonObject     JavaScript Object which is passed to be saved in message context
     * @return true
     */
    public boolean setJsonObject(MessageContext messageContext, Object jsonObject) {
        messageContext.setProperty(JSON_OBJECT, jsonObject);
        return true;
    }

    /**
     * Saves the JSON String to the message context.
     *
     * @param messageContext The message context of the sequence
     * @param jsonObject     JavaScript string which is passed to be saved in message context
     * @return false if messageContext is null return true otherwise
     */
    public boolean setJsonText(MessageContext messageContext, Object jsonObject) {
        if (messageContext == null) {
            return false;
        }
        messageContext.setProperty(JSON_TEXT, jsonObject);
        return true;
    }

    /**
     * Returns the JavaScript Object saved in this message context.
     *
     * @param messageContext The message context of the sequence
     * @return o JavaScript Object saved in this message context
     */
    public Object jsonObject(MessageContext messageContext) {
        if (messageContext == null) {
            return null;
        }
        Object jsonObject = messageContext.getProperty(JSON_OBJECT);
        if (jsonObject == null) {
            return emptyJsonObject;
        }
        return jsonObject;
    }

    /**
     * Set a script engine.
     *
     * @param scriptEngine a ScriptEngine instance
     */
    public void setScriptEngine(ScriptEngine scriptEngine) {
        this.scriptEngine = scriptEngine;
    }

    /**
     * Returns the parsed xml document.
     *
     * @param text xml string or document needed to be parser
     * @return parsed document
     */
    public Document parseXml(String text) throws ScriptException {
        Document doc;
        try {
            doc = ScriptUtils.parseXml(text);
        } catch (SAXException | IOException e) {
            ScriptException scriptException = new ScriptException("Failed to parse provided xml");
            scriptException.initCause(e);
            throw scriptException;
        }

        return doc;
    }

    /**
     * Returns the parsed xml document.
     *
     * @param stream input stream of xml string or document needed to be parsed
     * @return parsed document
     */
    public OMElement getParsedOMElement(InputStream stream) {
        OMXMLParserWrapper builder = OMXMLBuilderFactory.createOMBuilderWithSec(stream);
        return builder.getDocumentElement();
    }

    /**
     * Returns the Axiom xpath.
     *
     * @param expression Xpath expression
     * @return Axiom xpath is returned
     */
    public AXIOMXPath getXpathResult(String expression) throws JaxenException {
        return new AXIOMXPath(expression);
    }

    /**
     * Add a new SOAP header to the message.
     *
     * @param mustUnderstand the value for the <code>soapenv:mustUnderstand</code> attribute
     * @param content        the XML for the new header
     * @throws ScriptException if an error occurs when converting the XML to OM
     */
    public void addHeader(boolean mustUnderstand, Object content) throws ScriptException {
        SOAPEnvelope envelope = mc.getEnvelope();
        SOAPFactory factory = (SOAPFactory) envelope.getOMFactory();
        SOAPHeader header = envelope.getHeader();
        if (header == null) {
            header = factory.createSOAPHeader(envelope);
        }

        OMElement element = xmlHelper.toOMElement(content);
        // We can't add the element directly to the SOAPHeader. Instead, we need to copy the
        // information over to a SOAPHeaderBlock.
        SOAPHeaderBlock headerBlock = header.addHeaderBlock(element.getLocalName(), element.getNamespace());
        for (Iterator it = element.getAllAttributes(); it.hasNext(); ) {
            headerBlock.addAttribute((OMAttribute) it.next());
        }
        headerBlock.setMustUnderstand(mustUnderstand);
        OMNode child = element.getFirstOMChild();
        while (child != null) {
            // Get the next child before addChild will detach the node from its original place.
            OMNode next = child.getNextOMSibling();
            headerBlock.addChild(child);
            child = next;
        }
    }

    /**
     * Get the XML representation of the complete SOAP envelope.
     *
     * @return return an object that represents the payload in the current scripting language
     * @throws ScriptException in-case of an error in getting
     *                         the XML representation of SOAP envelope
     */
    public Object getEnvelopeXML() throws ScriptException {
        SOAPEnvelope envelope = mc.getEnvelope();
        return envelope.toString();
    }

    /**
     * {@inheritDoc}
     */
    public SynapseConfiguration getConfiguration() {
        return mc.getConfiguration();
    }

    /**
     * {@inheritDoc}
     */
    public void setConfiguration(SynapseConfiguration cfg) {
        mc.setConfiguration(cfg);
    }

    /**
     * {@inheritDoc}
     */
    public SynapseEnvironment getEnvironment() {
        return mc.getEnvironment();
    }

    /**
     * {@inheritDoc}
     */
    public void setEnvironment(SynapseEnvironment se) {
        mc.setEnvironment(se);
    }

    /**
     * {@inheritDoc}
     */
    public Map<String, Object> getContextEntries() {
        return mc.getContextEntries();
    }

    /**
     * {@inheritDoc}
     */
    public void setContextEntries(Map<String, Object> entries) {
        mc.setContextEntries(entries);
    }

    /**
     * {@inheritDoc}
     */
    public Object getProperty(String key) {
        return mc.getProperty(key);
    }

    /**
     * {@inheritDoc}
     */
    public Object getEntry(String key) {
        return mc.getEntry(key);
    }

    /**
     * {@inheritDoc}
     */
    public Object getLocalEntry(String key) {
        return mc.getLocalEntry(key);
    }

    /**
     * Add a new property to the message.
     *
     * @param key   unique identifier of property
     * @param value value of property
     */
    public void setProperty(String key, Object value) {
        try {
            OMElement omElement = xmlHelper.toOMElement(value);
            mc.setProperty(key, omElement);
        } catch (ScriptException e) {
            //Try to convert the value into OMElement if it fails it means value is not a representation of xml so
            // set as key value pair
            mc.setProperty(key, value);
        }
    }

    /**
     * Add a new property to the message.
     *
     * @param key   unique identifier of property
     * @param value value of property
     * @param scope scope of the property
     */
    public void setProperty(String key, Object value, String scope) {
        if (scope == null || XMLConfigConstants.SCOPE_DEFAULT.equals(scope)) {
            setProperty(key, value);
        } else if (XMLConfigConstants.SCOPE_AXIS2.equals(scope)) {
            //Setting property into the  Axis2 Message Context
            Axis2MessageContext axis2smc = (Axis2MessageContext) mc;
            org.apache.axis2.context.MessageContext axis2MessageCtx = axis2smc.getAxis2MessageContext();
            axis2MessageCtx.setProperty(key, value);
            handleSpecialProperties(key, value, axis2MessageCtx);

        } else if (XMLConfigConstants.SCOPE_TRANSPORT.equals(scope)) {
            //Setting Transport Headers
            Axis2MessageContext axis2smc = (Axis2MessageContext) mc;
            org.apache.axis2.context.MessageContext axis2MessageCtx = axis2smc.getAxis2MessageContext();
            Object headers = axis2MessageCtx.getProperty(org.apache.axis2.context.MessageContext.TRANSPORT_HEADERS);

            if (headers != null && headers instanceof Map) {
                Map headersMap = (Map) headers;
                headersMap.put(key, value);
            }
            if (headers == null) {
                Map headersMap = new HashMap();
                headersMap.put(key, value);
                axis2MessageCtx.setProperty(org.apache.axis2.context.MessageContext.TRANSPORT_HEADERS, headersMap);
            }
        } else if (XMLConfigConstants.SCOPE_OPERATION.equals(scope)) {
            Axis2MessageContext axis2smc = (Axis2MessageContext) mc;
            org.apache.axis2.context.MessageContext axis2MessageCtx = axis2smc.getAxis2MessageContext();
            axis2MessageCtx.getOperationContext().setProperty(key, value);
        }
    }

    /**
     * Remove property from the message.
     *
     * @param key   unique identifier of property
     * @param scope scope of the property
     */
    public void removeProperty(String key, String scope) {
        if (scope == null || XMLConfigConstants.SCOPE_DEFAULT.equals(scope)) {
            Set pros = mc.getPropertyKeySet();
            if (pros != null) {
                pros.remove(key);
            }
        } else if (XMLConfigConstants.SCOPE_AXIS2.equals(scope)) {
            //Removing property from the Axis2 Message Context
            Axis2MessageContext axis2smc = (Axis2MessageContext) mc;
            org.apache.axis2.context.MessageContext axis2MessageCtx = axis2smc.getAxis2MessageContext();
            axis2MessageCtx.removeProperty(key);

        } else if (XMLConfigConstants.SCOPE_TRANSPORT.equals(scope)) {
            // Removing transport headers
            Axis2MessageContext axis2smc = (Axis2MessageContext) mc;
            org.apache.axis2.context.MessageContext axis2MessageCtx = axis2smc.getAxis2MessageContext();
            Object headers = axis2MessageCtx.getProperty(org.apache.axis2.context.MessageContext.TRANSPORT_HEADERS);
            if (headers != null && headers instanceof Map) {
                Map headersMap = (Map) headers;
                headersMap.remove(key);
            }
        } else if (XMLConfigConstants.SCOPE_OPERATION.equals(scope)) {
            // Removing operation scope headers
            Axis2MessageContext axis2smc = (Axis2MessageContext) mc;
            org.apache.axis2.context.MessageContext axis2MessageCtx = axis2smc.getAxis2MessageContext();
            OperationContext axis2oc = axis2MessageCtx.getOperationContext();
            axis2oc.removeProperty(key);
        }

    }

    /**
     * Add special properties such as content type to the message context.
     *
     * @param key            unique identifier of property
     * @param value          value of property
     * @param messageContext Axis2 message context
     */
    private void handleSpecialProperties(String key, Object value,
                                         org.apache.axis2.context.MessageContext messageContext) {
        if (org.apache.axis2.Constants.Configuration.MESSAGE_TYPE.equals(key)) {
            messageContext.setProperty(org.apache.axis2.Constants.Configuration.CONTENT_TYPE, value);
            Object o = messageContext.getProperty(org.apache.axis2.context.MessageContext.TRANSPORT_HEADERS);
            Map headers = (Map) o;
            if (headers != null) {
                headers.put(HTTP.CONTENT_TYPE, value);
            }
        }
    }

    /**
     * {@inheritDoc}
     */
    public Set getPropertyKeySet() {
        return mc.getPropertyKeySet();
    }

    /**
     * {@inheritDoc}
     */
    public Mediator getMainSequence() {
        return mc.getMainSequence();
    }

    /**
     * {@inheritDoc}
     */
    public Mediator getFaultSequence() {
        return mc.getFaultSequence();
    }

    /**
     * {@inheritDoc}
     */
    public Mediator getSequence(String key) {
        return mc.getSequence(key);
    }

    /**
     * {@inheritDoc}
     */
    public OMElement getFormat(String s) {
        return mc.getFormat(s);
    }

    /**
     * {@inheritDoc}
     */
    public Endpoint getEndpoint(String key) {
        return mc.getEndpoint(key);
    }

    /**
     * {@inheritDoc}
     */
    public SOAPEnvelope getEnvelope() {
        return mc.getEnvelope();
    }

    /**
     * {@inheritDoc}
     */
    public void setEnvelope(SOAPEnvelope envelope) throws AxisFault {
        mc.setEnvelope(envelope);
    }

    /**
     * {@inheritDoc}
     */
    public EndpointReference getFaultTo() {
        return mc.getFaultTo();
    }

    /**
     * This is used to set the value which specifies the receiver of the faults relating to the message.
     *
     * @param reference specifies the receiver of the faults relating to the message
     */
    public void setFaultTo(String reference) {
        mc.setFaultTo(new EndpointReference(reference));
    }

    /**
     * {@inheritDoc}
     */
    public void setFaultTo(EndpointReference reference) {
        mc.setFaultTo(reference);
    }

    /**
     * {@inheritDoc}
     */
    public EndpointReference getFrom() {
        return mc.getFrom();
    }

    /**
     * This is used to set the value which specifies the sender of the message.
     *
     * @param reference specifies the sender of the message
     */
    public void setFrom(String reference) {
        mc.setFrom(new EndpointReference(reference));
    }

    /**
     * {@inheritDoc}
     */
    public void setFrom(EndpointReference reference) {
        mc.setFrom(reference);
    }

    /**
     * {@inheritDoc}
     */
    public String getMessageID() {
        return mc.getMessageID();
    }

    /**
     * {@inheritDoc}
     */
    public void setMessageID(String string) {
        mc.setMessageID(string);
    }

    /**
     * {@inheritDoc}
     */
    public RelatesTo getRelatesTo() {
        return mc.getRelatesTo();
    }

    /**
     * {@inheritDoc}
     */
    public void setRelatesTo(RelatesTo[] reference) {
        mc.setRelatesTo(reference);
    }

    /**
     * {@inheritDoc}
     */
    public EndpointReference getReplyTo() {
        return mc.getReplyTo();
    }

    /**
     * This is used to set the value which specifies the receiver of the replies to the message.
     *
     * @param reference specifies the receiver of the replies to the message
     */
    public void setReplyTo(String reference) {
        mc.setReplyTo(new EndpointReference(reference));
    }

    /**
     * {@inheritDoc}
     */
    public void setReplyTo(EndpointReference reference) {
        mc.setReplyTo(reference);
    }

    /**
     * {@inheritDoc}
     */
    public EndpointReference getTo() {
        return mc.getTo();
    }

    /**
     * This is used to set the value which specifies the receiver of the message.
     *
     * @param reference specifies the receiver of the message
     */
    public void setTo(String reference) {
        mc.setTo(new EndpointReference(reference));
    }

    /**
     * {@inheritDoc}
     */
    public void setTo(EndpointReference reference) {
        mc.setTo(reference);
    }

    /**
     * {@inheritDoc}
     */
    public String getWSAAction() {
        return mc.getWSAAction();
    }

    /**
     * {@inheritDoc}
     */
    public void setWSAAction(String actionURI) {
        mc.setWSAAction(actionURI);
    }

    /**
     * {@inheritDoc}
     */
    public String getSoapAction() {
        return mc.getSoapAction();
    }

    /**
     * {@inheritDoc}
     */
    public void setSoapAction(String string) {
        mc.setSoapAction(string);
    }

    /**
     * {@inheritDoc}
     */
    public String getWSAMessageID() {
        return mc.getWSAMessageID();
    }

    /**
     * {@inheritDoc}
     */
    public void setWSAMessageID(String messageID) {
        mc.setWSAMessageID(messageID);
    }

    /**
     * {@inheritDoc}
     */
    public boolean isDoingMTOM() {
        return mc.isDoingMTOM();
    }

    /**
     * {@inheritDoc}
     */
    public void setDoingMTOM(boolean b) {
        mc.setDoingMTOM(b);
    }

    /**
     * {@inheritDoc}
     */
    public boolean isDoingSWA() {
        return mc.isDoingSWA();
    }

    /**
     * {@inheritDoc}
     */
    public void setDoingSWA(boolean b) {
        mc.setDoingSWA(b);
    }

    /**
     * {@inheritDoc}
     */
    public boolean isDoingPOX() {
        return mc.isDoingPOX();
    }

    /**
     * {@inheritDoc}
     */
    public void setDoingPOX(boolean b) {
        mc.setDoingPOX(b);
    }

    /**
     * {@inheritDoc}
     */
    public boolean isDoingGET() {
        return mc.isDoingGET();
    }

    /**
     * {@inheritDoc}
     */
    public void setDoingGET(boolean b) {
        mc.setDoingGET(b);
    }

    /**
     * {@inheritDoc}
     */
    public boolean isSOAP11() {
        return mc.isSOAP11();
    }

    /**
     * {@inheritDoc}
     */
    public boolean isResponse() {
        return mc.isResponse();
    }

    /**
     * {@inheritDoc}
     */
    public void setResponse(boolean b) {
        mc.setResponse(b);
    }

    /**
     * {@inheritDoc}
     */
    public boolean isFaultResponse() {
        return mc.isFaultResponse();
    }

    /**
     * {@inheritDoc}
     */
    public void setFaultResponse(boolean b) {
        mc.setFaultResponse(b);
    }

    /**
     * {@inheritDoc}
     */
    public int getTracingState() {
        return mc.getTracingState();
    }

    /**
     * {@inheritDoc}
     */
    public void setTracingState(int tracingState) {
        mc.setTracingState(tracingState);
    }

    /**
     * {@inheritDoc}
     */
    public Stack<FaultHandler> getFaultStack() {
        return mc.getFaultStack();
    }

    /**
     * {@inheritDoc}
     */
    public void pushFaultHandler(FaultHandler fault) {
        mc.pushFaultHandler(fault);
    }

    /**
     * {@inheritDoc}
     */
    public void pushContinuationState(ContinuationState continuationState) {
        mc.pushContinuationState(continuationState);
    }

    /**
     * {@inheritDoc}
     */
    public Stack<ContinuationState> getContinuationStateStack() {
        return mc.getContinuationStateStack();
    }

    @Override
    public void restoreContinuationStateStack(Stack<ContinuationState> continuationStateStack) {
        mc.restoreContinuationStateStack(continuationStateStack);
    }

    /**
     * {@inheritDoc}
     */
    public boolean isContinuationEnabled() {
        return false;
    }

    /**
     * {@inheritDoc}
     */
    public void setContinuationEnabled(boolean contStateStackEnabled) {
    }

    /**
     * {@inheritDoc}
     */
    public Log getServiceLog() {
        return LogFactory.getLog(GraalVMJavaScriptMessageContext.class);
    }

    /**
     * {@inheritDoc}
     */
    public Mediator getSequenceTemplate(String key) {
        return mc.getSequenceTemplate(key);
    }

    /**
     * {@inheritDoc}
     */
    public Mediator getDefaultConfiguration(String arg0) {
        return mc.getDefaultConfiguration(arg0);
    }

    /**
     * {@inheritDoc}
     */
    public String getMessageString() {
        return mc.getMessageString();
    }

    /**
     * {@inheritDoc}
     */
    public int getMessageFlowTracingState() {
        return SynapseConstants.TRACING_OFF;
    }

    /**
     * {@inheritDoc}
     */
    public void setMessageFlowTracingState(int state) {
        mc.setMessageFlowTracingState(state);
    }

    @Override
    public Object getVariable(String key) {

        return mc.getVariable(key);
    }

    @Override
    public void setVariable(String key, Object value) {

        try {
            OMElement omElement = xmlHelper.toOMElement(value);
            mc.setVariable(key, omElement);
        } catch (ScriptException e) {
            // Try to convert the value into OMElement if it fails it means value is not a representation of xml so
            // set as key value pair
            mc.setVariable(key, value);
        }
    }

    @Override
    public Set getVariableKeySet() {

        return mc.getVariableKeySet();
    }
}
