/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2011, Red Hat, Inc., and individual contributors
 * as indicated by the @author tags. See the copyright.txt file in the
 * distribution for a full listing of individual contributors.
 *
 * This is free software; you can redistribute it and/or modify it
 * under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation; either version 2.1 of
 * the License, or (at your option) any later version.
 *
 * This software is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this software; if not, write to the Free
 * Software Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA
 * 02110-1301 USA, or see the FSF site: http://www.fsf.org.
 */

package org.jboss.as.clustering.infinispan;

import java.io.ObjectInputStream;
import java.util.Arrays;
import java.util.Collections;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.Predicate;
import java.util.stream.Collectors;

import javax.management.Attribute;
import javax.management.AttributeList;
import javax.management.AttributeNotFoundException;
import javax.management.DynamicMBean;
import javax.management.InstanceAlreadyExistsException;
import javax.management.InstanceNotFoundException;
import javax.management.IntrospectionException;
import javax.management.InvalidAttributeValueException;
import javax.management.ListenerNotFoundException;
import javax.management.MBeanAttributeInfo;
import javax.management.MBeanException;
import javax.management.MBeanInfo;
import javax.management.MBeanRegistration;
import javax.management.MBeanRegistrationException;
import javax.management.MBeanServer;
import javax.management.NotCompliantMBeanException;
import javax.management.NotificationFilter;
import javax.management.NotificationListener;
import javax.management.ObjectInstance;
import javax.management.ObjectName;
import javax.management.OperationsException;
import javax.management.QueryExp;
import javax.management.ReflectionException;
import javax.management.loading.ClassLoaderRepository;

import org.infinispan.commons.jmx.MBeanServerLookup;
import org.infinispan.jmx.ObjectNameKeys;

/**
 * @author Paul Ferraro
 */
public class MBeanServerProvider implements MBeanServerLookup {

    private final MBeanServer server;

    public MBeanServerProvider(MBeanServer server) {
        // Filter problematic attributes from specific component mbeans
        this.server = new FilteringMBeanServer(server, Collections.singletonMap("Statistics", new TreeSet<>(Arrays.asList("hitTimes", "missTimes", "storeTimes", "removeTimes", "statisticsEnabled"))));
    }

    /**
     * {@inheritDoc}
     * @see org.infinispan.jmx.MBeanServerLookup#getMBeanServer(java.util.Properties)
     */
    @Override
    public MBeanServer getMBeanServer(Properties properties) {
        return this.server;
    }

    private static class FilteredDynamicMBean implements DynamicMBean, MBeanRegistration, Predicate<Attribute> {
        private final DynamicMBean mbean;
        private final MBeanRegistration registration;
        private final Set<String> forbiddenAttributes;

        FilteredDynamicMBean(DynamicMBean mbean, Set<String> forbiddenAttributes) {
            this.mbean = mbean;
            this.registration = (MBeanRegistration) mbean;
            this.forbiddenAttributes = forbiddenAttributes;
        }

        @Override
        public void postDeregister() {
            this.registration.postDeregister();
        }

        @Override
        public void postRegister(Boolean registrationDone) {
            this.registration.postRegister(registrationDone);
        }

        @Override
        public void preDeregister() throws Exception {
            this.registration.preDeregister();
        }

        @Override
        public ObjectName preRegister(MBeanServer server, ObjectName name) throws Exception {
            return this.registration.preRegister(server, name);
        }

        @Override
        public Object getAttribute(String attribute) throws AttributeNotFoundException, MBeanException, ReflectionException {
            if (this.forbiddenAttributes.contains(attribute)) throw new AttributeNotFoundException();
            return this.mbean.getAttribute(attribute);
        }

        @Override
        public AttributeList getAttributes(String[] attributes) {
            return this.mbean.getAttributes(Arrays.asList(attributes).stream().filter(attribute -> !this.forbiddenAttributes.contains(attribute)).toArray(String[]::new));
        }

        @Override
        public MBeanInfo getMBeanInfo() {
            MBeanInfo info = this.mbean.getMBeanInfo();
            return new MBeanInfo(info.getClassName(), info.getDescription(), Arrays.asList(info.getAttributes()).stream().filter(attribute -> !this.forbiddenAttributes.contains(attribute.getName())).toArray(MBeanAttributeInfo[]::new), info.getConstructors(), info.getOperations(), info.getNotifications());
        }

        @Override
        public Object invoke(String actionName, Object[] params, String[] signature) throws MBeanException, ReflectionException {
            return this.mbean.invoke(actionName, params, signature);
        }

        @Override
        public void setAttribute(Attribute attribute) throws AttributeNotFoundException, InvalidAttributeValueException, MBeanException, ReflectionException {
            if (this.test(attribute)) {
                this.mbean.setAttribute(attribute);
            }
        }

        @Override
        public AttributeList setAttributes(AttributeList attributes) {
            return this.mbean.setAttributes(new AttributeList(attributes.asList().stream().filter(this).collect(Collectors.toList())));
        }

        @Override
        public boolean test(Attribute attribute) {
            return !this.forbiddenAttributes.contains(attribute.getName());
        }
    }

    private static class FilteringMBeanServer implements MBeanServer {
        private final MBeanServer server;
        private final Map<String, Set<String>> forbiddenAttributes;

        FilteringMBeanServer(MBeanServer server, Map<String, Set<String>> forbiddenAttributes) {
            this.server = server;
            this.forbiddenAttributes = forbiddenAttributes;
        }

        @Override
        public ObjectInstance registerMBean(Object object, ObjectName name) throws InstanceAlreadyExistsException, MBeanRegistrationException, NotCompliantMBeanException {
            String component = name.getKeyProperty(ObjectNameKeys.COMPONENT);
            Set<String> forbiddenAttributes = (component != null) ? this.forbiddenAttributes.get(component) : null;
            Object mbean = (forbiddenAttributes != null) && (object instanceof DynamicMBean) ? new FilteredDynamicMBean((DynamicMBean) object, forbiddenAttributes) : object;
            return this.server.registerMBean(mbean, name);
        }

        @Override
        public void addNotificationListener(ObjectName arg0, NotificationListener arg1, NotificationFilter arg2, Object arg3) throws InstanceNotFoundException {
            this.server.addNotificationListener(arg0, arg1, arg2, arg3);
        }

        @Override
        public void addNotificationListener(ObjectName arg0, ObjectName arg1, NotificationFilter arg2, Object arg3) throws InstanceNotFoundException {
            this.server.addNotificationListener(arg0, arg1, arg2, arg3);
        }

        @Override
        public ObjectInstance createMBean(String arg0, ObjectName arg1) throws ReflectionException, InstanceAlreadyExistsException, MBeanRegistrationException, MBeanException, NotCompliantMBeanException {
            return this.server.createMBean(arg0, arg1);
        }

        @Override
        public ObjectInstance createMBean(String arg0, ObjectName arg1, ObjectName arg2) throws ReflectionException, InstanceAlreadyExistsException, MBeanRegistrationException, MBeanException, NotCompliantMBeanException, InstanceNotFoundException {
            return this.server.createMBean(arg0, arg1, arg2);
        }

        @Override
        public ObjectInstance createMBean(String arg0, ObjectName arg1, Object[] arg2, String[] arg3) throws ReflectionException, InstanceAlreadyExistsException, MBeanRegistrationException, MBeanException, NotCompliantMBeanException {
            return this.server.createMBean(arg0, arg1, arg2, arg3);
        }

        @Override
        public ObjectInstance createMBean(String arg0, ObjectName arg1, ObjectName arg2, Object[] arg3, String[] arg4) throws ReflectionException, InstanceAlreadyExistsException, MBeanRegistrationException, MBeanException, NotCompliantMBeanException, InstanceNotFoundException {
            return this.server.createMBean(arg0, arg1, arg2, arg3, arg4);
        }

        @Deprecated
        @Override
        public ObjectInputStream deserialize(ObjectName arg0, byte[] arg1) throws InstanceNotFoundException, OperationsException {
            return this.server.deserialize(arg0, arg1);
        }

        @Deprecated
        @Override
        public ObjectInputStream deserialize(String arg0, byte[] arg1) throws OperationsException, ReflectionException {
            return this.server.deserialize(arg0, arg1);
        }

        @Deprecated
        @Override
        public ObjectInputStream deserialize(String arg0, ObjectName arg1, byte[] arg2) throws InstanceNotFoundException, OperationsException, ReflectionException {
            return this.server.deserialize(arg0, arg1, arg2);
        }

        @Override
        public Object getAttribute(ObjectName arg0, String arg1) throws MBeanException, AttributeNotFoundException, InstanceNotFoundException, ReflectionException {
            return this.server.getAttribute(arg0, arg1);
        }

        @Override
        public AttributeList getAttributes(ObjectName arg0, String[] arg1) throws InstanceNotFoundException, ReflectionException {
            return this.server.getAttributes(arg0, arg1);
        }

        @Override
        public ClassLoader getClassLoader(ObjectName arg0) throws InstanceNotFoundException {
            return this.server.getClassLoader(arg0);
        }

        @Override
        public ClassLoader getClassLoaderFor(ObjectName arg0) throws InstanceNotFoundException {
            return this.server.getClassLoaderFor(arg0);
        }

        @Override
        public ClassLoaderRepository getClassLoaderRepository() {
            return this.server.getClassLoaderRepository();
        }

        @Override
        public String getDefaultDomain() {
            return this.server.getDefaultDomain();
        }

        @Override
        public String[] getDomains() {
            return this.server.getDomains();
        }

        @Override
        public Integer getMBeanCount() {
            return this.server.getMBeanCount();
        }

        @Override
        public MBeanInfo getMBeanInfo(ObjectName arg0) throws InstanceNotFoundException, IntrospectionException, ReflectionException {
            return this.server.getMBeanInfo(arg0);
        }

        @Override
        public ObjectInstance getObjectInstance(ObjectName arg0) throws InstanceNotFoundException {
            return this.server.getObjectInstance(arg0);
        }

        @Override
        public Object instantiate(String arg0) throws ReflectionException, MBeanException {
            return this.server.instantiate(arg0);
        }

        @Override
        public Object instantiate(String arg0, ObjectName arg1) throws ReflectionException, MBeanException, InstanceNotFoundException {
            return this.server.instantiate(arg0, arg1);
        }

        @Override
        public Object instantiate(String arg0, Object[] arg1, String[] arg2) throws ReflectionException, MBeanException {
            return this.server.instantiate(arg0, arg1, arg2);
        }

        @Override
        public Object instantiate(String arg0, ObjectName arg1, Object[] arg2, String[] arg3) throws ReflectionException, MBeanException, InstanceNotFoundException {
            return this.server.instantiate(arg0, arg1, arg2, arg3);
        }

        @Override
        public Object invoke(ObjectName arg0, String arg1, Object[] arg2, String[] arg3) throws InstanceNotFoundException, MBeanException, ReflectionException {
            return this.server.invoke(arg0, arg1, arg2, arg3);
        }

        @Override
        public boolean isInstanceOf(ObjectName arg0, String arg1) throws InstanceNotFoundException {
            return this.server.isInstanceOf(arg0, arg1);
        }

        @Override
        public boolean isRegistered(ObjectName arg0) {
            return this.server.isRegistered(arg0);
        }

        @Override
        public Set<ObjectInstance> queryMBeans(ObjectName arg0, QueryExp arg1) {
            return this.server.queryMBeans(arg0, arg1);
        }

        @Override
        public Set<ObjectName> queryNames(ObjectName arg0, QueryExp arg1) {
            return this.server.queryNames(arg0, arg1);
        }

        @Override
        public void removeNotificationListener(ObjectName arg0, ObjectName arg1) throws InstanceNotFoundException, ListenerNotFoundException {
            this.server.removeNotificationListener(arg0, arg1);
        }

        @Override
        public void removeNotificationListener(ObjectName arg0, NotificationListener arg1) throws InstanceNotFoundException, ListenerNotFoundException {
            this.server.removeNotificationListener(arg0, arg1);
        }

        @Override
        public void removeNotificationListener(ObjectName arg0, ObjectName arg1, NotificationFilter arg2, Object arg3) throws InstanceNotFoundException, ListenerNotFoundException {
            this.server.removeNotificationListener(arg0, arg1, arg2, arg3);
        }

        @Override
        public void removeNotificationListener(ObjectName arg0, NotificationListener arg1, NotificationFilter arg2, Object arg3) throws InstanceNotFoundException, ListenerNotFoundException {
            this.server.removeNotificationListener(arg0, arg1, arg2, arg3);
        }

        @Override
        public void setAttribute(ObjectName arg0, Attribute arg1) throws InstanceNotFoundException, AttributeNotFoundException, InvalidAttributeValueException, MBeanException, ReflectionException {
            this.server.setAttribute(arg0, arg1);
        }

        @Override
        public AttributeList setAttributes(ObjectName arg0, AttributeList arg1) throws InstanceNotFoundException, ReflectionException {
            return this.server.setAttributes(arg0, arg1);
        }

        @Override
        public void unregisterMBean(ObjectName arg0) throws InstanceNotFoundException, MBeanRegistrationException {
            this.server.unregisterMBean(arg0);
        }
    }
}
