/**
 * JOnAS: Java(TM) Open Application Server
 * Copyright (C) 2012 Bull S.A.S.
 * Contact: jonas-team@ow2.org
 *
 * This library 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 library 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 library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA
 *
 * --------------------------------------------------------------------------
 * $Id$
 * --------------------------------------------------------------------------
 */

package org.ow2.jonas.lib.tenant.interceptor.jmx;

import org.ow2.jonas.jmx.Interceptor;
import org.ow2.jonas.lib.tenant.context.TenantContext;
import org.ow2.jonas.lib.tenant.context.TenantCurrent;

import javax.interceptor.InvocationContext;
import javax.management.MBeanServer;
import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;
import javax.management.QueryExp;
import java.util.HashMap;
import java.util.Hashtable;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * Intercept MBeanServer's call and add tenant-id attribute
 * @author Mohammed Boukada
 */
public class JMXTenantIdInterceptor implements Interceptor {

    /**
     * Tenant attribute nambe MBeans ObjectName
     */
    private String tenantIdAttributeName;

    /**
     * Whether tenants are allowed to access to platform MBeans
     */
    private boolean allowToAccessPlatformMBeans;

    /**
     * Map containing tenants names and their corresponding tenant-id
     */
    private Map<String, List<ObjectName>> tenantBeans = new HashMap<String, List<ObjectName>>();

    /**
     * Whether is a JMX session
     */
    private boolean isJMXSession;

    /**
     * Construct an Interceptor
     * @param tenantIdAttributeName
     * @param allowToAccessPlatformMBeans
     */
    public JMXTenantIdInterceptor(String tenantIdAttributeName, boolean allowToAccessPlatformMBeans) {
        this.tenantIdAttributeName = tenantIdAttributeName;
        this.allowToAccessPlatformMBeans = allowToAccessPlatformMBeans;
    }

    /**
     * Invoke interceptor method
     * @param invocationContext use to apply the interceptors sequence
     * @return invocation results
     * @throws Exception
     */
    public Object invoke(InvocationContext invocationContext) throws Exception {
        String methodName = invocationContext.getMethod().getName();
        Object[] parameters = invocationContext.getParameters();
        Integer index = null;
        boolean saveBean = false;
        isJMXSession = false;

        // Define the position of ObjectName parameter
        if (methodName.equals("registerMBean")
                || methodName.equals("createMBean")
                ){
            index = 1;
            saveBean = true;
        } else if (methodName.equals("unregisterMBean")
                || methodName.equals("getObjectInstance")
                || methodName.equals("queryMBeans")
                || methodName.equals("queryNames")
                || methodName.equals("isRegistered")
                || methodName.equals("getAttribute")
                || methodName.equals("getAttributes")
                || methodName.equals("setAttribute")
                || methodName.equals("setAttributes")
                || methodName.equals("invoke")
                || methodName.equals("addNotificationListener")
                || methodName.equals("removeNotificationListener")
                || methodName.equals("getMBeanInfo")
                || methodName.equals("isInstanceOf")) {

            index = 0;
        }

        if (index != null) {
            ObjectName oldObjectName = (ObjectName) parameters[index];
            ObjectName newObjectName = null;

            if (methodName.equals("queryMBeans") || methodName.equals("queryNames")) {
                newObjectName = updateNamePattern(oldObjectName, invocationContext);
            } else if (methodName.equals("registerMBean") || methodName.equals("createMBean")) {
                newObjectName = updateObjectName(oldObjectName, saveBean);
            } else {
                if (isTenantBean(oldObjectName)) {
                    newObjectName = updateObjectName(oldObjectName, saveBean);
                }
            }

            if (newObjectName != null) {
                parameters[index] = newObjectName;
                invocationContext.setParameters(parameters);
            }
        }
        return invocationContext.proceed();
    }

    /**
     * Update ObjectName by adding tenantId
     * @param old ObjectName to update
     * @param saveBean save the old ObjectName as registered by the current tenant
     * @return ObjectName updated
     */
    private ObjectName updateObjectName (ObjectName old, boolean saveBean) {
        if (isObjectNameUpdated(old) || old == null) {
            return null;
        }

        String tenantId = getTenantId(false);
        if (tenantId != null && !tenantId.equals(TenantContext.DEFAULT_TENANT_ID)) {
            try {
                if (saveBean) {
                    if(tenantBeans.get(tenantId) == null) {
                        List<ObjectName> list = new LinkedList<ObjectName>();
                        list.add(old);
                        tenantBeans.put(tenantId, list);
                    } else {
                        List<ObjectName> list = tenantBeans.get(tenantId);
                        list.add(old);
                    }
                }
                return new ObjectName(new StringBuilder(old.toString()).append("," + tenantIdAttributeName + "=").append(tenantId).toString());
            } catch (MalformedObjectNameException e) {
                throw new RuntimeException("Failed to format ObjectName [" + old + "]", e);
            }
        }
        return null;
    }

    /**
     * Update name pattern identifying the MBeans to be retrieved
     * @param old name pattern
     * @param invocationContext
     * @return new name pattern
     */
    private ObjectName updateNamePattern (ObjectName old, InvocationContext invocationContext) throws Exception {
        String tenantId = getTenantId(true);
        ObjectName newObjectName = null;
        // If tenantId is not set, then no customization
        if (tenantId == null || tenantId.equals(TenantContext.DEFAULT_TENANT_ID)) {
            return null;
        }

        // if is a JMX administration session
        // or a query on a name registered by the tenant
        if (isJMXSession || isTenantQuery(invocationContext)) {
            newObjectName = updateObjectName(old, false);
        }

        // if initial ObjectName pattern was null
        // null is used to retrieve all platform MBeans
        if (old == null) {
            try {
                newObjectName = new ObjectName("*:" + tenantIdAttributeName + "=" + tenantId + ",*");
            } catch (MalformedObjectNameException e) {
                throw new RuntimeException("Failed to format Name Pattern [" + old + "]", e);
            }
        }
        return newObjectName;
    }

    /**
     * Gets tenantId from TenantContext
     * @param query whether the call is a query
     * @return tenantId
     */
    private String getTenantId(boolean query) {
        String tenantId = null;
        if (TenantCurrent.getCurrent().getTenantContext() != null) {
            // If the call is a query
            if (query) {
                // Check if it is a connection on a JMX Connector (JConsole, ...)
                String jmxSessionTenantId = TenantCurrent.getCurrent().getTenantContext().getJmxSessionTenantId();
                if (jmxSessionTenantId != null) {
                    isJMXSession = true;
                    return jmxSessionTenantId;
                }
            }
            tenantId = TenantCurrent.getCurrent().getTenantContext().getTenantId();
        }
        return tenantId;
    }

    /**
     * Whether tenant-id attribute was set or not in the ObjectName
     * @param objectName objectName to test
     * @return
     */
    private boolean isObjectNameUpdated(ObjectName objectName) {
        if (objectName == null) {
            return false;
        }
        Hashtable properties = objectName.getKeyPropertyList();
        return properties.containsKey(tenantIdAttributeName);
    }

    /**
     * Whether the bean was registered by this tenant
     * @param name
     * @return
     */
    private boolean isTenantBean(ObjectName name) {
        String tenantId = getTenantId(false);
        if (tenantId == null) {
            return false;
        }
        List<ObjectName> list = tenantBeans.get(tenantId);
        if (list == null) {
            return false;
        }
        return list.contains(name);
    }

    /**
     * Whether the query is on tenant's MBeans
     * @param invocationContext
     * @return
     */
    private boolean isTenantQuery(InvocationContext invocationContext) throws Exception {
        // Execute the query and see if tenant-id property is
        // present on ObjectNames
        Object[] args = invocationContext.getParameters();
        ObjectName oldObjectName = (ObjectName) args[0];
        QueryExp queryExp = (QueryExp) args[1];
        Set<ObjectName> mbeans = ((MBeanServer) invocationContext.getTarget()).queryNames(oldObjectName, queryExp);

        // if 'tenant-id' property is present in one or more
        // ObjectNames, then customize the query
        for (ObjectName mbean : mbeans) {
            if (mbean.getKeyProperty(tenantIdAttributeName) != null) {
                return true;
            }
        }
        return false;
    }
}
