/*-
 * ~~~~~~licensing~~~~~~
 * system-commons
 * ==========
 * Copyright (C) 2020 - 2021 EmDev LLC
 * ==========
 * Licensed 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.
 * ~~~~~~/licensing~~~~~~
 */
package ru.entaxy.esb.system.common.validator;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.cxf.binding.soap.SoapMessage;
import org.apache.cxf.binding.soap.saaj.SAAJInInterceptor;
import org.apache.cxf.interceptor.Fault;
import org.apache.cxf.phase.AbstractPhaseInterceptor;
import org.apache.cxf.phase.Phase;
import org.apache.cxf.service.Service;
import org.apache.cxf.service.model.ServiceModelUtil;
import org.apache.cxf.ws.addressing.EndpointReferenceUtils;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.xml.sax.SAXException;

import javax.xml.soap.SOAPMessage;
import javax.xml.stream.XMLStreamException;
import javax.xml.transform.dom.DOMSource;
import javax.xml.validation.Schema;
import javax.xml.validation.Validator;
import javax.xml.xpath.XPathExpressionException;
import java.io.IOException;
import java.util.List;

public class ValidateInterceptor extends AbstractPhaseInterceptor<SoapMessage> {

    private static final Log LOG = LogFactory.getLog(ValidateInterceptor.class);

    private final SAAJInInterceptor saajIn;
    private final XmlParser xmlParser;
    private boolean schemaValidationEnabled;

    public ValidateInterceptor() {
        super(Phase.PRE_PROTOCOL);
        saajIn = new SAAJInInterceptor();
        xmlParser = new XmlParser();

        getAfter().add(SAAJInInterceptor.class.getName());
    }

    @Override
    public void handleMessage(SoapMessage message) throws Fault {
        try {
            Node body = getMessageBody((DOMSource) getSOAPMessage(message).getSOAPPart().getContent());
            if (body != null)
                validate(body, message);
            else
                throw new XMLStreamException("Can't find the tag \"Body\"");
        } catch (RuntimeException re) {
            throw re;
        } catch (Exception e) {
            throw new Fault(e);
        }
    }

    private Node getMessageBody(DOMSource source) throws XPathExpressionException {
        Node node = source.getNode().cloneNode(true);
        List<Node> nodeList = xmlParser.getNodes(node.getLastChild(), "Body");
        return !nodeList.isEmpty() ? nodeList.get(0) : null;
    }

    private void validate(Node node, SoapMessage soapMessage) throws IOException, SAXException, XPathExpressionException {
        Validator validator = getValidator(soapMessage);
        validator.validate(new DOMSource(getNodeForValidate(node)));
    }

    private Node getNodeForValidate(Node node) throws XPathExpressionException {
        if (schemaValidationEnabled && node.getLocalName().contains("packets")) {
            Element element = (Element) node;
            for (Node content : xmlParser.getNodes(element, "content")) {
                content.getParentNode().removeChild(content);
            }
            return element;
        }
        return node;
    }

    private Validator getValidator(SoapMessage soapMessage) {
        Service service = ServiceModelUtil.getService(soapMessage.getExchange());
        Schema schema = EndpointReferenceUtils.getSchema(service.getServiceInfos().get(0), soapMessage.getExchange().getBus());
        return schema.newValidator();
    }

    private SOAPMessage getSOAPMessage(SoapMessage smsg) {
        SOAPMessage soapMessage = smsg.getContent(SOAPMessage.class);
        if (soapMessage == null) {
            saajIn.handleMessage(smsg);
            soapMessage = smsg.getContent(SOAPMessage.class);
        }
        return soapMessage;
    }

    public void setSchemaValidationEnabled(boolean schemaValidationEnabled) {
        this.schemaValidationEnabled = schemaValidationEnabled;
    }
}
