import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.Proxy;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import javax.management.Attribute;
import javax.management.AttributeList;
import javax.management.MBeanServerConnection;
import javax.management.MBeanServerInvocationHandler;
import javax.management.ObjectName;

/**
 * A Factory class for generating MBean proxies
 */
public class ProxyFactory {
    
    /**
     * obtain a <code>SnapshotProxy</code> for a given MBean
     * @param connection MBeanServerConnection
     * @param objectName ObjectName of the MBean
     * @param interfaceClass MBean's interface class
     * @return a <code>SnapshotProxy</code> for the MBean
     */
    public Object getSnapshotProxy(MBeanServerConnection connection,
            ObjectName objectName,
            Class interfaceClass) {

        final InvocationHandler handler =
                new SnapshotProxy(connection,
                objectName,
                interfaceClass);
        
        final Class[] interfaces = new Class[] {interfaceClass};
        
        return Proxy.newProxyInstance(interfaceClass.getClassLoader(),
                interfaces,
                handler);
    }
    
    /**
     * Implementation class of the <code>SnapshotProxy</code>
     */
    private static class SnapshotProxy extends MBeanServerInvocationHandler {
        
        private MBeanServerConnection connection;
        private ObjectName objectName;
        private Class interfaceClass;
        private Map<String,Attribute>attributeMap = null;
        
        /**
         * Constructor
         * @param connection MBeanServerConnection
         * @param objectName ObjectName of the MBean
         * @param interfaceClass Class defining MBean's interface
         */
        private SnapshotProxy(MBeanServerConnection connection,
                ObjectName objectName,
                Class interfaceClass) {
            
            super(connection, objectName);
            this.connection = connection;
            this.objectName = objectName;
            this.interfaceClass = interfaceClass;
        }
        
        /* Javadoc inherited from MBeanServerInvocationHandler */
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            
            final String methodName = method.getName();
            final Class[] paramTypes = method.getParameterTypes();
            final Class returnType = method.getReturnType();
            
            /* Inexplicably, InvocationHandler specifies that args is null
               when the method takes no arguments rather than a
               zero-length array.  */
            final int nargs = (args == null) ? 0 : args.length;
            
            String attributeName = decodeAttributeName(methodName, nargs, returnType);
            
            /* We only implement the returning of attribute values */
            if (attributeName == null) {
                throw new RuntimeException("Proxy does not implement method "+method.getName());
            }
            
            /* If we've never been called before, fetch all the attributes */
            if (attributeMap == null) {
                fillAttributeMap();
            }
            
            if (attributeMap == null || !attributeMap.containsKey(attributeName)) {
                // If we were unable to snapshot the attribute due
                // to some kind of error, fallback on default proxy
                return super.invoke(proxy, method, args);
            }

            return attributeMap.get(attributeName);
            
        }
        
        /**
         * Fill the internal map of attribute name to attribute value, done
         * as a single snapshot using a call to <code>getAttributes</code>
         */
        private void fillAttributeMap() {
            
            String[] attributeNames = getAttributeNames(interfaceClass);
            try {
                
                Map<String, Attribute> myMap = new HashMap();
                
                AttributeList attributeList =
                        connection.getAttributes(objectName, attributeNames);
                for (Iterator<Attribute> i = attributeList.iterator(); i.hasNext(); ) {
                    Attribute attr = i.next();
                    myMap.put(attr.getName(), attr);
                }
                attributeMap = myMap;
            } catch (Exception e) {
                // leave map uninitialized in case of error
            }
        }
        
        /**
         * Return the names of all the attributes defined by this interface and
         * any interfaces this interface has extended.
         * @param interfaceClass The MBean's interface class
         * @return A String[] containing all the attribute names.
         */
        private String[] getAttributeNames(Class interfaceClass) {
            Set<String> nameSet = new HashSet<String>();
            
            addAttributeNames(nameSet, interfaceClass);
                
            // recurse over all other interfaces implemented by this interface
            Class[] interfaces = interfaceClass.getInterfaces();
            for (int i = 0; i < interfaces.length; i++) {
                addAttributeNames(nameSet, interfaces[i]);
            }
            return (String[])nameSet.toArray(new String[] {});
        }
        
        /**
         * Add the names of all the attributes defined in this interface to
         * a given <code>Set</code>
         * @param nameSet All discovered names are added into this set.
         * @param interfaceClass class to introspect
         */
        private void addAttributeNames(Set<String> nameSet, Class interfaceClass) {
            
            Method[] methods = interfaceClass.getDeclaredMethods();
            
            for (int i = 0 ; i < methods.length; i++) {
                
                // Only consider public attributes
                if (!Modifier.isPublic(methods[i].getModifiers()))
                    continue;
                
                String methodName = methods[i].getName();
                final Class[] paramTypes = methods[i].getParameterTypes();
                final Class returnType = methods[i].getReturnType();
                
                String attrName =
                        decodeAttributeName(methodName, paramTypes.length, returnType);
                if (attrName != null) {
                    nameSet.add(attrName);
                }
            }
            
            return;
        }
        
        /**
         * Maps from a method name to the attribute name, or returns null if
         * method is not an attribute "getter".
         * @param methodName name of method
         * @param argCount number of arguments the method takes
         * @param returnType the return type of the method
         * @return the JMX MBean's attribute name, if this was a read-only
         * accessor to an attribute, or null if not.
         */
        private String decodeAttributeName(String methodName, int argCount, Class returnType) {
            
            String attributeName = null;
            
            if (methodName.startsWith("get")
                    && methodName.length() > 3
                    && argCount == 0
                    && !returnType.equals(Void.TYPE)) {
                attributeName = methodName.substring(3);
            } else if (methodName.startsWith("is")
                    && methodName.length() > 2
                    && argCount == 0
                    && returnType.equals(Boolean.TYPE)) {
                attributeName = methodName.substring(2);
            }
            
            return attributeName;
        }
    }
}