SetElementData.java

/*
	Milyn - Copyright (C) 2006 - 2010

	This library is free software; you can redistribute it and/or
	modify it under the terms of the GNU Lesser General Public
	License (version 2.1) as published by the Free Software
	Foundation.

	This library is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

	See the GNU Lesser General Public License for more details:
	http://www.gnu.org/licenses/lgpl.txt
*/
package org.milyn.visitors.set;

import org.milyn.SmooksException;
import org.milyn.cdr.Parameter;
import org.milyn.cdr.SmooksResourceConfiguration;
import org.milyn.cdr.annotation.Config;
import org.milyn.cdr.annotation.ConfigParam;
import org.milyn.container.ExecutionContext;
import org.milyn.delivery.Filter;
import org.milyn.delivery.annotation.Initialize;
import org.milyn.delivery.dom.DOMVisitAfter;
import org.milyn.delivery.sax.DefaultSAXElementSerializer;
import org.milyn.delivery.sax.SAXElement;
import org.milyn.util.FreeMarkerTemplate;
import org.milyn.xml.DomUtils;
import org.w3c.dom.Element;

import javax.xml.XMLConstants;
import javax.xml.namespace.QName;
import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Set Element Data visitor.
 * @author <a href="mailto:tom.fennelly@gmail.com">tom.fennelly@gmail.com</a>
 */
public class SetElementData extends DefaultSAXElementSerializer implements DOMVisitAfter {

    protected static final String ATTRIBUTE_DATA = "attributeData";

    private String elementQName;
    private String elementName;
    private String elementNamespace;
    private String elementNamespacePrefix;
    private Map<QName, FreeMarkerTemplate> attributes = new LinkedHashMap<QName, FreeMarkerTemplate>();

    @Config
    private SmooksResourceConfiguration resourceConfig;

    @Initialize
    public void init() {
        if(elementQName != null) {
            int nsPrefixIdx = elementQName.indexOf(":");
            if(nsPrefixIdx != -1) {
                elementNamespacePrefix = elementQName.substring(0, nsPrefixIdx);
                elementName = elementQName.substring(nsPrefixIdx + 1);

                if(elementNamespacePrefix.equals(XMLConstants.XMLNS_ATTRIBUTE) && elementNamespace == null) {
                    elementNamespace = XMLConstants.XMLNS_ATTRIBUTE_NS_URI;
                }
            } else {
                elementName = elementQName;

                if(elementName.equals(XMLConstants.XMLNS_ATTRIBUTE) && elementNamespace == null) {
                    elementNamespace = XMLConstants.XMLNS_ATTRIBUTE_NS_URI;
                }
            }
        }

        List<Parameter> attributeDataParams = resourceConfig.getParameters(ATTRIBUTE_DATA);
        if(attributeDataParams != null && !attributeDataParams.isEmpty()) {
            extractAttributeData(attributeDataParams);
        }
    }

    private void extractAttributeData(List<Parameter> attributeDataParams) {
        for(Parameter attributeDataParam : attributeDataParams) {
            Element attributeElement = attributeDataParam.getXml();
            String name = attributeElement.getAttribute("name");
            String namespace = attributeElement.getAttribute("namespace");
            String value = attributeElement.getAttribute("value");
            int prefixTokIdx = name.indexOf(':');
            QName qName;

            if(prefixTokIdx != -1) {
                String prefix = name.substring(0, prefixTokIdx);

                if(prefix.equals(XMLConstants.XMLNS_ATTRIBUTE) && (namespace.equals(""))) {
                    namespace = XMLConstants.XMLNS_ATTRIBUTE_NS_URI;
                }
                qName = new QName(namespace, name.substring(prefixTokIdx + 1), prefix);
            } else {
                if(name.equals(XMLConstants.XMLNS_ATTRIBUTE) && (namespace.equals(""))) {
                    namespace = XMLConstants.XMLNS_ATTRIBUTE_NS_URI;
                }
                qName = new QName(namespace, name);
            }

            attributes.put(qName, new FreeMarkerTemplate(value));
        }
    }

    @ConfigParam (name = "name", use = ConfigParam.Use.OPTIONAL)
    public SetElementData setElementName(String elementQName) {
        this.elementQName = elementQName;
        return this;
    }

    @ConfigParam (name = "namespace", use = ConfigParam.Use.OPTIONAL)
    public SetElementData setElementNamespace(String namespace) {
        this.elementNamespace = namespace;
        return this;
    }

    public SetElementData setAttribute(QName name, String valueTemplate) {
        attributes.put(name, new FreeMarkerTemplate(valueTemplate));
        return this;
    }

    @Override
    protected void writeStart(SAXElement element) throws IOException {
        SAXElement reconstructedElement = reconstructElement(element);
        super.writeStart(reconstructedElement);
        element.setCache(this, reconstructedElement.getCache(this));
    }

    @Override
    protected void writeEnd(SAXElement element) throws IOException {
        SAXElement reconstructedElement = reconstructElement(element);
        reconstructedElement.setCache(this, element.getCache(this));
        super.writeEnd(reconstructedElement);
    }

    private SAXElement reconstructElement(SAXElement element) {
        QName qName = element.getName();

        if(elementQName != null || elementNamespace != null) {
            // Need to create a new QName for the element...
            String newElementName = (elementName != null? elementName : qName.getLocalPart());
            String newElementNamespace = (elementNamespace != null? elementNamespace : qName.getNamespaceURI());
            String newElementNamespacePrefix = (elementNamespacePrefix != null? elementNamespacePrefix : qName.getPrefix());

            qName = new QName(newElementNamespace, newElementName, newElementNamespacePrefix);
        }

        SAXElement newElement = new SAXElement(qName, element.getAttributes(), element.getParent());
        newElement.setWriter(element.getWriter(this), this);

        if(!attributes.isEmpty()) {
            Map<String, Object> beans = Filter.getCurrentExecutionContext().getBeanContext().getBeanMap();
            Set<Map.Entry<QName, FreeMarkerTemplate>> attributeSet = attributes.entrySet();

            for(Map.Entry<QName, FreeMarkerTemplate> attributeConfig : attributeSet) {
                QName attribName = attributeConfig.getKey();
                FreeMarkerTemplate valueTemplate = attributeConfig.getValue();
                String namespaceURI = attribName.getNamespaceURI();

                if(namespaceURI != null) {
                    String prefix = attribName.getPrefix();
                    if(prefix != null && prefix.length() > 0) {
                        newElement.setAttributeNS(namespaceURI, prefix + ":" + attribName.getLocalPart(), valueTemplate.apply(beans));
                    } else {
                        newElement.setAttributeNS(namespaceURI, attribName.getLocalPart(), valueTemplate.apply(beans));
                    }
                } else {
                    newElement.setAttribute(attribName.getLocalPart(), valueTemplate.apply(beans));
                }
            }
        }

        return newElement;
    }

    public void visitAfter(Element element, ExecutionContext executionContext) throws SmooksException {
        if(elementQName != null || elementNamespace != null) {
            String newElementName = (elementQName != null? elementQName : element.getTagName());
            String newElementNamespace = (elementNamespace != null? elementNamespace : element.getNamespaceURI());

            element = DomUtils.renameElementNS(element, newElementName, newElementNamespace, true, true);
        }

        if(!attributes.isEmpty()) {
            Map<String, Object> beans = Filter.getCurrentExecutionContext().getBeanContext().getBeanMap();
            Set<Map.Entry<QName, FreeMarkerTemplate>> attributeSet = attributes.entrySet();

            for(Map.Entry<QName, FreeMarkerTemplate> attributeConfig : attributeSet) {
                QName attribName = attributeConfig.getKey();
                FreeMarkerTemplate valueTemplate = attributeConfig.getValue();
                String namespaceURI = attribName.getNamespaceURI();

                if(namespaceURI != null) {
                    String prefix = attribName.getPrefix();
                    if(prefix != null && prefix.length() > 0) {
                        element.setAttributeNS(namespaceURI, prefix + ":" + attribName.getLocalPart(), valueTemplate.apply(beans));
                    } else {
                        element.setAttributeNS(namespaceURI, attribName.getLocalPart(), valueTemplate.apply(beans));
                    }
                } else {
                    element.setAttribute(attribName.getLocalPart(), valueTemplate.apply(beans));
                }
            }
        }
    }
}