/*
 * JBoss, the OpenSource J2EE webOS
 *
 * Distributable under LGPL license.
 * See terms of license at gnu.org.
 *
 */
package org.jboss.resource.adapter.jdbc.remote;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.UndeclaredThrowableException;
import java.sql.CallableStatement;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ParameterMetaData;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.Map;

import javax.management.MBeanServer;
import javax.management.ObjectName;
import javax.naming.BinaryRefAddr;
import javax.naming.InitialContext;
import javax.naming.NamingException;
import javax.naming.Reference;
import javax.naming.StringRefAddr;
import javax.resource.Referenceable;
import javax.sql.DataSource;

import org.jboss.deployment.DeploymentException;
import org.jboss.invocation.Invocation;
import org.jboss.invocation.Invoker;
import org.jboss.invocation.InvokerInterceptor;
import org.jboss.invocation.MarshalledInvocation;
import org.jboss.logging.Logger;
import org.jboss.naming.NonSerializableFactory;
import org.jboss.naming.Util;
import org.jboss.proxy.ClientMethodInterceptor;
import org.jboss.proxy.GenericProxyFactory;
import org.jboss.resource.connectionmanager.ConnectionFactoryBindingService;
import org.jboss.system.Registry;
import org.jboss.system.ServiceMBeanSupport;


/** An mbean service that pvovides the detached invoker ops for the
 * javax.sql.DataSource and related java.sql.* interfaces.
 *
 *
 * @jmx:mbean name="jboss.jca:service=ConnectionFactoryBinding"
 *            extends="org.jboss.resource.connectionmanager.ConnectionFactoryBindingServiceMBean"
 *
 * @author Scott.Stark@jboss.org
 * @author Tom.Elrod@jboss.org
 * @author adrian@jboss.com
 * @version $Revision: 1.7 $
 */
public class WrapperDataSourceService extends ConnectionFactoryBindingService
   implements WrapperDataSourceServiceMBean
{
   // Constants -----------------------------------------------------

   private static Logger log = Logger.getLogger(WrapperDataSourceService.class);
   
   // Attributes ----------------------------------------------------

   private ObjectName jmxInvokerName;
   private Invoker delegateInvoker;
   private Object theProxy;
   private HashMap marshalledInvocationMapping = new HashMap();
   private HashMap connectionMap = new HashMap();
   private HashMap statementConnectionRelation = new HashMap();
   private HashMap statementMap = new HashMap();
   private HashMap resultSetMap = new HashMap();
   private HashMap lobMap = new HashMap();
   private HashMap databaseMetaDataMap = new HashMap();
   private boolean trace = log.isTraceEnabled();

   // Static --------------------------------------------------------

   // Constructors --------------------------------------------------

   // Public --------------------------------------------------------
   
   // ServiceMBeanSupport overrides ---------------------------------

   protected void startService() throws Exception
   {
      determineBindName();
      createConnectionFactory();
      createProxy();
      calculateMethodHases();
      bindConnectionFactory();
   }

   protected void stopService() throws Exception
   {
      unbindConnectionFactory();
      destroyProxy();
   }
   
   // ConnectionFactoryBindingService overrides ---------------------
   
   protected void bindConnectionFactory() throws Exception
   {
      InitialContext ctx = new InitialContext();
      try
      {
         log.debug("Binding object '" + cf + "' into JNDI at '" + bindName + "'");
         // Associated the local cf with the NonSerializable factory
         NonSerializableFactory.rebind(bindName, cf);
         /* Create a reference that uses the the DataSourceFactory as the
         reference factory class. This class detects whether the lookup
         is being done locally or remotely and returns either the just bound
         connection factory, or a DataSource proxy that uses the detached
         invoker framework to expose remote proxies to the server side
         DataSource and related elements.
         */
         Referenceable referenceable = (Referenceable) cf;
         // Set the DataSource proxy as the ProxyData ref address
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         ObjectOutputStream oos = new ObjectOutputStream(baos);
         oos.writeObject(theProxy);
         oos.close();
         byte[] proxyBytes = baos.toByteArray();
         BinaryRefAddr dsAddr = new BinaryRefAddr("ProxyData", proxyBytes);
         String factory = DataSourceFactory.class.getName();
         Reference dsRef = new Reference("javax.sql.DataSource", dsAddr, factory, null);
         referenceable.setReference(dsRef);
         // Set the VMID as the address local/remote indicator
         baos.reset();
         ObjectOutputStream oos2 = new ObjectOutputStream(baos);
         oos2.writeObject(DataSourceFactory.vmID);
         oos2.close();
         byte[] id = baos.toByteArray();
         BinaryRefAddr localAddr = new BinaryRefAddr("VMID", id);
         dsRef.add(localAddr);
         /* Bind the Referenceable connection factory into JNDI and set the
         JndiName value of the reference address for use by the DataSourceFactory
         when looking up the local factory from the NonSerializableFactory.
         */
         StringRefAddr jndiRef = new StringRefAddr("JndiName", bindName);
         dsRef.add(jndiRef);
         Util.rebind(ctx, bindName, cf);
         log.info("Bound connection factory for resource adapter for ConnectionManager '" + serviceName + " to JNDI name '" + bindName + "'");
      }
      catch (NamingException ne)
      {
         throw new DeploymentException("Could not bind ConnectionFactory into jndi: " + bindName);
      }
      finally
      {
         ctx.close();
      }
   }
   
   // WrapperDataSourceServiceMBean implementation ------------------

   /**
    * Get the transport
    * 
    * @jmx:managed-attribute
    * @return the transport
    */
   public ObjectName getJMXInvokerName()
   {
      return jmxInvokerName;
   }

   /**
    * Set the transport
    * 
    * @jmx:managed-attribute
    * @param jndiName the jndi name
    */
   public void setJMXInvokerName(ObjectName jmxInvokerName)
   {
      this.jmxInvokerName = jmxInvokerName;
   }

   /** Expose the DataSource via JMX to invokers.
    *
    * @jmx:managed-operation
    *
    * @param invocation    A pointer to the invocation object
    * @return              Return value of method invocation.
    *
    * @throws Exception    Failed to invoke method.
    */
   public Object invoke(Invocation invocation) throws Exception
   {
      // Set the method hash to Method mapping
      if (invocation instanceof MarshalledInvocation)
      {
         MarshalledInvocation mi = (MarshalledInvocation) invocation;
         mi.setMethodMap(marshalledInvocationMapping);
      }
      // Invoke the Naming method via reflection
      Method method = invocation.getMethod();
      Class methodClass = method.getDeclaringClass();
      Object[] args = invocation.getArguments();
      Object value = null;

      try
      {
         if( methodClass.isAssignableFrom(DataSource.class) )
         {
            InitialContext ctx = new InitialContext();
            DataSource ds = (DataSource) ctx.lookup(bindName);
            value = doDataSourceMethod(ds, method, args);
         }
         else if( methodClass.isAssignableFrom(Connection.class) )
         {
            Integer id = (Integer) invocation.getId();
            Connection conn = (Connection) connectionMap.get(id);
            if( conn == null )
            {
               throw new IllegalAccessException("Failed to find connection: "+id);
            }
            value = doConnectionMethod(conn, method, args);
         }
         else if( methodClass.isAssignableFrom(Statement.class) ||
            methodClass.isAssignableFrom(PreparedStatement.class) ||
            methodClass.isAssignableFrom(CallableStatement.class))
         {
            Integer id = (Integer) invocation.getId();
            Statement stmt = (Statement) statementMap.get(id);
            if( stmt == null )
            {
               throw new SQLException("Failed to find Statement: " + id);
            }
            value = doStatementMethod(stmt, method, args);
         }
         else if( methodClass.isAssignableFrom(ResultSet.class) )
         {
            Integer id = (Integer) invocation.getId();
            ResultSet results = (ResultSet) resultSetMap.get(id);
            if( results == null )
            {
               throw new IllegalAccessException("Failed to find ResultSet: "+id);
            }
            value = doResultSetMethod(results, method, args);
         }
         else if (methodClass.isAssignableFrom(DatabaseMetaData.class))
         {
            Integer id = (Integer) invocation.getId();
            DatabaseMetaData dbMetaData = (DatabaseMetaData) databaseMetaDataMap.get(id);
            if(dbMetaData == null)
            {
               throw new IllegalAccessException("Failed to find DatabaseMetaData: " + id);
            }
            value = doDatabaseMetaDataMethod(dbMetaData, method, args);
         }
         else if( methodClass.isAssignableFrom(java.sql.Blob.class) || methodClass.isAssignableFrom(java.sql.Clob.class))
         {
            Integer id = (Integer) invocation.getId();
            Object results = lobMap.get(id);
            if( results == null )
            {
               throw new IllegalAccessException("Failed to find LOB: "+id);
            }
            value = doLobMethod(results, method, args);
         }
         else
         {
            throw new UnsupportedOperationException("Do not know how to handle method="+method);
         }
      }
      catch (InvocationTargetException e)
      {
         Throwable t = e.getTargetException();
         if (t instanceof Exception)
            throw (Exception) t;
         else
            throw new UndeclaredThrowableException(t, method.toString());
      }

      return value;
   }

   // Package protected ---------------------------------------------

   // Protected -----------------------------------------------------
   
   /**
    * Create the proxy
    * 
    * @todo this should be external configuration
    */
   protected void createProxy() throws Exception
   {
      /* Create an JRMPInvokerProxy that will be associated with a naming JMX
      invoker given by the jmxInvokerName.
      */
     delegateInvoker = (Invoker) Registry.lookup(jmxInvokerName);
     log.debug("Using delegate: " + delegateInvoker
        + " for invoker=" + jmxInvokerName);
     ObjectName targetName = getServiceName();
     Integer nameHash = new Integer(targetName.hashCode());
     Registry.bind(nameHash, targetName);

     Object cacheID = null;
     String proxyBindingName = null;
     String jndiName = null;
     Class[] ifaces = {javax.sql.DataSource.class};
     /* Initialize interceptorClasses with default client interceptor list
        if no client interceptor configuration was provided */
     ArrayList interceptorClasses = new ArrayList();
     interceptorClasses.add(ClientMethodInterceptor.class);
     interceptorClasses.add(InvokerInterceptor.class);
     ClassLoader loader = Thread.currentThread().getContextClassLoader();
     GenericProxyFactory proxyFactory = new GenericProxyFactory();
     theProxy = proxyFactory.createProxy(cacheID, targetName,
        delegateInvoker, jndiName, proxyBindingName, interceptorClasses,
        loader, ifaces);
     log.debug("Created proxy for invoker=" + jmxInvokerName
        + ", targetName=" + targetName + ", nameHash=" + nameHash);
   }
   
   /**
    * Destroy the proxy
    */
   protected void destroyProxy() throws Exception
   {
      ObjectName name = getServiceName();
      Integer nameHash = new Integer(name.hashCode());
      Registry.unbind(nameHash);
   }
   
   /**
    * Calculate the method hashes
    */
   protected void calculateMethodHases() throws Exception
   {
      Method[] methods = DataSource.class.getMethods();
      for(int m = 0; m < methods.length; m ++)
      {
         Method method = methods[m];
         Long hash = new Long(MarshalledInvocation.calculateHash(method));
         marshalledInvocationMapping.put(hash, method);
      }

      // Get the Long to Method mappings
      Map m = MarshalledInvocation.methodToHashesMap(Connection.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
      m = MarshalledInvocation.methodToHashesMap(Statement.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
      m = MarshalledInvocation.methodToHashesMap(CallableStatement.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
      m = MarshalledInvocation.methodToHashesMap(PreparedStatement.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
      m = MarshalledInvocation.methodToHashesMap(ResultSet.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
      m = MarshalledInvocation.methodToHashesMap(DatabaseMetaData.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
      m = MarshalledInvocation.methodToHashesMap(java.sql.Blob.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
      m = MarshalledInvocation.methodToHashesMap(java.sql.Clob.class);
      displayHashes(m);
      marshalledInvocationMapping.putAll(m);
   }
   
   // Private -------------------------------------------------------

   private Object doDataSourceMethod(DataSource ds, Method method, Object[] args)
      throws InvocationTargetException, IllegalAccessException
   {
      Object value = method.invoke(ds, args);
      if( value instanceof Connection )
      {
         value = createConnectionProxy(value);
      }
      else if( value != null && (value instanceof Serializable) == false )
      {
         throw new IllegalAccessException("Method="+method+" does not return Serializable");
      }
      return value;
   }

   private Object doConnectionMethod(Connection conn, Method method, Object[] args)
         throws InvocationTargetException, IllegalAccessException, SQLException
   {
      if( trace )
      {
         log.trace("doConnectionMethod, conn="+conn+", method="+method);
      }
      /*
      if( method.getName().equals("close") )
      {
         Integer id = new Integer(conn.hashCode());
         connectionMap.remove(id);
         log.debug("Closed Connection="+id);
      }
      */
      Object value = method.invoke(conn, args);
      if( value instanceof Statement )
      {
         value = createStatementProxy(value);
      }
      else if(value instanceof DatabaseMetaData)
      {
         value = createDatabaseMetaData(value);
      }
      else if( value != null && (value instanceof Serializable) == false )
      {
         throw new IllegalAccessException("Method="+method+" does not return Serializable");
      }
      return value;
   }

   private Object doStatementMethod(Statement stmt, Method method, Object[] args)
         throws InvocationTargetException, IllegalAccessException, SQLException
   {
      if( trace )
      {
         log.trace("doStatementMethod, conn="+stmt+", method="+method);
      }

      if( method.getName().equals("close") )
      {
         Integer id = new Integer(stmt.hashCode());
         statementMap.remove(id);
         log.debug("Closed Statement="+id);
      }
	  for (int i=0; i < args.length;++i)
	  {
	  	if ((args[i] instanceof java.sql.Blob || args[i] instanceof java.sql.Clob)&& args[i] instanceof org.jboss.proxy.IClientContainer)
	  	{
	  		Integer id = (Integer) ((org.jboss.proxy.IClientContainer)args[i]).getInvocationContext().getCacheId();
            Object results = lobMap.get(id);
            if( results == null )
               throw new IllegalAccessException("Failed to find LOB: "+id);
            args[i] = results;
	  	}
	  }
      Object value = method.invoke(stmt, args);
      if( value instanceof ResultSet )
      {
         value = createResultSetProxy(value);
      }
      else if( value instanceof ResultSetMetaData )
      {
         ResultSetMetaData rmd = (ResultSetMetaData) value;
         value = new SerializableResultSetMetaData(rmd);
      }
      else if ( value instanceof ParameterMetaData )
      {
         ParameterMetaData pmd = (ParameterMetaData) value;
         value = new SerializableParameterMetaData(pmd);
      }
      else if (value instanceof java.sql.Blob || value instanceof java.sql.Clob)
      {
      	value = createLobProxy(value);
      }
      else if (value instanceof Connection)
      {
      	//hack
      	value = createConnectionProxy(value);
      }
      else if( value != null && (value instanceof Serializable) == false )
      {
         throw new IllegalAccessException("Method="+method+" does not return Serializable");
      }
      return value;
   }

   private Object doResultSetMethod(ResultSet results, Method method, Object[] args)
         throws InvocationTargetException, IllegalAccessException, SQLException, IOException
   {
      if( trace )
      {
         log.trace("doResultSetMethod, results="+results+", method="+method);
      }

      if( method.getName().equals("close") )
      {
         Integer id = new Integer(results.hashCode());
         resultSetMap.remove(id);
         log.debug("Closed ResultSet="+id);
      }

      Object value = method.invoke(results, args);
      if( value instanceof ResultSetMetaData )
      {
         ResultSetMetaData rmd = (ResultSetMetaData) value;
         value = new SerializableResultSetMetaData(rmd);
      }
      // Need to create serializable version of ascii stream returned by result set
      if(("getBinaryStream".equals(method.getName()) || "getAsciiStream".equals(method.getName())) && value instanceof InputStream)
      {
         InputStream ins = (InputStream)value;
         value = new SerializableInputStream(ins);
      }
      else if ("getCharacterStream".equals(method.getName()) && value instanceof java.io.Reader)
      {
         java.io.Reader ins = (java.io.Reader)value;
         value = new SerializableReader(ins);
      }
      else if("getClob".equals(method.getName()) || "getBlob".equals(method.getName()))
      {
      	value = createLobProxy(value);
      }
      	
      if( value != null && (value instanceof Serializable) == false )
      {
         throw new IllegalAccessException("Method="+method+" does not return Serializable");
      }
      return value;
   }

  private Object doLobMethod(Object results, Method method, Object[] args)
     throws InvocationTargetException, IllegalAccessException, SQLException, IOException
  {
      if( trace )
      {
         log.trace("doLobMethod, results="+results+", method="+method);
      }

      Object value = method.invoke(results, args);
      if( value instanceof InputStream )
      {
         InputStream ins = (InputStream)value;
         value = new SerializableInputStream(ins);
      }
      else if (value instanceof java.io.Reader)
      {
         java.io.Reader ins = (java.io.Reader)value;
         value = new SerializableReader(ins);
      	
      }
      if( value != null && (value instanceof Serializable) == false )
      {
         throw new IllegalAccessException("Method="+method+" does not return Serializable");
      }
      return value;
  }
   private Object doDatabaseMetaDataMethod(DatabaseMetaData dbMetaData, Method method, Object[] args)
         throws InvocationTargetException, IllegalAccessException
   {
      if( trace )
      {
         log.trace("doDatabaseMetaDataMethod, dbMetaData="+dbMetaData+", method="+method);
      }

      Object value = method.invoke(dbMetaData, args);
      if( value instanceof ResultSet )
      {
         value = createResultSetProxy(value);
      }
      else if( value instanceof Connection )
      {
         value = createConnectionProxy(value);
      }
      if( value != null && (value instanceof Serializable) == false )
      {
         throw new IllegalAccessException("Method="+method+" does not return Serializable");
      }
      return value;
   }

   private Object createConnectionProxy(Object conn)
   {
      Object cacheID = new Integer(conn.hashCode());
      ObjectName targetName = getServiceName();
      String proxyBindingName = null;
      String jndiName = null;
      Class[] ifaces = {java.sql.Connection.class};
      ArrayList interceptorClasses = new ArrayList();
      interceptorClasses.add(ClientMethodInterceptor.class);
      interceptorClasses.add(InvokerInterceptor.class);
      ClassLoader loader = Thread.currentThread().getContextClassLoader();
      GenericProxyFactory proxyFactory = new GenericProxyFactory();
      Object connProxy = proxyFactory.createProxy(cacheID, targetName,
         delegateInvoker, jndiName, proxyBindingName, interceptorClasses,
         loader, ifaces);
      connectionMap.put(cacheID, conn);
      log.debug("Created Connection proxy for invoker=" + jmxInvokerName
         + ", targetName=" + targetName + ", cacheID=" + cacheID);
      return connProxy;
   }

   private Object createStatementProxy(Object stmt)
   {
      Object cacheID = new Integer(stmt.hashCode());
      ObjectName targetName = getServiceName();
      String proxyBindingName = null;
      String jndiName = null;
      Class[] ifaces = stmt.getClass().getInterfaces();
      // Filter out all but java* interfaces
      ArrayList tmp = new ArrayList();
      for(int i = 0; i < ifaces.length; i ++)
      {
         Class c = ifaces[i];
         if( c.getName().startsWith("java") )
            tmp.add(c);
      }
      ifaces = new Class[tmp.size()];
      tmp.toArray(ifaces);
      ArrayList interceptorClasses = new ArrayList();
      interceptorClasses.add(StatementInterceptor.class);
      interceptorClasses.add(ClientMethodInterceptor.class);
      interceptorClasses.add(InvokerInterceptor.class);
      ClassLoader loader = Thread.currentThread().getContextClassLoader();
      GenericProxyFactory proxyFactory = new GenericProxyFactory();
      Object stmtProxy = proxyFactory.createProxy(cacheID, targetName,
         delegateInvoker, jndiName, proxyBindingName, interceptorClasses,
         loader, ifaces);
      statementMap.put(cacheID, stmt);
      log.debug("Created Statement proxy for invoker=" + jmxInvokerName
         + ", targetName=" + targetName + ", cacheID=" + cacheID);
      return stmtProxy;
   }

   private Object createResultSetProxy(Object results)
   {
      Object cacheID = new Integer(results.hashCode());
      ObjectName targetName = getServiceName();
      String proxyBindingName = null;
      String jndiName = null;
      Class[] ifaces = results.getClass().getInterfaces();
      ArrayList interceptorClasses = new ArrayList();
      interceptorClasses.add(ClientMethodInterceptor.class);
      interceptorClasses.add(InvokerInterceptor.class);
      ClassLoader loader = Thread.currentThread().getContextClassLoader();
      GenericProxyFactory proxyFactory = new GenericProxyFactory();
      Object resultsProxy = proxyFactory.createProxy(cacheID, targetName,
         delegateInvoker, jndiName, proxyBindingName, interceptorClasses,
         loader, ifaces);
      resultSetMap.put(cacheID, results);
      log.debug("Created ResultSet proxy for invoker=" + jmxInvokerName
         + ", targetName=" + targetName + ", cacheID=" + cacheID);
      return resultsProxy;
   }

	private Object createLobProxy(Object results)
   {
      Object cacheID = new Integer(results.hashCode());
      ObjectName targetName = getServiceName();
      String proxyBindingName = null;
      String jndiName = null;
      Class[] ifaces = results.getClass().getInterfaces();
      ArrayList interceptorClasses = new ArrayList();
      interceptorClasses.add(ClientMethodInterceptor.class);
      interceptorClasses.add(InvokerInterceptor.class);
      ClassLoader loader = Thread.currentThread().getContextClassLoader();
      GenericProxyFactory proxyFactory = new GenericProxyFactory();
      Object resultsProxy = proxyFactory.createProxy(cacheID, targetName,
         delegateInvoker, jndiName, proxyBindingName, interceptorClasses,
         loader, ifaces);
      lobMap.put(cacheID, results);
      log.debug("Created LOB proxy for invoker=" + jmxInvokerName
         + ", targetName=" + targetName + ", cacheID=" + cacheID);
      return resultsProxy;
   }
   private Object createDatabaseMetaData(Object dbMetaData)
   {
      Object cacheID = new Integer(dbMetaData.hashCode());
      ObjectName targetName = getServiceName();
      String proxyBindingName = null;
      String jndiName = null;
      Class[] ifaces = {java.sql.DatabaseMetaData.class};
      ArrayList interceptorClasses = new ArrayList();
      interceptorClasses.add(ClientMethodInterceptor.class);
      interceptorClasses.add(InvokerInterceptor.class);
      ClassLoader loader = Thread.currentThread().getContextClassLoader();
      GenericProxyFactory proxyFactory = new GenericProxyFactory();
      Object dbMetaDataProxy = proxyFactory.createProxy(cacheID, targetName,
         delegateInvoker, jndiName, proxyBindingName, interceptorClasses,
         loader, ifaces);
      databaseMetaDataMap.put(cacheID, dbMetaData);
      log.debug("Created DatabaseMetadata proxy for invoker=" + jmxInvokerName
         + ", targetName=" + targetName + ", cacheID=" + cacheID);
      return dbMetaDataProxy;

   }

   private void displayHashes(Map m)
   {
      if( trace == false )
         return;

      Iterator keys = m.keySet().iterator();
      while( keys.hasNext() )
      {
         Long key = (Long) keys.next();
         log.trace(key+"="+m.get(key));
      }
   }

   // Inner classes -------------------------------------------------
}