/**
 * Distributable under LGPL license.
 * See terms of license at gnu.org.
 *
 * Copyright (C)
 * <a href="http://www.inab.org">Spanish National Institute of Bioinformatics (INB)</a>
 * <a href="http://www.bsc.es">Barcelona Supercomputing Center (BSC)</a>
 * <a href="http://inb.bsc.es">Computational Node 6</a>
 */

package org.inb.biomoby.client;

import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.Logger;
import org.inb.biomoby.MobyMessageContext;
import javax.xml.bind.JAXBContext;
import javax.xml.bind.JAXBElement;
import javax.xml.bind.JAXBException;
import javax.xml.bind.Marshaller;
import javax.xml.bind.Unmarshaller;
import javax.xml.bind.util.JAXBSource;
import javax.xml.namespace.QName;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.soap.MessageFactory;
import javax.xml.soap.Name;
import javax.xml.soap.SOAPBody;
import javax.xml.soap.SOAPBodyElement;
import javax.xml.soap.SOAPConstants;
import javax.xml.soap.SOAPElement;
import javax.xml.soap.SOAPEnvelope;
import javax.xml.soap.SOAPException;
import javax.xml.soap.SOAPFault;
import javax.xml.soap.SOAPHeader;
import javax.xml.soap.SOAPHeaderElement;
import javax.xml.soap.SOAPMessage;
import javax.xml.soap.SOAPPart;
import javax.xml.stream.XMLStreamException;
import javax.xml.transform.Source;
import javax.xml.transform.stream.StreamSource;
import javax.xml.ws.BindingProvider;
import javax.xml.ws.Dispatch;
import javax.xml.ws.Service;
import javax.xml.ws.soap.AddressingFeature;
import javax.xml.ws.soap.MTOMFeature;
import javax.xml.ws.soap.SOAPBinding;
import javax.xml.ws.wsaddressing.W3CEndpointReference;
import org.inb.biomoby.shared.registry.Service.CATEGORY;
import org.inb.biomoby.shared.message.MobyDataElement;
import org.inb.biomoby.shared.message.MobyMessage;
import org.inb.biomoby.shared.wsrf.MobyEndpointReference;
import org.inb.biomoby.shared.wsrf.PropertyResponseWrapper;
import org.inb.biomoby.shared.wsrf.SubmitResponseWrapper;
import org.inb.lsae.AnalysisEvent;
import org.inb.wsrf.rp2.GetMultipleResourceProperties;
import org.inb.wsrf.rp2.GetMultipleResourcePropertiesResponse;
import org.inb.wsrf.rp2.GetResourcePropertyResponse;
import org.inb.wsrf.rp2.QueryExpressionType;
import org.inb.wsrf.rp2.QueryResourceProperties;
import org.inb.wsrf.rp2.QueryResourcePropertiesResponse;
import org.w3c.dom.Element;
import org.w3c.dom.Node;

/**
 * @author Dmitry Repchevsky
 */

public class MobyDispatch
{
    public final static String NAMESPACE = "http://biomoby.org/";
    
    private Service service;
    private org.inb.biomoby.shared.registry.Service mobyService;
    
    private final String user;
    private final String password;
    
    public MobyDispatch(org.inb.biomoby.shared.registry.Service mobyService)
    {
        this(mobyService, null, null);
    }
    
    public MobyDispatch(org.inb.biomoby.shared.registry.Service mobyService, String user, String password)
    {
        this.mobyService = mobyService;
        
        this.user = user;
        this.password = password;
        
        QName sName = new QName(NAMESPACE, mobyService.getName());
        
        this.service = Service.create(sName);
        this.service.addPort(sName, SOAPBinding.SOAP11HTTP_MTOM_BINDING, mobyService.getUrl());
        
    }

    /**
     * Method to synchronously invoke BioMoby service (provided in constructor)
     *
     * @param message - message to be sent.
     * @return message sent by the service.
     * @throws Exception
     */
    public MobyMessage invoke(MobyMessage message) throws Exception
    {
        CATEGORY category = mobyService.getCategory();

        if (category == CATEGORY.doc_literal ||
            category == CATEGORY.doc_literal_async)
        {
            return invokeDocument(message);
        }
        else if (category == CATEGORY.moby ||
                 category == CATEGORY.moby_async)
        {
            return invokeRPC(message);
        }

        throw new Exception("Unknown BioMoby service category (" + category.name() + ")");
    }
    
    public MobyWSRFResource invokeAsync(MobyMessage message) throws Exception
    {
        CATEGORY category = mobyService.getCategory();

        if (category == CATEGORY.doc_literal_async)
        {
            return invokeDocumentAsync(message);
        }
        else if (category == CATEGORY.moby_async)
        {
            return invokeRPCAsync(message);
        }

        throw new Exception("Unknown (or not valid) BioMoby service category (" + category.name() + ")");
    }
    
    private MobyMessage invokeDocument(MobyMessage message) throws JAXBException
    {
        QName pName = new QName(NAMESPACE, mobyService.getName());
        
        Dispatch dispatch = service.createDispatch(pName,
                                                   MobyMessageContext.getContext(),
                                                   Service.Mode.PAYLOAD, new MTOMFeature());

        SOAPBinding binding = (SOAPBinding)dispatch.getBinding();
        binding.setMTOMEnabled(true);

        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, Boolean.TRUE); //use SOAPAction!
        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, service.getServiceName().getNamespaceURI() + '#' + mobyService.getName());
        
        if (user != null)
        {
            //dispatch.getRequestContext().put(Context.SECURITY_AUTHENTICATION, "none");

//            dispatch.getRequestContext().put(BindingProvider.USERNAME_PROPERTY, "pepino");
//            dispatch.getRequestContext().put(BindingProvider.PASSWORD_PROPERTY, "password");
        }

        MobyMessage response = (MobyMessage)dispatch.invoke(message);
        
        return response;
    }

    private MobyWSRFResource invokeDocumentAsync(MobyMessage message) throws JAXBException, ParserConfigurationException
    {
        QName pName = new QName(NAMESPACE, mobyService.getName());
        
        Dispatch<Source> dispatch = service.createDispatch(pName,
                                                           Source.class,
                                                           Service.Mode.PAYLOAD);

        SOAPBinding binding = (SOAPBinding)dispatch.getBinding();
        binding.setMTOMEnabled(true);

        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, Boolean.TRUE); //use SOAPAction!
        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, service.getServiceName().getNamespaceURI() + '#' + mobyService.getName() + "_submit");

        JAXBContext ctx = MobyMessageContext.getContext();
        
        JAXBSource source = new JAXBSource(ctx, message);
        Source response = dispatch.invoke(source);
                
        Unmarshaller u = ctx.createUnmarshaller();

        JAXBElement<SubmitResponseWrapper> element = u.unmarshal(response, SubmitResponseWrapper.class);
        
        SubmitResponseWrapper wrapper = element.getValue();
        W3CEndpointReference ref = wrapper.getEndpointReference();
        
        return new MobyWSRFResource(new MobyEndpointReference(ref));
    }
    
    private MobyMessage invokeRPC(MobyMessage message) throws Exception
    {
        QName sName = service.getServiceName();
        
        SOAPMessage request = marshal(sName.getLocalPart(), message);
     
        Dispatch<SOAPMessage> dispatch = service.createDispatch(sName, SOAPMessage.class, Service.Mode.MESSAGE);

        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, Boolean.TRUE); //use SOAPAction!
        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, service.getServiceName().getNamespaceURI() + '#' + mobyService.getName());
        
        SOAPMessage response = dispatch.invoke(request);
        
        return MobyDispatch.unmarshal(response);
    }

    private MobyWSRFResource invokeRPCAsync(MobyMessage message) throws Exception
    {
        QName sName = service.getServiceName();
        
        SOAPMessage request = marshal(sName.getLocalPart(), message, true);
     
        Dispatch<SOAPMessage> dispatch = service.createDispatch(sName, SOAPMessage.class, Service.Mode.MESSAGE);

        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, Boolean.TRUE); //use SOAPAction!
        dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, service.getServiceName().getNamespaceURI() + '#' + mobyService.getName() + "_submit");
        
        SOAPMessage responseMessage = dispatch.invoke(request);

        JAXBContext ctx = MobyMessageContext.getContext();
        Unmarshaller u = ctx.createUnmarshaller();
        
        JAXBElement<SubmitResponseWrapper> element = u.unmarshal(responseMessage.getSOAPBody().getFirstChild(), SubmitResponseWrapper.class);
        
        SubmitResponseWrapper wrapper = element.getValue();
        W3CEndpointReference ref = wrapper.getEndpointReference();
        
        return new MobyWSRFResource(new MobyEndpointReference(ref));
    }

    public static SOAPMessage marshal(String serviceName, Object object) throws SOAPException, IOException, JAXBException, XMLStreamException
    {
        return marshal(serviceName, object, false);
    }
    
    public static SOAPMessage marshal(String serviceName, Object object, boolean isAsync) throws SOAPException, IOException, JAXBException, XMLStreamException
    {
        MessageFactory mf = MessageFactory.newInstance(SOAPConstants.SOAP_1_1_PROTOCOL);
        SOAPMessage msg = mf.createMessage();
        
        SOAPPart sp = msg.getSOAPPart();

        SOAPEnvelope env = sp.getEnvelope();
        env.addNamespaceDeclaration("xsi","http://www.w3.org/2001/XMLSchema-instance");
        env.addNamespaceDeclaration("xsd", "http://www.w3.org/2001/XMLSchema");

        SOAPBody bd = env.getBody();

        Name sName;
        if (isAsync)
        {
            sName = env.createName(serviceName + "_submit", "nsl", NAMESPACE);
        }
        else
        {
            sName = env.createName(serviceName, "nsl", NAMESPACE);
        }
        
        SOAPElement sElement = bd.addChildElement(sName);
        sElement.setEncodingStyle("http://schemas.xmlsoap.org/soap/encoding/");

        SOAPElement aElement = sElement.addChildElement("arg0", "nsl");
        
        aElement.addNamespaceDeclaration("soapenc", "http://schemas.xmlsoap.org/soap/encoding/");
        aElement.setAttribute("xsi:type", "soapenc:string");

        if (object != null)
        {
            String moby = MobyMessageContext.marshall(object);
            aElement.addTextNode(moby);

            Logger.getLogger(MobyDispatch.class.getCanonicalName()).finest("marshalled message:\n" + moby);
        }
        
	msg.saveChanges();

        return msg;
    }

    public static MobyMessage unmarshal(SOAPMessage message) throws JAXBException, IOException, XMLStreamException, SOAPException
    {
        SOAPBody body = message.getSOAPBody();

        Node method = body.getFirstChild(); // check mobyService.getName() + "Response"?
        Node argument = method.getFirstChild();

        Unmarshaller u = MobyMessageContext.getContext().createUnmarshaller();

        JAXBElement jel = u.unmarshal(argument, byte[].class);

        Object escaped = jel.getValue();

        String moby = escaped instanceof String ? (String)escaped : new String((byte[])escaped);

        Logger.getLogger(MobyDispatch.class.getCanonicalName()).finest("unmarshalled message:\n" + moby);

        return (MobyMessage)u.unmarshal(new StreamSource(new StringReader(moby)));
    }

    public static class MobyWSRFResource
    {
        private final static String QUERY_EXPRESSION_DIALECT = "http://www.w3.org/TR/1999/REC-xpath-19991116";

        private final Service service;
        public final MobyEndpointReference ref;
        
        public MobyWSRFResource(MobyEndpointReference ref)
        {
            QName sName = new QName(NAMESPACE, "anonimous");
            
            service = Service.create(sName);
            service.addPort(sName, SOAPBinding.SOAP11HTTP_BINDING, ref.url);

            
            this.ref = ref;
        }

        public void destroy() throws Exception
        {
            Dispatch<SOAPMessage> dispatch = service.createDispatch(service.getServiceName(),
                                                                    SOAPMessage.class,
                                                                    Service.Mode.MESSAGE, new AddressingFeature());
            
            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, true);
            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, "http://docs.oasis-open.org/wsrf/rlw-2/ImmediateResourceTermination/DestroyRequest");

            MessageFactory mf = MessageFactory.newInstance(SOAPConstants.SOAP_1_1_PROTOCOL);
            SOAPMessage requestMessage = mf.createMessage();

            SOAPHeader header = requestMessage.getSOAPHeader();

            SOAPHeaderElement serviceInvocationId = header.addHeaderElement(new QName(NAMESPACE, "ServiceInvocationId"));
            serviceInvocationId.setTextContent(ref.invocationId);

            SOAPMessage responseMessage = dispatch.invoke(requestMessage);

            SOAPBody body = responseMessage.getSOAPBody();
            
            if (body.hasFault())
            {
                SOAPFault fault = body.getFault();
                throw new Exception(fault.getFaultString());
            }
        }

        public AnalysisEvent getResourcePropertyStatus(String queryId) throws Exception
        {
            return getResourceProperty("status", queryId);
        }

        public MobyMessage getResourcePropertyResult(String queryId) throws Exception
        {
            return getResourceProperty("result", queryId);
        }

        private <T> T getResourceProperty(String property, String queryId) throws Exception
        {
            Dispatch<SOAPMessage> dispatch = service.createDispatch(service.getServiceName(),
                                                                    SOAPMessage.class,
                                                                    Service.Mode.MESSAGE, new AddressingFeature());
            
            SOAPBinding binding = (SOAPBinding)dispatch.getBinding();
            binding.setMTOMEnabled(true);
            
            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, true);
            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, "http://docs.oasis-open.org/wsrf/rpw-2/GetResourceProperty/GetResourcePropertyRequest");

            MessageFactory mf = MessageFactory.newInstance(SOAPConstants.SOAP_1_1_PROTOCOL);
            SOAPMessage requestMessage = mf.createMessage();
            requestMessage.getSOAPPart().getEnvelope().addNamespaceDeclaration("mobyws", NAMESPACE);

            SOAPHeader header = requestMessage.getSOAPHeader();

            SOAPHeaderElement serviceInvocationId = header.addHeaderElement(new QName(NAMESPACE, "ServiceInvocationId"));
            serviceInvocationId.setTextContent(ref.invocationId);

            SOAPBody request = requestMessage.getSOAPBody();

            SOAPBodyElement getResourceProperty = request.addBodyElement(new QName("http://docs.oasis-open.org/wsrf/rp-2", "GetResourceProperty"));
            
            getResourceProperty.setTextContent("mobyws:" + property + '_' + queryId);

            SOAPMessage responseMessage = dispatch.invoke(requestMessage);

            SOAPBody body = responseMessage.getSOAPBody();
            
            if (body.hasFault())
            {
                SOAPFault fault = body.getFault();
                throw new Exception(fault.getFaultString());
            }
            
            JAXBContext ctx = MobyMessageContext.getContext();
            Unmarshaller u = ctx.createUnmarshaller();
            
            GetResourcePropertyResponse response = (GetResourcePropertyResponse)u.unmarshal(body.extractContentAsDocument());
            
            return (T)getResponseObject(response);
        }
        
        public List<AnalysisEvent> getMultipleResourcePropertiesStatus(List<String> queryIDs) throws Exception
        {
            return getMultipleResourceProperties("status", queryIDs);
        }

        public List<MobyMessage> getMultipleResourcePropertiesResult(List<String> queryIDs) throws Exception
        {
            return getMultipleResourceProperties("result", queryIDs);
        }

        private <T> List<T> getMultipleResourceProperties(String property, List<String> queryIDs) throws Exception
        {
            Dispatch<SOAPMessage> dispatch = service.createDispatch(service.getServiceName(),
                                                                    SOAPMessage.class,
                                                                    Service.Mode.MESSAGE, new AddressingFeature());
            
            SOAPBinding binding = (SOAPBinding)dispatch.getBinding();
            binding.setMTOMEnabled(true);

            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, true);
            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, "http://docs.oasis-open.org/wsrf/rpw-2/GetMultipleResourceProperties/GetMultipleResourcePropertiesRequest");

            MessageFactory mf = MessageFactory.newInstance(SOAPConstants.SOAP_1_1_PROTOCOL);
            SOAPMessage requestMessage = mf.createMessage();
            requestMessage.getSOAPPart().getEnvelope().addNamespaceDeclaration("mobyws", NAMESPACE);
            
            SOAPHeader header = requestMessage.getSOAPHeader();
            SOAPHeaderElement serviceInvocationId = header.addHeaderElement(new QName(NAMESPACE, "ServiceInvocationId"));
            serviceInvocationId.setTextContent(ref.invocationId);

            SOAPBody request = requestMessage.getSOAPBody();

            GetMultipleResourceProperties getMultipleResourceProperties = new GetMultipleResourceProperties();
            
            for (String queryId : queryIDs)
            {
                getMultipleResourceProperties.addProperty(new QName(NAMESPACE, property + '_' + queryId, "mobyws"));
            }
            
            Marshaller m = MobyMessageContext.getContext().createMarshaller();
            m.marshal(getMultipleResourceProperties, request);

            SOAPMessage responseMessage = dispatch.invoke(requestMessage);

            SOAPBody body = responseMessage.getSOAPBody();
            
            if (body.hasFault())
            {
                SOAPFault fault = body.getFault();
                throw new Exception(fault.getFaultString());
            }
            
            JAXBContext ctx = MobyMessageContext.getContext();
            Unmarshaller u = ctx.createUnmarshaller();
            
            GetMultipleResourcePropertiesResponse response = (GetMultipleResourcePropertiesResponse)u.unmarshal(body.getFirstChild());
            
            List<T> objects = new ArrayList<T>();
            
            List responses = response.getResponses();
            
            if (responses.isEmpty())
            {
                return objects;
            }
            
            for(Object object : responses)
            {
                if (object instanceof GetResourcePropertyResponse)
                {
                    GetResourcePropertyResponse getResourcePropertyResponse = (GetResourcePropertyResponse)object;
                    
                    T t = (T)getResponseObject(getResourcePropertyResponse);
                    
                    if (t != null)
                    {
                        objects.add(t);
                    }
                }
            }

            return objects;           
        }

        public List<MobyDataElement> queryResourcePropertyResult(String expression) throws Exception
        {
            Dispatch<SOAPMessage> dispatch = service.createDispatch(service.getServiceName(),
                                                                    SOAPMessage.class,
                                                                    Service.Mode.MESSAGE, new AddressingFeature());
            
            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_USE_PROPERTY, true);
            dispatch.getRequestContext().put(BindingProvider.SOAPACTION_URI_PROPERTY, "http://docs.oasis-open.org/wsrf/rpw-2/QueryResourceProperties/QueryResourcePropertiesRequest");

            MessageFactory mf = MessageFactory.newInstance(SOAPConstants.SOAP_1_1_PROTOCOL);
            SOAPMessage requestMessage = mf.createMessage();
            requestMessage.getSOAPPart().getEnvelope().addNamespaceDeclaration("mobyws", NAMESPACE);

            SOAPHeader header = requestMessage.getSOAPHeader();

            SOAPHeaderElement serviceInvocationId = header.addHeaderElement(new QName(NAMESPACE, "ServiceInvocationId"));
            serviceInvocationId.setTextContent(ref.invocationId);

            SOAPBody request = requestMessage.getSOAPBody();
            
            QueryExpressionType eType = new QueryExpressionType();
            eType.setDialect(QUERY_EXPRESSION_DIALECT);
            eType.addContent(expression);
            
            QueryResourceProperties queryResourceProperties = new QueryResourceProperties();
            queryResourceProperties.setQueryExpression(eType);
            
            Marshaller m = MobyMessageContext.getContext().createMarshaller();
            m.marshal(queryResourceProperties, request);

            SOAPMessage responseMessage = dispatch.invoke(requestMessage);

            SOAPBody body = responseMessage.getSOAPBody();
            
            if (body.hasFault())
            {
                SOAPFault fault = body.getFault();
                throw new Exception(fault.getFaultString());
            }
            
            JAXBContext ctx = MobyMessageContext.getContext();
            Unmarshaller u = ctx.createUnmarshaller();
            
            QueryResourcePropertiesResponse response = (QueryResourcePropertiesResponse)u.unmarshal(body.extractContentAsDocument());

            List<MobyDataElement> elements = new ArrayList<MobyDataElement>();
            
            List list = response.getContent();
            
            if (!list.isEmpty())
            {
                for(Object object : list)
                {
                    if (object instanceof MobyDataElement)
                    {
                        elements.add((MobyDataElement) object);
                    }
                }
            }
            
            return elements;
            
        }

        private Object getResponseObject(GetResourcePropertyResponse response) throws Exception
        {
            List properties = response.getProperties();
            
            if (properties.isEmpty())
            {
                return null; // or generate en exception???
            }

            // According WSRF a property may be any valid xml...
            // To do not break this contract parse it manually...
            Object object = properties.get(0);
            
            if (object instanceof Element)
            {
                Element element = (Element)object;
                Unmarshaller u = MobyMessageContext.getContext().createUnmarshaller();
            
                JAXBElement<PropertyResponseWrapper> wrapper = u.unmarshal(element, PropertyResponseWrapper.class);
                
                return wrapper.getValue().getResponseObject();
            }
            
            throw new Exception("Unknown element: " + object.getClass().getSimpleName());
        }
    }
}