package com.ibm.ws.webservices.axis.ws.providers.java;

import java.io.ByteArrayInputStream;
import java.io.ObjectInputStream;
import java.util.Iterator;

import javax.ejb.EJBObject;
import javax.ejb.Handle;
import javax.xml.soap.SOAPEnvelope;
import javax.xml.soap.SOAPHeader;
import javax.xml.soap.SOAPHeaderElement;
import javax.xml.soap.SOAPMessage;
import javax.xml.soap.SOAPPart;

import com.ibm.ws.webservices.engine.MessageContext;
import com.ibm.ws.webservices.engine.encoding.Base64;

public class StatefulEjbProvider extends WASEjbProvider
{
    protected Object makeNewServiceObject(MessageContext messagecontext, String s)
        throws Exception
    {
		String id = getId(messagecontext);
		if(id != null){
			/* return stateful service */
			return this.deSerialize(id);
		}
		else
/* return stateless service */
			return super.makeNewServiceObject(messagecontext,s);
    }
    /* Your Factory service should have the matching serialize method available*/
	private EJBObject deSerialize(String id) {
	  EJBObject ejbObject= null;
	  try {
		byte[] idBytes = Base64.decode(id);
		ByteArrayInputStream byteStream = new ByteArrayInputStream(idBytes);
		ObjectInputStream objStream = new ObjectInputStream(byteStream);
		Handle ejbHandle = (Handle)objStream.readObject();
		objStream.close();
		ejbObject = ejbHandle.getEJBObject();
	  } catch (Exception e) {e.printStackTrace();}
	  return ejbObject;
	}

	private String getId(MessageContext messagecontext) throws Exception{
		String id = null;
		SOAPMessage message = messagecontext.getMessage();
		SOAPPart sp = message.getSOAPPart();
		SOAPEnvelope se = sp.getEnvelope();
		SOAPHeader header =  se.getHeader();
		Iterator it = header.getChildElements();
		while (it.hasNext()){
			SOAPHeaderElement he = (SOAPHeaderElement) it.next();
			if (he.getElementName().getLocalName().equals("id"))
				id =  he.getValue();
		}
		return id;
	}
}

public class TestClient {

	public static void main(String[] args) {
		try {
  		/* Locate factory service */
		FactoryServiceLocator factoryLocator = new FactoryServiceLocator();
		Factory factoryService= factoryLocator.getFactory();
		/* Create instance of stateful service and get a serviceId for it */
		String statefulService1Handle = factoryService.createStatefulService1("param");
		String statefulService2Handle = factoryService.createStatefulService2("param",2);

  		/* Locate stateful service */
		StatefulService1ServiceLocator statefulService1ServiceLocator =
		new StatefulService1ServiceLocator();

		StatefulService1SoapBindingStub statefulService1=
		(StatefulService1SoapBindingStub) 
			statefulService1ServiceLocator.getStatefulService1();

        /* Set serviceId to locate the specific instance of stateful service*/
		statefulService1.setServiceId(statefulService1Handle);

		/* Call a method on the stateful service */
		System.out.println("StatefulService1. state : " + statefulService1.getStateS1());

		/* Destroy the instance of stateful service*/
		factoryService.removeStatefulService(statefulService1Handle);

		StatefulService2ServiceLocator statefulService2ServiceLocator =
 		new StatefulService2ServiceLocator();

		StatefulService2SoapBindingStub statefulService2=
		(StatefulService2SoapBindingStub) 
			statefulService2ServiceLocator.getStatefulService2();

		/* Set serviceId to locate the specific instance of stateful service*/
		statefulService2.setServiceId(statefulService2Handle);

		System.out.println("StatefulService2. state : " + statefulService2.getStateS2());

		factoryService.removeStatefulService(statefulService2Handle);

	} catch (Exception e) {
		e.printStackTrace();
	}
}

public class StatefulService1SoapBindingStub extends org.apache.axis.client.Stub
		implements com.ibm.StatefulService1 {
    private String serviceId;
    private org.apache.axis.client.Call createCall() throws java.rmi.RemoteException {
        try {
            org.apache.axis.client.Call _call =
                    (org.apache.axis.client.Call) super.service.createCall();
			     ...
			/* set serviceId header entry */
	if (this.serviceId != null){
		SOAPHeaderElement element = new SOAPHeaderElement("http://ibm.com", "id");
		try {
			element.setObjectValue(this.serviceId);
		} catch (SOAPException e) {...}
	}
	_call.addHeader(element);
	}
            return _call;
        }
        catch (java.lang.Throwable t) {...}
    }

    public void setStateS1(java.lang.String newState) throws java.rmi.RemoteException {...}

    public java.lang.String getStateS1() throws java.rmi.RemoteException {...}

     public void setServiceId(String string) {
serviceId = string;
     }
}