/*
 * JBoss, Home of Professional Open Source
 * Copyright 2007, Red Hat Middleware LLC, and individual contributors
 * as indicated by the @author tags.
 * See the copyright.txt in the distribution for a full listing
 * of individual contributors.
 * This copyrighted material is made available to anyone wishing to use,
 * modify, copy, or redistribute it subject to the terms and conditions
 * of the GNU General Public License, v. 2.0.
 * This program is distributed in the hope that it will be useful, but WITHOUT A
 * WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
 * PARTICULAR PURPOSE.  See the GNU General Public License for more details.
 * You should have received a copy of the GNU General Public License,
 * v. 2.0 along with this distribution; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA  02110-1301, USA.
 *
 * (C) 2007,
 * @author JBoss Inc.
 */
package com.jboss.ba.demo.client;

import org.apache.log4j.Logger;

import javax.xml.ws.handler.MessageContext;
import javax.xml.ws.handler.soap.SOAPHandler;
import javax.xml.ws.handler.soap.SOAPMessageContext;
import javax.xml.namespace.QName;
import javax.xml.soap.*;

import java.util.Set;
import java.util.HashSet;
import java.util.Iterator;

import com.arjuna.mw.wsc.context.Context;
import com.arjuna.webservices.wscoor.CoordinationConstants;
import com.arjuna.webservices.wscoor.CoordinationContextType;
import com.arjuna.webservices.wsat.AtomicTransactionConstants;
import com.arjuna.webservices.wsba.BusinessActivityConstants;
import com.arjuna.mw.wst.common.CoordinationContextHelper;
import com.arjuna.mw.wst.common.SOAPUtil;
import com.arjuna.mw.wst.*;

/**
 * Created by IntelliJ IDEA.
 * User: jhalli
 * Date: Apr 20, 2007
 * Time: 1:06:28 PM
 * To change this template use File | Settings | File Templates.
 */
public class JaxWSClientHeaderContextProcessor implements SOAPHandler
{
	private static Logger log = Logger.getLogger(JaxWSClientHeaderContextProcessor.class);

	private String handlerName;

	public String getHandlerName()
	{
	   return handlerName;
	}

	public void setHandlerName(String handlerName)
	{
	   this.handlerName = handlerName;
	}

	public boolean handleMessage(MessageContext msgContext)
	{
		Boolean outbound = (Boolean)msgContext.get(MessageContext.MESSAGE_OUTBOUND_PROPERTY);
		if (outbound == null)
		   throw new IllegalStateException("Cannot obtain required property: " + MessageContext.MESSAGE_OUTBOUND_PROPERTY);

		return outbound ? handleOutbound(msgContext) : handleInbound(msgContext);
	}


	public boolean handleFault(MessageContext messagecontext)
	{
		resumeTransaction(messagecontext) ;
		return true;
	}

	public void close(MessageContext messageContext)
	{
	}

	public String toString()
	{
	   return (handlerName != null ? handlerName : super.toString());
	}



   // The header blocks that can be processed by this Handler instance
   private Set<QName> headers = new HashSet<QName>();

   /** Gets the header blocks that can be processed by this Handler instance.
    */
   public Set<QName> getHeaders()
   {
	   log.info("getHeaders");

	   Set<QName> headerSet = new HashSet<QName>();
	   headerSet.add(new QName(CoordinationConstants.WSCOOR_NAMESPACE, CoordinationConstants.WSCOOR_ELEMENT_COORDINATION_CONTEXT));

	  return headerSet;
   }

   /** Sets the header blocks that can be processed by this Handler instance.
    */
   public void setHeaders(Set<QName> headers)
   {
      this.headers = headers;
   }



   protected boolean handleInbound(MessageContext msgContext)
   {
      log.info("handleInbound");
	   resumeTransaction(msgContext) ;


        return true ;
   }


   protected boolean handleOutbound(MessageContext msgContext)
   {
	   log.info("handleOutbound");

	   final SOAPMessageContext soapMessageContext = (SOAPMessageContext)msgContext ;
	   final SOAPMessage soapMessage = soapMessageContext.getMessage() ;

	   if (soapMessage == null)
	   {
		   return true ;
	   }

	   try
	   {
		   /*
			* There should either be an Atomic Transaction *or* a Business Activity
			* associated with the thread.
			*/
		   final TransactionManager transactionManager = TransactionManagerFactory.transactionManager() ;
		   final BusinessActivityManager businessActivityManager = BusinessActivityManagerFactory.businessActivityManager() ;

		   final Context atContext ;
		   if (transactionManager != null)
		   {
			   final com.arjuna.mwlabs.wst.at.context.TxContextImple txContext =
				   (com.arjuna.mwlabs.wst.at.context.TxContextImple)transactionManager.currentTransaction() ;
			   atContext = (txContext == null ? null : txContext.context()) ;
		   }
		   else
		   {
			   atContext = null ;
		   }

		   final Context baContext ;
		   if (businessActivityManager != null)
		   {
			   final com.arjuna.mwlabs.wst.ba.context.TxContextImple txContext =
				   (com.arjuna.mwlabs.wst.ba.context.TxContextImple)businessActivityManager.currentTransaction() ;
			   baContext = (txContext == null ? null : txContext.context()) ;
		   }
		   else
		   {
			   baContext = null ;
		   }

		   final CoordinationContextType coordinationContext ;
		   if (atContext != null)
		   {
			   coordinationContext = atContext.getCoordinationContext() ;
		   }
		   else if (baContext != null)
		   {
			   coordinationContext = baContext.getCoordinationContext() ;
		   }
		   else
		   {
			   coordinationContext = null ;
		   }

		   if (coordinationContext != null)
		   {
			   final SOAPEnvelope env = soapMessage.getSOAPPart().getEnvelope() ;
			   SOAPHeader header = env.getHeader() ;
			   if (header == null)
			   {
				   header = env.addHeader() ;
			   }
			   final Name name = env.createName(CoordinationConstants.WSCOOR_ELEMENT_COORDINATION_CONTEXT, CoordinationConstants.WSCOOR_PREFIX, CoordinationConstants.WSCOOR_NAMESPACE) ;
			   final SOAPHeaderElement headerElement = header.addHeaderElement(name) ;
			   headerElement.addNamespaceDeclaration(CoordinationConstants.WSCOOR_PREFIX, CoordinationConstants.WSCOOR_NAMESPACE) ;
			   headerElement.setMustUnderstand(true) ;
			   CoordinationContextHelper.serialise(env, headerElement, coordinationContext) ;
		   }
	   }
	   catch (final Throwable th)
	   {
		   log.warn("handleOutbound: ", th);
	   }

      return true;
   }

	private void resumeTransaction(final MessageContext messageContext)
	{
		final SOAPMessageContext soapMessageContext = (SOAPMessageContext)messageContext ;
		final SOAPMessage soapMessage = soapMessageContext.getMessage() ;

		if (soapMessage != null)
		{
			try
			{
				final SOAPEnvelope soapEnvelope = soapMessage.getSOAPPart().getEnvelope() ;
				final SOAPHeaderElement soapHeaderElement = getHeaderElement(soapEnvelope, CoordinationConstants.WSCOOR_NAMESPACE, CoordinationConstants.WSCOOR_ELEMENT_COORDINATION_CONTEXT) ;

				if (soapHeaderElement != null)
				{
					final CoordinationContextType cc = CoordinationContextHelper.deserialise(soapEnvelope, soapHeaderElement) ;
					if (cc != null)
					{
						final String coordinationType = cc.getCoordinationType().getValue() ;
						if (AtomicTransactionConstants.WSAT_PROTOCOL.equals(coordinationType))
						{
							final TxContext txContext = new com.arjuna.mwlabs.wst.at.context.TxContextImple(cc) ;
							TransactionManagerFactory.transactionManager().resume(txContext) ;
						}
						else if (BusinessActivityConstants.WSBA_PROTOCOL_ATOMIC_OUTCOME.equals(coordinationType))
						{
							final TxContext txContext = new com.arjuna.mwlabs.wst.ba.context.TxContextImple(cc);
							BusinessActivityManagerFactory.businessActivityManager().resume(txContext) ;
						}
						else
						{
							log.warn("resumeTransaction: unknown Coordtype");
						}
					}
				}
			}
			catch (final Throwable th)
			{
				log.warn("resumeTransaction: ", th);
			}
		}
	}

	private SOAPHeaderElement getHeaderElement(final SOAPEnvelope soapEnvelope, final String uri, final String name)
		throws SOAPException
	{
		final SOAPHeader soapHeader = soapEnvelope.getHeader() ;
		if (soapHeader != null)
		{
			final Iterator headerIter = SOAPUtil.getChildElements(soapHeader) ;
			while(headerIter.hasNext())
			{
				final SOAPHeaderElement current = (SOAPHeaderElement)headerIter.next() ;
				final Name currentName = current.getElementName() ;
				if ((currentName != null) &&
					match(name, currentName.getLocalName()) &&
					match(uri, currentName.getURI()))
				{
					return current ;
				}
			}
		}
		return null ;
	}

	/**
	 * Do the two references match?
	 * @param lhs The first reference.
	 * @param rhs The second reference.
	 * @return true if the references are both null or if they are equal.
	 */
	private boolean match(final Object lhs, final Object rhs)
	{
		if (lhs == null)
		{
			return (rhs == null) ;
		}
		else
		{
			return lhs.equals(rhs) ;
		}
	}
}
